From 423f34326ed420e5acf0b1fb862a1b5ffb7e0fa7 Mon Sep 17 00:00:00 2001 From: Youyuan Li <105433145+Darcy-Lee@users.noreply.github.com> Date: Thu, 1 Oct 2026 13:04:31 +0300 Subject: [PATCH 01/22] [Model] Add YuE2-3B text-to-music (model + /v1/audio/speech adapter) (#7886) Signed-off-by: Youyuan Li <105433145+Darcy-Lee@users.noreply.github.com> Signed-off-by: princepride Co-authored-by: princepride --- docs/models/supported_models.md | 1 + examples/offline_inference/yue2/end2end.py | 177 ++ recipes/m-a-p/YuE2-3B.md | 85 + tests/assets/yue2/twinkle_golden.abc | 6 + tests/config/test_environment_variables.py | 4 +- .../test_invalid_audio_speech.py | 80 + tests/e2e/offline_inference/test_yue2.py | 82 + tests/e2e/online_serving/test_yue2.py | 237 +++ tests/e2e/online_serving/yue2_metal_verify.py | 185 ++ .../openai_api/test_serving_speech.py | 32 + .../openai_api/test_tts_detection.py | 1 + tests/worker/test_gpu_ar_model_runner.py | 18 +- .../config/environment_variable_inventory.py | 1 + vllm_omni/config/pipeline_registry.py | 2 + vllm_omni/deploy/yue2.yaml | 55 + .../entrypoints/openai/serving_speech.py | 18 + .../openai/tts_adapters/__init__.py | 1 + .../entrypoints/openai/tts_adapters/base.py | 9 + .../entrypoints/openai/tts_adapters/yue2.py | 267 +++ vllm_omni/model_executor/models/registry.py | 6 + .../model_executor/models/yue2/__init__.py | 7 + .../model_executor/models/yue2/pipeline.py | 53 + vllm_omni/model_executor/models/yue2/yue2.py | 1750 +++++++++++++++++ vllm_omni/tokenizers/yue2_prompt.py | 97 + vllm_omni/tokenizers/yue2_tokenizer.py | 50 + .../transformers_utils/configs/__init__.py | 3 + vllm_omni/transformers_utils/configs/yue2.py | 73 + 27 files changed, 3294 insertions(+), 6 deletions(-) create mode 100644 examples/offline_inference/yue2/end2end.py create mode 100644 recipes/m-a-p/YuE2-3B.md create mode 100644 tests/assets/yue2/twinkle_golden.abc create mode 100644 tests/e2e/offline_inference/test_yue2.py create mode 100644 tests/e2e/online_serving/test_yue2.py create mode 100644 tests/e2e/online_serving/yue2_metal_verify.py create mode 100644 vllm_omni/deploy/yue2.yaml create mode 100644 vllm_omni/entrypoints/openai/tts_adapters/yue2.py create mode 100644 vllm_omni/model_executor/models/yue2/__init__.py create mode 100644 vllm_omni/model_executor/models/yue2/pipeline.py create mode 100644 vllm_omni/model_executor/models/yue2/yue2.py create mode 100644 vllm_omni/tokenizers/yue2_prompt.py create mode 100644 vllm_omni/tokenizers/yue2_tokenizer.py create mode 100644 vllm_omni/transformers_utils/configs/yue2.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 5b8bf3464..336b41d8b 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -83,6 +83,7 @@ th { | `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | — | | `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | — | | `MiniMaxMusic3ForConditionalGeneration` | MiniMax Music 3 (text-to-music) | `MiniMaxAI/MiniMax-Music3` | ✅︎ | | | | — | +| `Yue2ForCausalLM` | YuE2-3B (text-to-music with ABC plan) | `m-a-p/YuE2-3B` | ✅︎ | | | | — | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | ✅︎ | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | diff --git a/examples/offline_inference/yue2/end2end.py b/examples/offline_inference/yue2/end2end.py new file mode 100644 index 000000000..39fddd371 --- /dev/null +++ b/examples/offline_inference/yue2/end2end.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""YuE2-3B offline end-to-end text-to-music. + +One song is one or two engine requests against the single-stage pipeline: +``cot=full|melody`` first asks the model to write an ABC score (the abc phase, +tokens only), then the semantic phase generates codec tokens and the model +solves the acoustic ODE + decodes 48 kHz stereo as the request finishes. +``cot=off`` and a user-supplied ``--abc-file`` skip the first request. +Prefix caching lets the second request reuse the first one's blocks. + +Example: + python end2end.py --model /models/YuE2-3B --vae /models/YuE2-Vae \ + --style "gentle lo-fi hip-hop, 85 BPM, F minor" \ + --lyrics "[Verse]\nHumming on a quiet street\n[Chorus]\nStay a while" \ + --cot full --seed 831001 --max-frames 200 --output song.wav +""" + +from __future__ import annotations + +import argparse +import copy +from pathlib import Path + +import soundfile as sf +import torch + +from vllm_omni import Omni +from vllm_omni.model_executor.models.yue2.yue2 import ( + ABC_SAMPLING, + KEY_MAX_AUDIO_FRAMES, + KEY_MIN_TOKENS, + KEY_PENALTY_WINDOW, + KEY_PHASE, + KEY_PREFIX_IDS, + KEY_REPETITION_PENALTY, + KEY_SEED, + KEY_SKIP_SYNTHESIS, + KEY_TEMPERATURE, + KEY_TOP_K, + KEY_TOP_P, + SEMANTIC_SAMPLING, + STOP_TOKEN_IDS, +) +from vllm_omni.tokenizers.yue2_prompt import ( + abc_ids_from_generated, + abc_prefix_ids, + semantic_prefix_ids, +) +from vllm_omni.tokenizers.yue2_tokenizer import YuE2TextTokenizer + + +def sampling_params(engine, *, phase, seed, max_frames, prompt_ids): + """Deep-copy the stage defaults, then set the yue2_* keys for one phase.""" + params = copy.deepcopy(engine.resolve_sampling_params_list(None))[0] + preset = ABC_SAMPLING if phase == "abc" else SEMANTIC_SAMPLING + params.extra_args = { + **(params.extra_args or {}), + KEY_PHASE: phase, + KEY_SEED: seed, + KEY_TEMPERATURE: preset["temperature"], + KEY_TOP_P: preset["top_p"], + KEY_TOP_K: preset["top_k"], + KEY_REPETITION_PENALTY: preset["repetition_penalty"], + KEY_PENALTY_WINDOW: preset["penalty_window"], + KEY_MIN_TOKENS: preset["min_tokens"], + KEY_MAX_AUDIO_FRAMES: max_frames, + KEY_SKIP_SYNTHESIS: phase == "abc", + # Full prompt ids: under a KV prefix-cache hit the engine schedules + # only the uncached tail, so the model cannot rebuild its NAR + # conditioning prefix from the scheduled tokens alone. + KEY_PREFIX_IDS: list(prompt_ids), + } + params.max_tokens = preset["max_tokens"] if phase == "abc" else max_frames + 1 + params.stop_token_ids = list(STOP_TOKEN_IDS) + params.detokenize = False + return params + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", required=True, help="m-a-p/YuE2-3B checkpoint dir or repo id") + parser.add_argument("--vae", default=None, help="m-a-p/YuE2-Vae dir or repo id (or $YUE2_VAE)") + parser.add_argument("--style", required=True) + parser.add_argument("--lyrics", required=True) + parser.add_argument("--cot", choices=["off", "melody", "full"], default="full") + parser.add_argument("--seed", type=int, default=831001) + parser.add_argument("--max-frames", type=int, default=200, help="semantic frame budget (25 frames = 1 s)") + parser.add_argument("--abc-file", default=None, help="external ABC score (requires cot=melody/full)") + parser.add_argument("--output", default="yue2_song.wav") + parser.add_argument( + "--dump-tokens", + default=None, + metavar="PATH", + help="write prompt + generated token ids to this json (default: .tokens.json)", + ) + parser.add_argument( + "--gpu-memory-utilization", + type=float, + default=None, + help="override the deploy yaml's per-stage gpu_memory_utilization " + "(necessary on cards shared with other tenants: the yaml fraction is " + "of the WHOLE card, not of the free remainder)", + ) + args = parser.parse_args() + + if args.vae: + import os + + os.environ["YUE2_VAE"] = args.vae + + tokenizer = YuE2TextTokenizer(Path(args.model) / "qwen.tiktoken") + engine_kwargs: dict = {} + if args.gpu_memory_utilization is not None: + engine_kwargs["gpu_memory_utilization"] = args.gpu_memory_utilization + engine = Omni(model=args.model, **engine_kwargs) + + abc_ids = None + if args.abc_file is not None: + if args.cot == "off": + raise SystemExit("--abc-file requires --cot melody|full") + abc_ids = tokenizer.encode(Path(args.abc_file).read_text()) + print(f"Using supplied ABC score ({len(abc_ids)} tokens)") + elif args.cot != "off": + prompt_ids = abc_prefix_ids(tokenizer.encode, args.style, args.lyrics, args.cot) + prompt = {"prompt_token_ids": prompt_ids} + params = sampling_params(engine, phase="abc", seed=args.seed, max_frames=args.max_frames, prompt_ids=prompt_ids) + outputs = engine.generate([prompt], [params]) + generated = list(outputs[0].outputs[0].token_ids) + abc_ids = abc_ids_from_generated(generated) + print(f"Generated ABC score ({len(abc_ids)} tokens):\n{tokenizer.decode(abc_ids)}\n") + + prompt_ids = semantic_prefix_ids(tokenizer.encode, args.style, args.lyrics, args.cot, abc_ids=abc_ids) + prompt = {"prompt_token_ids": prompt_ids} + params = sampling_params( + engine, phase="semantic", seed=args.seed, max_frames=args.max_frames, prompt_ids=prompt_ids + ) + outputs = engine.generate([prompt], [params]) + output = outputs[0].outputs[0] + generated_ids = list(output.token_ids) + if args.dump_tokens: + import json + + payload = { + "prompt_token_ids": prompt_ids, + "generated_token_ids": generated_ids, + "seed": args.seed, + "cot": args.cot, + "max_frames": args.max_frames, + } + dump_path = args.dump_tokens if args.dump_tokens.endswith(".json") else args.output + ".tokens.json" + Path(dump_path).write_text(json.dumps(payload)) + print(f"Dumped {len(generated_ids)} generated token ids to {dump_path}") + mm = output.multimodal_output or {} + audio = mm.get("audio") + if audio is None and "model_outputs" in mm: + audio = mm["model_outputs"] + if audio is None: + raise SystemExit("No audio in the model output; check server logs") + sr = int(mm.get("sr").item()) if mm.get("sr") is not None else 48000 + truncated = False + meta = mm.get("meta") or {} + if "truncated" in meta: + truncated = bool(torch.as_tensor(meta["truncated"]).reshape(-1)[-1].item()) + + waveform = audio.float() + # soundfile wants [frames, channels] float in [-1, 1]; torchaudio has no + # wheel matching vllm 0.29.0's torch pin, so the driver uses libsndfile. + sf.write(args.output, waveform.T.contiguous(), sr) + print( + f"Saved {args.output}: {waveform.shape[-1] / sr:.1f}s @ {sr} Hz stereo, " + f"truncated={truncated}, generated {len(output.token_ids)} semantic tokens" + ) + + +if __name__ == "__main__": + main() diff --git a/recipes/m-a-p/YuE2-3B.md b/recipes/m-a-p/YuE2-3B.md new file mode 100644 index 000000000..8ad3b008d --- /dev/null +++ b/recipes/m-a-p/YuE2-3B.md @@ -0,0 +1,85 @@ +# YuE2-3B + +> Text-to-music: style + lyrics (optionally an ABC score) in, a 48 kHz stereo song out + +## Summary + +- Vendor: m-a-p (Multimodal Art Projection) +- Model: `m-a-p/YuE2-3B` + `m-a-p/YuE2-Vae` +- Task: Text-to-music generation with an editable symbolic plan (ABC notation) +- Mode: Offline generation and online `/v1/audio/speech` serving (adapter `tts_adapters/yue2.py`) +- Integration issue: [#7661](https://github.com/vllm-project/vllm-omni/issues/7661) + +## When to use this recipe + +Use it to generate songs from a style description and lyrics. Unlike MiniMax +Music 3, YuE2-3B first writes an ABC score (chords and melody, `cot=full` / +`cot=melody`) that you can inspect, edit and resubmit (`--abc-file`), then +renders it with the lyrics. `cot=off` skips the score and generates music +directly. Weights are **CC BY-NC 4.0** (non-commercial); this recipe loads +user-downloaded weights and bundles none. + +## References + +- Model card: [m-a-p/YuE2-3B](https://huggingface.co/m-a-p/YuE2-3B) +- Upstream: [multimodal-art-projection/YuE](https://github.com/multimodal-art-projection/YuE) (pinned `bd90e4c`) +- Offline example: `examples/offline_inference/yue2/end2end.py` +- Deploy config: `vllm_omni/deploy/yue2.yaml` (auto-discovered for `model_type=yue2`) + +## Architecture + +| Component | Spec | +| --- | --- | +| Backbone | Qwen3-1.7B-class decoder, 28 layers, hidden 2048, GQA 16/8, q/k-norm, vocab 184,704, context 24,576 | +| MoT | NAR path (`nar_self_attn`/`nar_mlp` per layer) shares embed_tokens, final norm and lm_head with the AR path | +| Frames | Single codebook: one codec token per frame, 25 frames/s, span [151853, 184521) | +| Sampling | temperature/top-k/top-p + windowed repetition penalty, seeded; model-owned (phase-masked vocab) | +| Guidance | Off by default (1.0 full/melody, 1.01 off); not exercised | +| Acoustic | 32-step midpoint flow matching over 64-dim latents on the NAR path, chunked ~12k frames | +| VAE | `m-a-p/YuE2-Vae` decoder-only, tiled (1024 core + 16 halo frames), FP32 | +| Output | 48 kHz stereo, whole song delivered when the semantic request finishes | + +Generation is one or two requests on a single AR stage: the abc phase writes +the score, the semantic phase generates codec frames and, on its last step, +solves the ODE and decodes the whole song in-engine (weights load once; the +VAE loads at startup from `$YUE2_VAE` or the hub id). + +## Running + +```bash +export YUE2_VAE=/models/YuE2-Vae # or omit and let it resolve from the hub +python examples/offline_inference/yue2/end2end.py \ + --model /models/YuE2-3B \ + --style "A melancholic lo-fi hip-hop track at 85 BPM in F minor" \ + --lyrics "[Verse] +Walking down the empty street at midnight +[Chorus] +And I keep on walking" \ + --cot full --seed 831001 --max-frames 200 --output song.wav +``` + +`--max-frames 200` caps the song at 8 s (25 frames/s). The end token is +masked for the preset's first 200 steps, so a 200-frame budget always +truncates; a longer budget lets the song end on its own end token, and +`truncated` in the report line says which happened. Serving follows the +same rule: omit `max_new_tokens` (or set it above 200) for a natural +ending. + +## Notes + +- **Memory (RTX 4090, 24 GB, `gpu_memory_utilization: 0.70`):** the engine + reserves ~15.8 GiB after startup; with 4 concurrent requests pinned to the + 9000-frame cap (worst case), peak usage during the terminal NAR/VAE + finishing pass reaches ~20.9 GiB, leaving ~3 GiB of headroom. At 0.85 the + same workload OOMs inside the finishing pass, so the deploy yaml pins 0.70. +- **Known limitations:** 4 concurrent full-length requests are the verified + shape (`max_num_seqs: 4`); CUDA graph capture is a + follow-up; the abc phase runs eagerly + after prefill (its tokens are the product, not audio). +- Audio is not expected to match the upstream torch reference bit-for-bit + (fused vs eager kernels). Measured against the upstream torch reference + (bd90e4c, same style/lyrics/seed 831001, cot=off, 200 frames): prompt token + ids 63/63 bit-identical, and the sampled semantic stream agrees + bit-for-bit for the first 55 frames before kernel numerics diverge — well + past the "first frames agree" acceptance bar of the MiniMax Music 3 + precedent. diff --git a/tests/assets/yue2/twinkle_golden.abc b/tests/assets/yue2/twinkle_golden.abc new file mode 100644 index 000000000..9d2ec067e --- /dev/null +++ b/tests/assets/yue2/twinkle_golden.abc @@ -0,0 +1,6 @@ +X:1 +T:Twinkle Twinkle Little Star (SheetSage2 transcription of the Wikimedia piano rendition, single voice) +M:4/4 +L:1/16 +K:C +C4C4G4G4|A4A4G8|F4F4E4E4|D4D4C8|G4G4F4F4|E4E4D8|G4G4F4F4|E4E4D8|C4C4G4G4|A4A4G8|F4F4E4E4|D4D4C8| diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index b66c930fc..6f004b6b6 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: 32, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29, - EnvironmentVariableCategory.MODEL_SPECIFIC: 74, + EnvironmentVariableCategory.MODEL_SPECIFIC: 75, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21, EnvironmentVariableCategory.INTERNAL: 6, } @@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC ) assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == { - ModelEnvironmentVariableDisposition.PROMOTE: 45, + ModelEnvironmentVariableDisposition.PROMOTE: 46, ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5, ModelEnvironmentVariableDisposition.EXTERNAL: 0, ModelEnvironmentVariableDisposition.INTERNALIZE: 17, diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py index 98b4184d9..6cf93526a 100644 --- a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py @@ -1048,3 +1048,83 @@ def test_minimax_music3_speech_invalid_field_values( online_client.send_audio_speech_http_request( {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message} ) + + +# ─── POST /v1/audio/speech · YuE2-3B ─── + +# Text-to-music like MiniMax Music 3, with one twist: the ABC score span is +# request-controlled via extra_params.cot/abc and the two must agree, and +# sampling is pinned by the checkpoint preset so sampling knobs are rejected. +_YUE2_SPEECH = [ + pytest.param( + OmniServerParams( + model="m-a-p/YuE2-3B", + stage_config_path=get_deploy_config_path("yue2.yaml"), + server_args=_SPEECH_SERVER_ARGS, + ), + id="yue2", + marks=hardware_marks(res={"cuda": "H100"}), + ), +] + +_YUE2_LYRICS = "一闪一闪亮晶晶\n满天都是小星星" +_YUE2_CAPTION = "Chinese heavy metal, distorted electric guitars, 140 BPM, key of C." +_YUE2_ABC = "X:1\nT:Twinkle\nM:4/4\nL:1/4\nK:C\nC C G G|A A G2|" +# 25 frames per second; the semantic preset bounds the frame budget. +_YUE2_MAX_FRAMES = 9000 + + +@pytest.mark.parametrize( + "overrides, err_message", + [ + pytest.param({"instructions": None}, ("instructions",), id="caption_missing"), + pytest.param({"instructions": " "}, ("instructions",), id="caption_blank"), + pytest.param({"input": " "}, ("input",), id="lyrics_blank"), + pytest.param({"voice": "alloy"}, ("voice",), id="voice_rejected"), + pytest.param({"speed": 1.5}, ("speed",), id="speed_rejected"), + pytest.param( + {"max_new_tokens": _YUE2_MAX_FRAMES + 1}, + ("max_new_tokens", str(_YUE2_MAX_FRAMES)), + id="max_new_tokens_over_cap", + ), + pytest.param({"max_new_tokens": 199}, ("max_new_tokens",), id="max_new_tokens_below_min"), + pytest.param( + {"extra_params": {"cot": "melody"}}, + ("abc",), + id="cot_melody_without_abc", + ), + pytest.param( + {"extra_params": {"abc": _YUE2_ABC}}, + ("cot",), + id="abc_without_cot", + ), + pytest.param( + {"extra_params": {"temperature": 0.3}}, + ("temperature",), + id="temperature_rejected", + ), + ], +) +@pytest.mark.parametrize("omni_server", _YUE2_SPEECH, indirect=True) +def test_yue2_speech_invalid_field_values( + omni_server: OmniServer, + online_client: OnlineOmniClient, + overrides: dict[str, object], + err_message: str | tuple[str, ...], +) -> None: + body: dict[str, Any] = { + "model": omni_server.model, + "input": _YUE2_LYRICS, + "instructions": _YUE2_CAPTION, + "seed": 1, + "max_new_tokens": 250, + "response_format": "wav", + } + for key, value in overrides.items(): + if value is None: + body.pop(key, None) + else: + body[key] = value + online_client.send_audio_speech_http_request( + {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message} + ) diff --git a/tests/e2e/offline_inference/test_yue2.py b/tests/e2e/offline_inference/test_yue2.py new file mode 100644 index 000000000..996f42e9f --- /dev/null +++ b/tests/e2e/offline_inference/test_yue2.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""End-to-end regressions for YuE2 offline WAV export.""" + +import os +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest +import soundfile as sf + +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import get_model_prefix +from vllm_omni.transformers_utils.repo_utils import hf_api + +pytestmark = [ + pytest.mark.full_model, + pytest.mark.tts, + *hardware_marks(res={"cuda": ["L4", "H100", "B200"]}, num_cards=1), +] + +REPO_ROOT = Path(__file__).resolve().parents[3] + + +@pytest.fixture(scope="module") +def vae_path() -> Path: + path = Path(os.environ.get("YUE2_VAE", get_model_prefix() + "m-a-p/YuE2-Vae")) + if not path.is_dir(): + path = Path(hf_api().snapshot_download(str(path), allow_patterns=["*.json", "*.safetensors"])) + return path + + +@pytest.mark.parametrize("max_frames,truncated", [(200, True), (9000, False)], ids=["budget", "natural-eos"]) +def test_offline_example_saves_audio(vae_path: Path, tmp_path: Path, max_frames: int, truncated: bool) -> None: + model = Path(os.environ.get("YUE2_MODEL_DIR", get_model_prefix() + "m-a-p/YuE2-3B")) + if not model.is_dir(): + model = Path( + hf_api().snapshot_download(str(model), allow_patterns=["*.json", "*.safetensors", "qwen.tiktoken"]) + ) + output = tmp_path / "song.wav" + env = os.environ.copy() + env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + env["PYTHONPATH"] = str(REPO_ROOT) + os.pathsep + env.get("PYTHONPATH", "") + completed = subprocess.run( + [ + sys.executable, + str(REPO_ROOT / "examples/offline_inference/yue2/end2end.py"), + "--model", + str(model), + "--vae", + str(vae_path), + "--style", + "gentle acoustic folk, clear English vocals, 90 BPM", + "--lyrics", + "[Verse]\nMorning light across the bay\nCarry all my dreams away\n[Outro]", + "--cot", + "off", + "--seed", + "831001", + "--max-frames", + str(max_frames), + "--output", + str(output), + ], + cwd=REPO_ROOT, + env=env, + capture_output=True, + text=True, + timeout=900, + ) + assert completed.returncode == 0, completed.stdout + completed.stderr + assert f"truncated={truncated}" in completed.stdout + audio, sample_rate = sf.read(output, always_2d=True) + assert sample_rate == 48000 and audio.shape[1] == 2 + assert np.isfinite(audio).all() and np.abs(audio).max() > 0 + duration = len(audio) / sample_rate + if truncated: + assert duration == pytest.approx(max_frames / 25, abs=0.1) + else: + assert 8 <= duration < max_frames / 25 diff --git a/tests/e2e/online_serving/test_yue2.py b/tests/e2e/online_serving/test_yue2.py new file mode 100644 index 000000000..35b25b1b0 --- /dev/null +++ b/tests/e2e/online_serving/test_yue2.py @@ -0,0 +1,237 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""E2E online tests for YuE2-3B text-to-music on ``/v1/audio/speech``. + +Scenario: heavy-metal "Twinkle Twinkle Little Star" — Chinese lyrics plus a +metal style caption plus an ABC score of the melody (``extra_params.cot = +"melody"``). The golden score is a SheetSage2 transcription of the Wikimedia +piano rendition, reduced to a single voice (``tests/assets/yue2``). + +Two programmatic content checks, both calibrated against a same-seed control +request with ``cot="off"`` (no ABC conditioning): + +- Style: CLAP zero-shot probability of "heavy metal" must clear 0.9. +- Melody: the generated song is transcribed back with SheetSage2 and its + pitch-class sequence is LCS-matched against the golden score (max over + major/minor variants and 12 transpositions, tolerating pedal-tone + insertions). SheetSage2's transcription of distorted metal vocals carries + roughly ±0.05 noise, so the bar sits at 0.80; the ABC-conditioned request + must also beat the control. + +The content checks run in a dedicated verification interpreter because +SheetSage2 wants a dependency set (mir_eval, pretty_midi, numpy<2) that must +not leak into the serving environment. They are enabled by three env vars: + +- ``YUE2_E2E_VERIFY_PYTHON``: python with sheetsage2 + CLAP deps installed +- ``YUE2_E2E_CLAP_DIR``: local laion/clap-htsat-unfused snapshot (~0.6 GB) +- ``YUE2_E2E_SHEETSAGE2_DIR``: local m-a-p/SheetSage2 snapshot (~0.25 GB, + plus its MERT-v2 backbone ~2.5 GB from the HF cache) + +Without them the test still asserts the structural floor (non-empty payload +above the byte floor) and passes on that alone, following the MiniMax Music 3 +precedent. + +Weekly rather than per-PR: each request is a 16 s song and the verification +checkpoints add ~5.5 GB of downloads. Validated on a single RTX 4090 +(24 GB); the hardware mark targets an H100 because the CI fleet has no +consumer cards. +""" + +import concurrent.futures +import io +import json +import os +import subprocess +import wave + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import httpx +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.media import get_asset_path +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path + +MODEL = "m-a-p/YuE2-3B" +DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 900.0 + +# 400 frames = 16 s of song. The byte floor sits far below the observed +# payload (~3.0 MB of 48 kHz stereo 16-bit) to leave room for an early EOS. +_FRAMES = 400 +# Concurrent pair stays short: it exists to exercise row re-indexing, not +# to double the wall time. +_CONCURRENT_FRAMES = 250 +_MIN_BYTES = 1_000_000 + +LYRICS = "一闪一闪亮晶晶\n满天都是小星星\n挂在天空放光明\n好像千万小眼睛" +CAPTION = ( + "Chinese heavy metal, distorted electric guitars, aggressive double-kick drums, powerful vocals, 140 BPM, key of C" +) +SEED = 831001 + +_VERIFY_PYTHON = os.environ.get("YUE2_E2E_VERIFY_PYTHON") +_CLAP_DIR = os.environ.get("YUE2_E2E_CLAP_DIR") +_SHEETSAGE2_DIR = os.environ.get("YUE2_E2E_SHEETSAGE2_DIR") +_CONTENT_CHECKS_ENABLED = all([_VERIFY_PYTHON, _CLAP_DIR, _SHEETSAGE2_DIR]) +_VERIFY_HELPER = os.path.join(os.path.dirname(__file__), "yue2_metal_verify.py") + + +tts_server_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=get_deploy_config_path("yue2.yaml"), + server_args=["--disable-log-stats"], + ), + id="yue2", + ) +] + + +def _http_speech(omni_server, *, seed: int, frames: int) -> bytes: + """Raw-HTTP speech request returning (and structurally checking) the WAV. + + The OpenAI SDK client swallows response headers, so the truncation flag is + cross-checked here: ``X-Audio-Truncated`` must agree with the audio — the + full frame budget (a 25 frames/s song of exactly frames/25 seconds) means + truncated, a shorter song means a natural end token. + """ + response = httpx.post( + f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech", + json={ + "model": omni_server.model, + "input": LYRICS, + "instructions": CAPTION, + "seed": seed, + "max_new_tokens": frames, + "stream": False, + "response_format": "wav", + "extra_params": {"cot": "off"}, + }, + timeout=DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + ) + assert response.status_code == 200, response.text[:500] + truncated = response.headers.get("X-Audio-Truncated") + assert truncated in ("true", "false"), f"X-Audio-Truncated missing/invalid: {truncated!r}" + audio = response.content + with wave.open(io.BytesIO(audio)) as wav: + assert wav.getnchannels() == 2 + assert wav.getframerate() == 48000 + duration = wav.getnframes() / wav.getframerate() + hit_budget = duration >= frames / 25 - 0.5 + assert truncated == ("true" if hit_budget else "false"), ( + f"X-Audio-Truncated={truncated} but duration={duration:.2f}s vs {frames / 25:.1f}s budget" + ) + return audio + + +@pytest.mark.slow +@pytest.mark.tts +@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1}) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_yue2_metal_twinkle_001(omni_server, openai_client, tmp_path) -> None: + """Heavy-metal Twinkle: ABC-conditioned request vs same-seed control. + + Deploy Setting: yue2.yaml + Input Modal: text (lyrics) + text (style caption) + ABC score (extra_params) + Output Modal: audio, 48 kHz stereo + Input Setting: stream=False, seed pinned, 400 frames + """ + golden_abc = get_asset_path("yue2/twinkle_golden.abc").read_text() + + def _request(extra_params: dict) -> bytes: + responses = openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": LYRICS, + "instructions": CAPTION, + "seed": SEED, + "max_new_tokens": _FRAMES, + "stream": False, + "response_format": "wav", + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_BYTES, + # Sung metal vocals are not whisper-transcribable; skip the + # full_model transcript check (our content checks do the work). + "transcript_expected_text": "", + "extra_params": extra_params, + } + ) + audio = responses[0].audio_bytes + assert audio, "empty audio payload" + return audio + + metal_path = tmp_path / "metal.wav" + control_path = tmp_path / "control.wav" + metal_bytes = _request({"cot": "melody", "abc": golden_abc}) + control_bytes = _request({"cot": "off"}) + metal_path.write_bytes(metal_bytes) + control_path.write_bytes(control_bytes) + + # CI runs without the verification assets, so the structural checks below + # are the only guard against the two real bugs this PR fixed during review: + # mono/half-speed output (WAV header + duration bounds) and the int32 + # crash under concurrent row re-indexing (two simultaneous requests). + # Duration bounds: the end token is masked for the first 200 steps + # (min_tokens), and the 400-frame budget caps the song, so a healthy + # request lands in [8, 16.5] s; the half-speed bug produced ~32 s. + for label, audio in (("metal", metal_bytes), ("control", control_bytes)): + with wave.open(io.BytesIO(audio)) as wav: + assert wav.getnchannels() == 2, f"{label}: {wav.getnchannels()} channels, want stereo" + assert wav.getframerate() == 48000, f"{label}: {wav.getframerate()} Hz, want 48 kHz" + duration = wav.getnframes() / wav.getframerate() + assert _FRAMES / 50 <= duration <= _FRAMES / 25 + 0.5, ( + f"{label}: {duration:.2f}s outside [{_FRAMES / 50}, {_FRAMES / 25 + 0.5}]s" + ) + + # A pinned-seed request must flag truncation consistently with the audio + # (see _http_speech); the OpenAI SDK response does not expose headers. + _http_speech(omni_server, seed=SEED, frames=_FRAMES) + + # Two concurrent requests with different seeds: both must succeed and + # produce different songs (the model-owned sampler re-indexes rows across + # concurrent requests; a dtype mismatch here crashed the engine before). + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + futures = [pool.submit(_http_speech, omni_server, seed=SEED + i, frames=_CONCURRENT_FRAMES) for i in range(2)] + pair = [f.result() for f in futures] + assert pair[0] != pair[1], "different seeds produced identical audio" + + # Without the verification assets the structural assertions above + # (non-empty payload, byte floor) are the whole test — pass on them. + if not _CONTENT_CHECKS_ENABLED: + return + assert _VERIFY_PYTHON is not None + assert _CLAP_DIR is not None + assert _SHEETSAGE2_DIR is not None + + proc = subprocess.run( + [ + _VERIFY_PYTHON, + _VERIFY_HELPER, + "--clap-dir", + _CLAP_DIR, + "--sheetsage2-dir", + _SHEETSAGE2_DIR, + "--work-dir", + str(tmp_path), + str(metal_path), + str(control_path), + ], + capture_output=True, + text=True, + timeout=1800, + ) + assert proc.returncode == 0, f"verification helper failed:\n{proc.stderr}" + report = json.loads(proc.stdout.strip().splitlines()[-1]) + + metal_prob = report["metal_probability"] + assert metal_prob >= 0.9, f"style check: P(heavy metal)={metal_prob:.3f} < 0.9" + + melody_score = report["melody_score"] + control_score = report["control_score"] + assert melody_score >= 0.80, f"melody check: score={melody_score:.3f} < 0.80" + assert melody_score > control_score, ( + f"calibration: ABC-conditioned {melody_score:.3f} not above control {control_score:.3f}" + ) diff --git a/tests/e2e/online_serving/yue2_metal_verify.py b/tests/e2e/online_serving/yue2_metal_verify.py new file mode 100644 index 000000000..8c7553e73 --- /dev/null +++ b/tests/e2e/online_serving/yue2_metal_verify.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Content-gate helper for the YuE2 heavy-metal Twinkle e2e test. + +Run as ``__main__`` by a dedicated verification interpreter (see +``YUE2_E2E_VERIFY_PYTHON`` in test_yue2.py), NOT by the serving environment: +SheetSage2 wants its own dependency set (mir_eval, pretty_midi, numpy<2), +which must not leak into the serving venv. + +Usage: + python yue2_metal_verify.py --clap-dir DIR --sheetsage2-dir DIR \ + metal.wav control.wav + +Prints one JSON object on stdout: + {"metal_probability": float, "melody_score": float, "control_score": float} + +Gates (applied by the caller): +- metal_probability >= 0.9 (CLAP zero-shot heavy-metal style) +- melody_score >= 0.80 and melody_score > control_score + (SheetSage2 re-transcription, pitch-class LCS vs the golden score, max over + major/minor variants and 12 transpositions, pedal-tone tolerant) +""" + +import argparse +import json +import re +import wave +from typing import Any + +import numpy as np + +_NOTE_BASE = {"C": 0, "D": 2, "E": 4, "F": 5, "G": 7, "A": 9, "B": 11} +_GOLDEN_MAJOR = "C4C4G4G4|A4A4G8|F4F4E4E4|D4D4C8|G4G4F4F4|E4E4D8|G4G4F4F4|E4E4D8|C4C4G4G4|A4A4G8|F4F4E4E4|D4D4C8|" +# Metal arrangements commonly minor-ize the melody; accept both modes. +_GOLDEN_MINOR = ( + "C4C4G4G4|_A4_A4G8|F4F4_E4_E4|D4D4C8|G4G4F4F4|_E4_E4D8|G4G4F4F4|_E4_E4D8|C4C4G4G4|_A4_A4G8|F4F4_E4_E4|D4D4C8|" +) + +_STYLE_CANDIDATES = [ + "heavy metal music with distorted guitars and aggressive drums", + "a soft children's lullaby", + "classical piano music", + "pop music", +] + + +def _abc_note_to_midi(token): + m = re.match(r"([\^_=]*)([A-Ga-g])([,']*)(\d*)", token) + if not m: + return None + acc, letter, oct_marks, _ = m.groups() + pc = _NOTE_BASE[letter.upper()] + acc.count("^") - acc.count("_") + octave = 5 if letter.isupper() else 6 + octave += oct_marks.count("'") - oct_marks.count(",") + return 12 * octave + pc + + +def _collapse(seq): + out: list[Any] = [] + for x in seq: + if not out or x != out[-1]: + out.append(x) + return out + + +def _golden_events(body): + return _collapse([_abc_note_to_midi(t) for t in re.findall(r"[\^_=]*[A-G]\d*", body)]) + + +def _melody_events(score_abc): + """Note events of the voice carrying the most notes in a SheetSage2 score.""" + voices: dict[str | None, list[int]] = {} + current = None + for line in score_abc.splitlines(): + line = line.strip() + if line.startswith("V:"): + parts = line.split() + current = parts[1] if len(parts) > 1 else None + continue + if current is None or not line or line.startswith(("%", "X:", "T:", "M:", "L:", "Q:", "K:")): + continue + for tok in re.findall(r"[\^_=]*[A-Ga-g][,']*\d*", line.replace("|", " ")): + midi = _abc_note_to_midi(tok) + if midi is not None: + voices.setdefault(current, []).append(midi) + return _collapse(max(voices.values(), key=len, default=[])) + + +def _lcs(a, b): + dp = [0] * (len(b) + 1) + for x in a: + prev = 0 + for j, y in enumerate(b, 1): + tmp = dp[j] + dp[j] = prev + 1 if x == y else max(dp[j], dp[j - 1]) + prev = tmp + return dp[-1] + + +def melody_score(score_abc): + gen = [m % 12 for m in _melody_events(score_abc)] + if len(gen) < 4: + return 0.0 + best = 0 + for body in (_GOLDEN_MAJOR, _GOLDEN_MINOR): + gold = [m % 12 for m in _golden_events(body)] + best = max(best, max(_lcs(gen, [(p + r) % 12 for p in gold]) for r in range(12))) + return best / len(_golden_events(_GOLDEN_MAJOR)) + + +def _wav_to_mono(path): + with wave.open(path) as wf: + sr = wf.getframerate() + channels = wf.getnchannels() + frames = wf.readframes(wf.getnframes()) + data = np.frombuffer(frames, dtype=np.int16).astype(np.float32) / 32768.0 + if channels > 1: + data = data.reshape(-1, channels).mean(axis=1) + return data, sr + + +def metal_probability(wav_path, clap_dir): + import torch + from transformers import ClapModel, ClapProcessor + + model = ClapModel.from_pretrained(clap_dir) + proc = ClapProcessor.from_pretrained(clap_dir) + audio, sr = _wav_to_mono(wav_path) + if sr != 48000: + n = int(len(audio) * 48000 / sr) + audio = np.interp(np.linspace(0, len(audio), n), np.arange(len(audio)), audio).astype(np.float32) + audio = audio[: 48000 * 10] + inputs = proc( + text=_STYLE_CANDIDATES, + audios=[audio] * len(_STYLE_CANDIDATES), + sampling_rate=48000, + return_tensors="pt", + padding=True, + ) + with torch.no_grad(): + probs = model(**inputs).logits_per_audio[0].softmax(dim=-1) + return float(probs[0]) + + +def transcribe_melody_abc(wav_path, sheetsage2_dir, output_dir): + import torch + from transformers import AutoModel + + audio, sr = _wav_to_mono(wav_path) + model = AutoModel.from_pretrained(sheetsage2_dir, trust_remote_code=True) + result = model.transcribe( + torch.from_numpy(audio[None, :]), + sampling_rate=sr, + output_dir=str(output_dir), + melody_only=True, + dtype="fp32", + ) + abc = result.get("abc") + if not abc: + raise RuntimeError(f"SheetSage2 produced no melody ABC: {result.get('abc_error')}") + return abc + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--clap-dir", required=True) + parser.add_argument("--sheetsage2-dir", required=True) + parser.add_argument("metal_wav") + parser.add_argument("control_wav") + parser.add_argument("--work-dir", default=None) + args = parser.parse_args() + + work = args.work_dir or "." + report = {"metal_probability": metal_probability(args.metal_wav, args.clap_dir)} + report["melody_score"] = melody_score( + transcribe_melody_abc(args.metal_wav, args.sheetsage2_dir, f"{work}/metal_score") + ) + report["control_score"] = melody_score( + transcribe_melody_abc(args.control_wav, args.sheetsage2_dir, f"{work}/control_score") + ) + print(json.dumps(report)) + + +if __name__ == "__main__": + main() diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 6504359e4..dc49a7e21 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -1105,6 +1105,38 @@ class TestTTSMethods: assert isinstance(audio_obj, CreateAudio) assert audio_obj.speed == 1.0 + @pytest.mark.asyncio + async def test_audio_synthesis_error_flag_raises_tts_generation_error( + self, + speech_server, + mocker: MockerFixture, + ): + """A model-flagged synthesis failure surfaces as a non-retryable + TTSGenerationError, so create_speech answers 500 instead of shipping a + zero-length WAV or crashing on an unpacked Response.""" + + async def mock_generate(): + yield create_mock_audio_output_for_test() + + mocker.patch.object( + speech_server, + "_prepare_speech_generation", + new=mocker.AsyncMock(return_value=("speech-synth-err", mock_generate(), {})), + ) + adapter = mocker.MagicMock() + adapter.collect_response_metadata = lambda _audio_output, collect: collect.__setitem__( + "audio_synthesis_error", True + ) + mocker.patch.object(speech_server, "_get_tts_adapter", return_value=adapter) + + with pytest.raises(TTSGenerationError, match="failed to synthesize audio") as exc_info: + await speech_server._generate_audio_bytes( + OpenAICreateSpeechRequest(input="Hello"), + collect={}, + ) + + assert exc_info.value.retryable is False + def test_is_tts_detection_with_tts_stage(self, mocker: MockerFixture): """Test TTS model detection when TTS stage exists.""" mock_engine_client = mocker.MagicMock() diff --git a/tests/entrypoints/openai_api/test_tts_detection.py b/tests/entrypoints/openai_api/test_tts_detection.py index 3ef55f673..ca27c6e37 100644 --- a/tests/entrypoints/openai_api/test_tts_detection.py +++ b/tests/entrypoints/openai_api/test_tts_detection.py @@ -183,6 +183,7 @@ _PIPELINE_STAGES = [ "token2wav", "tts", "encoder", + "yue2", ] _STAGES = [*_PIPELINE_STAGES, None, "vae", "not_a_real_stage"] diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index ef9712eb9..79d8f610d 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/tests/worker/test_gpu_ar_model_runner.py @@ -15,6 +15,7 @@ import numpy as np import pytest import torch from vllm.sampling_params import SamplingParams +from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.worker import gpu_input_batch from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch @@ -103,6 +104,8 @@ def test_post_sample_talker_mtp_uses_current_sample_and_hidden() -> None: ) assert received["req_ids"] == ["ready"] + assert isinstance(received["input_ids"], torch.Tensor) + assert isinstance(received["hidden_states"], torch.Tensor) assert received["input_ids"].tolist() == [101] assert received["hidden_states"].tolist() == [[1.0, 2.0]] assert received["req_infos"] == [{"duplex": {"data_plane": True}}] @@ -139,6 +142,8 @@ def test_post_sample_talker_mtp_uses_gpu_token_with_async_scheduling() -> None: ) assert received["req_ids"] == ["ready"] + assert isinstance(received["input_ids"], torch.Tensor) + assert isinstance(received["hidden_states"], torch.Tensor) assert received["input_ids"].tolist() == [101] assert received["hidden_states"].tolist() == [[1.0, 2.0]] assert received["req_infos"] == [{"duplex": {"data_plane": True}}] @@ -275,7 +280,12 @@ def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatc ) input_batch.sampling_metadata = input_batch._make_sampling_metadata() - received = [] + received: list[SamplingMetadata] = [] + + def sample(logits: torch.Tensor, metadata: SamplingMetadata) -> str: + received.append(metadata) + return "model-sampler" + runner = object.__new__(GPUARModelRunner) runner._omni_cache_policy = ModelCachePolicy(needs_full_hidden_states=True) runner._pooler_payload_include_hidden_flag = True @@ -283,7 +293,7 @@ def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatc runner.model = SimpleNamespace( prefer_model_sampler=True, skips_model_sampler_output_token_history=True, - sample=lambda logits, metadata: received.append(metadata) or "model-sampler", + sample=sample, ) runner.sampler = SimpleNamespace() logits = torch.zeros((1, 4)) @@ -1670,7 +1680,6 @@ class TestMergeModelKvTransferMetadata: # ...while the engine-shared originals are untouched. assert original["r1"]["custom_metadata"] is original_meta assert original_meta == {"a": 1} - assert "talker_codes" not in original["r1"]["custom_metadata"] def test_merge_without_model_meta_passes_entry_through(self): class _Model: @@ -1713,7 +1722,7 @@ class TestDownstreamPayloadMemoization: return runner def test_missing_marker_defaults_true_without_memoizing(self): - stages = {"r1": None} + stages: dict[str, int | None] = {"r1": None} runner = self._runner(stages) assert runner._request_needs_downstream_stage_payload("r1") is True @@ -1861,6 +1870,7 @@ class TestPreferModelSamplerNoneFallback: "minicpmo_4_5", "minimax_music3", "nemotron_voicechat", + "yue2", # Always supplies SamplerOutput, including empty/prefill steps. } assert declarers == expected, ( "The set of models declaring `prefer_model_sampler` changed:\n" diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index aa94a6879..69d9306a7 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -189,6 +189,7 @@ _MODEL_PROMOTE = ( "VLLM_OMNI_SEEDVR2_SHARDED_CLIP_PIXELS", "VLLM_OMNI_SEEDVR2_SHARDED_FRAME_PIXELS", "VLLM_OMNI_SENSENOVA_PAGED_DECODE", + "YUE2_VAE", # VAE checkpoint selection belongs in typed model configuration. ) _MODEL_REQUEST_SCOPE = ( diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 686ef1bf0..f016b7f7d 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -125,6 +125,7 @@ from vllm_omni.model_executor.models.step_audio2.pipeline import ( from vllm_omni.model_executor.models.voxcpm2.pipeline import VOXCPM2_PIPELINE from vllm_omni.model_executor.models.voxtral_tts.pipeline import VOXTRAL_TTS_PIPELINE from vllm_omni.model_executor.models.wan2_2.pipeline import WAN2_2_TI2V_PIPELINE +from vllm_omni.model_executor.models.yue2.pipeline import YUE2_PIPELINE logger = init_logger(__name__) @@ -208,6 +209,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "voxcpm2": VOXCPM2_PIPELINE, "voxtral_tts": VOXTRAL_TTS_PIPELINE, "wan2_2_ti2v": WAN2_2_TI2V_PIPELINE, + "yue2": YUE2_PIPELINE, } diff --git a/vllm_omni/deploy/yue2.yaml b/vllm_omni/deploy/yue2.yaml new file mode 100644 index 000000000..502d950c5 --- /dev/null +++ b/vllm_omni/deploy/yue2.yaml @@ -0,0 +1,55 @@ +# YuE2-3B deploy: single-stage native-AR text-to-music (~3.6B bf16 + VAE). +# Target 1x 24 GB card; verified on a single RTX 4090. +# +# The architectures pin is required: the checkpoint's config.json names +# "YuE2ForCausalLM" with an auto_map that only makes sense under +# trust_remote_code, so the deploy pins vLLM to this package's model class. +# model_type stays "yue2" so the registered Yue2Config is used; no +# trust_remote_code is needed (the registered config/model win over the +# checkpoint's auto_map). +pipeline: yue2 +async_chunk: false +dtype: bfloat16 + +stages: + - stage_id: 0 + max_num_seqs: 4 + # Measured on a 24 GB RTX 4090 with 4 concurrent requests pinned to the + # 9000-frame cap (worst case): the engine reserves ~15.8 GiB at 0.70 and + # peak usage during the terminal NAR/VAE finishing pass reaches 20.9 GiB, + # leaving ~3 GiB of headroom. 0.85 reserves ~19.4 GiB and the same + # workload OOMs inside the finishing pass. + gpu_memory_utilization: 0.70 + # Phase 1 keeps the backbone eager: the model-owned sampler runs after the + # forward, and the terminal NAR/VAE pass runs inside the finishing step. + # CUDA-graph capture of the standard Qwen3 backbone is a perf follow-up. + enforce_eager: true + # Prefix caching stays off: the omni prefix-cache save snapshots the + # step's mm payload before the model-owned sampler runs, but YuE2 + # appends the decoded waveform inside sample() (the terminal NAR/VAE + # pass), so the cached payload would ship empty audio. Re-enable once + # the runner saves post-sample. + enable_prefix_caching: false + max_num_batched_tokens: 24576 + max_model_len: 24576 + devices: "0" + # The checkpoint ships no HF tokenizer files (only qwen.tiktoken, which + # the offline driver reads itself); every prompt arrives as token ids, + # so the engine builds no tokenizer. Placeholder tokenizer dir follows + # the indextts2 precedent. + skip_tokenizer_init: true + tokenizer: gpt2 + hf_overrides: + architectures: ["Yue2ForCausalLM"] + default_sampling_params: + # Placeholder values: the model's own sampler overrides everything from + # extra_args. Their presence turns on has_sampling_extra_args, which is + # what threads the per-request yue2_* keys into the model's forward. + extra_args: + yue2_phase: "semantic" + yue2_seed: 831001 + yue2_max_audio_frames: 200 + detokenize: false + # Engine-side backstop; the model stops on its own end token first. + max_tokens: 9000 + temperature: 1.0 diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index b6c495a64..3a9520574 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -2275,6 +2275,22 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if ts is not None: collect["word_timestamps"] = ts + # Let the adapter fold engine-side metadata (e.g. YuE2's + # meta.truncated) into ``collect`` for response headers. + if collect is not None and (adapter := self._get_tts_adapter()) is not None: + adapter.collect_response_metadata(audio_output, collect) + + # A model can flag a per-request synthesis failure through the + # adapter (e.g. YuE2's terminal NAR/VAE pass OOMing on one + # request); answer 500 instead of shipping a zero-length WAV. + # Raising (not returning a Response) keeps this function's + # tuple contract; create_speech maps TTSGenerationError to 500. + if collect is not None and collect.get("audio_synthesis_error"): + raise TTSGenerationError( + "The model failed to synthesize audio for this request", + retryable=False, + ) + audio_tensor = audio_output[audio_key] sr_raw = audio_output.get("sr", 24000) sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw @@ -2715,6 +2731,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): "(use the WebSocket streaming path for long transcripts)", len(ts_json), ) + if collect.get("audio_truncated") is not None: + headers["X-Audio-Truncated"] = "true" if collect["audio_truncated"] else "false" return Response(content=audio_bytes, media_type=media_type, headers=headers) except asyncio.CancelledError: diff --git a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py index f98fa87df..d6e7ff186 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py @@ -124,6 +124,7 @@ from vllm_omni.entrypoints.openai.tts_adapters import ( # noqa: E402,F401 step_audio2, voxcpm2, voxtral, + yue2, ) __all__ = [ diff --git a/vllm_omni/entrypoints/openai/tts_adapters/base.py b/vllm_omni/entrypoints/openai/tts_adapters/base.py index b2297d826..6ad7dcdcc 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/base.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/base.py @@ -320,6 +320,15 @@ class TTSModelAdapter(ABC): to unrelated TTS models. """ + def collect_response_metadata(self, audio_output: Mapping[str, Any], collect: dict) -> None: + """Fold engine-side metadata from the mm payload into ``collect``. + + Non-streaming only. The orchestrator surfaces select ``collect`` + entries as response headers (precedent: ``word_timestamps`` -> + ``X-Word-Timestamps``, ``audio_truncated`` -> ``X-Audio-Truncated``); + adapters opt in by setting the keys it knows. Default: no-op. + """ + async def warmup(self) -> None: return diff --git a/vllm_omni/entrypoints/openai/tts_adapters/yue2.py b/vllm_omni/entrypoints/openai/tts_adapters/yue2.py new file mode 100644 index 000000000..bcacc34e7 --- /dev/null +++ b/vllm_omni/entrypoints/openai/tts_adapters/yue2.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""YuE2-3B serving adapter for ``/v1/audio/speech``. + +This is a text-to-music model on a speech endpoint, so most of the speech +contract does not apply. ``input`` carries the lyrics and ``instructions`` the +style caption (genre, instrumentation, tempo, mood). There is no speaker to +select, no reference audio, and no sampling knobs: the request-local preset is +pinned by the checkpoint's reference config. Model-specific controls travel in +``extra_params``: + +- ``cot``: ``"off"`` (default) | ``"melody"`` | ``"full"`` — how the ABC + score span in the prompt is handled. +- ``abc``: an ABC score as text, required for ``cot=melody|full``. Serving v1 + does not run the model-generated ABC phase (that is two sequential engine + requests; use the offline example for it). + +Unsupported parameters are rejected rather than silently ignored, following +the MiniMax Music 3 precedent. +""" + +import secrets +from collections.abc import Mapping +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch + +from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter +from vllm_omni.entrypoints.openai.tts_adapters.base import ( + ARTTSAdapter, + PreparedRequest, + resolve_stage_model_path, +) +from vllm_omni.model_executor.models.yue2.yue2 import ( + CONTEXT, + FRAMES_PER_SECOND, + KEY_MAX_AUDIO_FRAMES, + KEY_MIN_TOKENS, + KEY_PENALTY_WINDOW, + KEY_PHASE, + KEY_PREFIX_IDS, + KEY_REPETITION_PENALTY, + KEY_SEED, + KEY_SKIP_SYNTHESIS, + KEY_TEMPERATURE, + KEY_TOP_K, + KEY_TOP_P, + SEMANTIC_SAMPLING, + STOP_TOKEN_IDS, +) +from vllm_omni.tokenizers.yue2_prompt import semantic_prefix_ids + +if TYPE_CHECKING: + from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest + +_COT_MODES = ("off", "melody", "full") + +# Without max_new_tokens the budget is the preset ceiling, so the model's +# own end token sets the song length (upstream reference behavior). +_DEFAULT_MAX_FRAMES = int(SEMANTIC_SAMPLING["max_tokens"]) + +# Sampling is fixed by the checkpoint's reference preset. Accepting these +# would imply a control the model does not honor. +_UNSUPPORTED_SAMPLING_PARAMS = ("temperature", "top_p", "top_k", "repetition_penalty") + + +@register_tts_adapter +class Yue2Adapter(ARTTSAdapter): + """Lyrics plus a style caption in, one 48 kHz stereo song out.""" + + name = "yue2" + stage_keys = frozenset({"yue2"}) + model_archs = frozenset({"Yue2ForCausalLM"}) + max_new_tokens_min = int(SEMANTIC_SAMPLING["min_tokens"]) + max_new_tokens_max = int(SEMANTIC_SAMPLING["max_tokens"]) + + def __init__(self, ctx: Any) -> None: + super().__init__(ctx) + self._cached_tokenizer: Any = None + + def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: + if not request.input or not request.input.strip(): + return "YuE2 requires non-empty 'input' carrying the lyrics" + + instructions = request.instructions + if not isinstance(instructions, str) or not instructions.strip(): + return ( + "YuE2 requires 'instructions' describing the music " + "(genre, instrumentation, tempo, mood); it is what decides the arrangement" + ) + + extra: dict[str, Any] = request.extra_params or {} + cot = extra.get("cot", "off") + if cot not in _COT_MODES: + return f"extra_params.cot must be one of {_COT_MODES}; got {cot!r}" + abc = extra.get("abc") + if cot == "off": + if abc is not None: + return "extra_params.abc requires cot=melody|full" + elif not isinstance(abc, str) or not abc.strip(): + return ( + f"cot={cot} requires extra_params.abc carrying an ABC score. " + "Model-generated ABC is two engine requests and is not served " + "in v1; use examples/offline_inference/yue2 for that." + ) + + 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}" + if request.max_new_tokens > self.max_new_tokens_max: + return ( + f"'max_new_tokens' counts audio frames at {FRAMES_PER_SECOND} per second " + f"and cannot exceed {self.max_new_tokens_max}" + ) + + voice = (request.voice or "").strip().lower() + if voice not in ("", "default"): + return ( + "YuE2 has no speaker to select; the vocal comes from " + f"'instructions'. Omit 'voice' or pass 'default'; got {request.voice!r}" + ) + if request.ref_audio is not None or request.ref_text is not None: + return "YuE2 does not support reference-audio conditioning" + if request.language is not None: + return "YuE2 has no language tag; the language follows the lyrics" + if request.task_type is not None: + return "YuE2 does not support 'task_type'" + if request.speed is not None and float(request.speed) != 1.0: + return "YuE2 only supports speed=1.0; put the tempo in 'instructions', e.g. 'at 120 BPM'" + if request.is_streaming(): + return ( + "YuE2 does not support streaming: the song is decoded in one " + "terminal NAR/VAE pass after AR decoding finishes. Use stream=false." + ) + + rejected = sorted(key for key in _UNSUPPORTED_SAMPLING_PARAMS if extra.get(key) is not None) + if rejected: + return "YuE2 has fixed sampling and does not accept: " + ", ".join(rejected) + return None + + async def warmup(self) -> None: + # Resolve (and possibly download) qwen.tiktoken at startup: build() + # runs on the request path and must not block on a hub download, and a + # bad path should surface before the first request. + self._tokenizer() + + async def build( + self, + request: "OpenAICreateSpeechRequest", + sampling_params_list: list, + has_inline_ref_audio: bool, + ) -> PreparedRequest: + del sampling_params_list, has_inline_ref_audio + extra: dict[str, Any] = request.extra_params or {} + cot = extra.get("cot", "off") + encode = self._tokenizer().encode + abc_ids = encode(extra["abc"]) if cot != "off" else None + # validate() has already rejected blank lyrics/caption. + instructions = request.instructions or "" + lyrics = request.input or "" + prefix = semantic_prefix_ids(encode, instructions, lyrics, cot, abc_ids=abc_ids) + max_frames = int(request.max_new_tokens or _DEFAULT_MAX_FRAMES) + if len(prefix) + max_frames > CONTEXT: + return_err = ( + f"YuE2 prompt is {len(prefix)} tokens and the {max_frames}-frame budget " + f"exceeds the {CONTEXT}-token context. Shorten the lyrics/style or " + "lower max_new_tokens." + ) + raise ValueError(return_err) + return PreparedRequest( + prompt={"prompt_token_ids": prefix}, + tts_params={"max_audio_frames": [max_frames]}, + model_type=self.name, + ) + + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: "OpenAICreateSpeechRequest", + prompt: dict[str, Any] | None = None, + request_id: str | None = None, + ) -> list: + import copy + + sampling_params_list = copy.deepcopy(sampling_params_list) + params = sampling_params_list[0] + max_frames = int(request.max_new_tokens or _DEFAULT_MAX_FRAMES) + # A service must not give every caller the same song; without an + # explicit seed, draw one per request. + seed = request.seed if request.seed is not None else secrets.randbelow(2**31) + assert prompt is not None # the engine always passes the scheduled prompt + params.extra_args = { + **(params.extra_args or {}), + KEY_PHASE: "semantic", + KEY_SEED: seed, + KEY_TEMPERATURE: SEMANTIC_SAMPLING["temperature"], + KEY_TOP_P: SEMANTIC_SAMPLING["top_p"], + KEY_TOP_K: SEMANTIC_SAMPLING["top_k"], + KEY_REPETITION_PENALTY: SEMANTIC_SAMPLING["repetition_penalty"], + KEY_PENALTY_WINDOW: SEMANTIC_SAMPLING["penalty_window"], + KEY_MIN_TOKENS: SEMANTIC_SAMPLING["min_tokens"], + KEY_MAX_AUDIO_FRAMES: max_frames, + KEY_SKIP_SYNTHESIS: False, + # Full prompt ids: under a KV prefix-cache hit the engine schedules + # only the uncached tail, so the model cannot rebuild its NAR + # conditioning prefix from the scheduled tokens alone. + KEY_PREFIX_IDS: list(prompt["prompt_token_ids"]), + } + # The model needs one decode step beyond the last frame to deliver the + # terminal NAR/VAE pass, so do not cap the engine at exactly max_frames. + params.max_tokens = max_frames + 1 + params.stop_token_ids = list(STOP_TOKEN_IDS) + params.detokenize = False + return sampling_params_list + + @staticmethod + def _last_meta_flag(audio_output: Mapping[str, Any], name: str) -> Any: + """Read a per-request 0/1 meta flag in all three wire shapes. + + The runner flattens nested meta.* to dotted keys (flatten_payload) + and the wire payload is tensor-only (_ensure_tensor_values), so by + serving time the flag is usually a 0-d "0"/"1" tensor under a dotted + key; accept the nested/offline shapes too. + """ + flags = audio_output.get(f"meta.{name}") + if flags is None: + meta = audio_output.get("meta") + if isinstance(meta, dict): + flags = meta.get(name) + if isinstance(flags, torch.Tensor): + flags = flags.reshape(-1)[-1].item() if flags.numel() else None + elif isinstance(flags, (list, tuple)): + flags = flags[-1] if flags else None + return flags + + def collect_response_metadata(self, audio_output: Mapping[str, Any], collect: dict) -> None: + """Surface meta.truncated / meta.error as response headers.""" + truncated = self._last_meta_flag(audio_output, "truncated") + if truncated is not None: + collect["audio_truncated"] = truncated is True or str(truncated) == "1" + # meta.error=1 marks a failed NAR/VAE pass (empty clip); serving turns + # it into a 500 via this generic key. Absent/0 means success. + error = self._last_meta_flag(audio_output, "error") + if error is not None and (error is True or str(error) == "1"): + collect["audio_synthesis_error"] = True + + def _tokenizer(self) -> Any: + """The checkpoint-native tiktoken BPE, loaded once per process.""" + if self._cached_tokenizer is None: + from vllm_omni.tokenizers.yue2_tokenizer import YuE2TextTokenizer + + model_path = resolve_stage_model_path(self.ctx.engine_client) + if model_path is None: + raise RuntimeError("YuE2 tokenizer needs a resolvable stage model path") + merge_file = Path(model_path).expanduser() / "qwen.tiktoken" + if not merge_file.is_file(): + # model_path is an HF repo id: pull just the merge file into + # the hub cache (vae.py does the same for the VAE snapshot). + from vllm_omni.transformers_utils.repo_utils import hf_api + + merge_file = Path(hf_api().hf_hub_download(model_path, "qwen.tiktoken")) + self._cached_tokenizer = YuE2TextTokenizer(merge_file) + return self._cached_tokenizer + + +__all__ = ["Yue2Adapter"] diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index a0a812343..ca0fba7f8 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -483,6 +483,12 @@ _OMNI_MODELS = { "acoustic", "MiniMaxMusic3AcousticForConditionalGeneration", ), + ## YuE2-3B (text-to-music; single-stage native-AR with terminal NAR/VAE) + "Yue2ForCausalLM": ( + "yue2", + "yue2", + "Yue2ForCausalLM", + ), } diff --git a/vllm_omni/model_executor/models/yue2/__init__.py b/vllm_omni/model_executor/models/yue2/__init__.py new file mode 100644 index 000000000..9db16ee2d --- /dev/null +++ b/vllm_omni/model_executor/models/yue2/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""YuE2-3B text-to-music model package.""" + +from .pipeline import YUE2_PIPELINE + +__all__ = ["YUE2_PIPELINE"] diff --git a/vllm_omni/model_executor/models/yue2/pipeline.py b/vllm_omni/model_executor/models/yue2/pipeline.py new file mode 100644 index 000000000..0a1fa51bd --- /dev/null +++ b/vllm_omni/model_executor/models/yue2/pipeline.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""YuE2-3B pipeline topology. + +Single-stage native-AR text-to-music (gepard pattern). The whole +Mixture-of-Transformers checkpoint — AR path, NAR path, projection heads — +lives in one ``model.safetensors``, so both the autoregressive sampling and +the terminal flow-matching + VAE pass run inside one stage and no weights are +duplicated across stages. The VAE decoder is a separate repository +(``m-a-p/YuE2-Vae``), loaded in ``load_weights`` from ``$YUE2_VAE`` or the +default hub id. + +Sampling is model-owned. ``sampling_constraints`` pins ``detokenize`` and the +union of both phase end tokens; ``extra_args`` carry the phase and the +request-local sampling preset, and must be present in the stage's +``default_sampling_params`` so ``has_sampling_extra_args`` turns on and the +per-request args reach the model's forward. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +from .yue2 import STOP_TOKEN_IDS + +YUE2_PIPELINE = PipelineConfig( + model_type="yue2", + default_deploy_config_name="yue2.yaml", + model_arch="Yue2ForCausalLM", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="yue2", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=True, + final_output_type="audio", + owns_tokenizer=True, + engine_output_type="audio", + # Whole-song audio ships once, on the finishing step; the audio + # accumulator concatenates over time and a single tensor just + # passes through. + sampling_constraints={ + "detokenize": False, + "stop_token_ids": list(STOP_TOKEN_IDS), + }, + ), + ), +) + +__all__ = ["YUE2_PIPELINE"] diff --git a/vllm_omni/model_executor/models/yue2/yue2.py b/vllm_omni/model_executor/models/yue2/yue2.py new file mode 100644 index 000000000..02549cf84 --- /dev/null +++ b/vllm_omni/model_executor/models/yue2/yue2.py @@ -0,0 +1,1750 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""YuE2-3B text-to-music: single-stage native-AR model, ported from upstream +YuE2 (bd90e4c, Apache-2.0). + +One file per the single-stage layout (see auk/): constants, the AR model with +its model-owned sampler, the terminal NAR flow-matching pass, weight routing +and the tiled VAE decoder all live here because they all run inside stage 0. +The text tokenizer and token-id prompt construction live in +``vllm_omni/tokenizers/`` (driver-side only, never used by the model). +""" + +from __future__ import annotations + +import json +import math +import os +import time +from collections.abc import Callable, Iterable, Sequence +from dataclasses import dataclass, field +from numbers import Integral +from pathlib import Path +from typing import Any, Literal + +import torch +import torch.nn.functional as F +from torch import nn +from torch.nn.utils import weight_norm +from transformers import PretrainedConfig, PreTrainedModel +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.models.qwen3 import Qwen3Model +from vllm.model_executor.models.utils import AutoWeightsLoader, maybe_prefix +from vllm.v1.outputs import SamplerOutput + +from vllm_omni.model_executor.models.output_templates import OmniOutput + +# -------------------- token/protocol constants -------------------- + + +"""Token/protocol constants for YuE2-3B, pinned by the checkpoint. + +Values mirror the upstream ``yue2/protocol.py`` (bd90e4c) exactly; changing any +of them desynchronizes generation from the reference implementation. The model +is single-codebook: one codec token per latent frame, 25 frames per second, +so ``max_tokens`` in engine terms equals frames/25 seconds of music. +""" + +EOD = 151643 +ABC_START, ABC_END = 151847, 151848 +MUSIC_START, MUSIC_END = 151851, 151852 +CODEC_OFFSET, CODEC_SIZE = 151853, 32768 +LATENT_START, LATENT_END, LATENT_PAD = 184621, 184622, 184623 +VOCAB_SIZE, CONTEXT = 184704, 24576 + +# Engine-side stop ids: the union of both phase ends. The model's own sampler +# only ever emits the end token of the request's current phase, so a request +# never stops on the other phase's end by accident. +STOP_TOKEN_IDS = [ABC_END, MUSIC_END] + +SAMPLE_RATE = 48000 +AUDIO_CHANNELS = 2 +FRAMES_PER_SECOND = 25 +SAMPLES_PER_FRAME = 1920 +LATENT_DIM = 64 + +# Reference sampling presets (upstream yue2_generation_config.json). +ABC_SAMPLING = { + "temperature": 0.7, + "top_p": 0.9, + "top_k": 30, + "repetition_penalty": 1.005, + "penalty_window": 100, + "min_tokens": 32, + "max_tokens": 4096, +} +SEMANTIC_SAMPLING = { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 100, + "repetition_penalty": 1.2, + "penalty_window": 50, + "min_tokens": 200, + "max_tokens": 9000, +} + +ODE_STEPS = 32 +ODE_METHOD = "midpoint" +VAE_CORE_FRAMES = 1024 +VAE_HALO_FRAMES = 16 +DEFAULT_VAE_ID = "m-a-p/YuE2-Vae" + +# extra_args keys the model reads per request. +KEY_PHASE = "yue2_phase" # "abc" | "semantic" +KEY_TEMPERATURE = "yue2_temperature" +KEY_TOP_P = "yue2_top_p" +KEY_TOP_K = "yue2_top_k" +KEY_REPETITION_PENALTY = "yue2_repetition_penalty" +KEY_PENALTY_WINDOW = "yue2_penalty_window" +KEY_MIN_TOKENS = "yue2_min_tokens" +KEY_MAX_AUDIO_FRAMES = "yue2_max_audio_frames" +KEY_SEED = "yue2_seed" +KEY_SKIP_SYNTHESIS = "yue2_skip_synthesis" # abc phase: tokens only, no audio +# Full prompt token ids, shipped by the driver. The NAR conditioning needs the +# whole prefix, but under a KV prefix-cache hit the engine schedules only the +# uncached tail, so the scheduled input_ids slice cannot rebuild it. +KEY_PREFIX_IDS = "yue2_prefix_ids" + + +# -------------------- checkpoint weight routing -------------------- + + +"""AR/side routing for YuE2's single-file MoT checkpoint. + +The checkpoint interleaves AR-path, NAR-path and projection tensors under one +namespace. AR tensors go through the vLLM backbone loader unchanged; NAR and +projection tensors are hand-loaded, with ``model.layers.N.nar_*`` names +remapped onto the ``nar_layers.N.*`` module layout. +""" + +_SIDE_TOP_LEVEL = {"vae2llm", "llm2vae", "time_embedder"} + + +def partition_checkpoint_weights( + weights: Iterable[tuple[str, torch.Tensor]], +) -> tuple[list[tuple[str, torch.Tensor]], list[tuple[str, torch.Tensor]]]: + """Split checkpoint tensors into (backbone AR pairs, remapped side pairs).""" + ar: list[tuple[str, torch.Tensor]] = [] + side: list[tuple[str, torch.Tensor]] = [] + for name, tensor in weights: + if name == "latent_pos_embed.pe": + continue # deterministic sinusoid, rebuilt in __init__ + if ".nar_" in name or name.split(".", 1)[0] in _SIDE_TOP_LEVEL: + new = name + if name.startswith("model.layers."): + rest = name[len("model.layers.") :] + layer_no, _, tail = rest.partition(".") + if tail.startswith("nar_"): + tail = tail[len("nar_") :] + new = f"nar_layers.{layer_no}.{tail}" + side.append((new, tensor)) + continue + ar.append((name, tensor)) + return ar, side + + +# -------------------- model-owned sampler -------------------- + + +"""Model-owned sampling for YuE2-3B, ported from upstream ``sampling.py``. + +The model claims ``prefer_model_sampler`` and reproduces the reference +request-local arithmetic exactly: phase masking (abc = text vocabulary + its +end; semantic = codec span + its end), windowed repetition penalty over the +request's own history, then temperature/top-k/top-p with a seeded multinomial. +Default CFG is off (guidance 1.0 for full/melody, 1.01 for off); the upstream +torch backend treats those the same, so Phase 1 samples a single row. +""" + + +def window_penalty(logits: torch.Tensor, recent_ids: list[int], penalty: float) -> torch.Tensor: + """Upweight/downweight ids seen in the window, upstream arithmetic.""" + if penalty == 1.0 or not recent_ids: + return logits + recent = torch.as_tensor(recent_ids, dtype=torch.long, device=logits.device) + if logits.dim() > 1: + recent = recent.reshape(1, -1) + freq = torch.zeros_like(logits) + freq.scatter_add_(-1, recent, torch.ones_like(recent, dtype=logits.dtype)) + alpha = penalty**freq + return torch.where(logits < 0, logits * alpha, logits / alpha) + + +def distribution( + logits: torch.Tensor, + *, + temperature: float, + top_p: float, + top_k: int, + repetition_penalty: float, + penalty_window: int, + history: list[int], + step: int, + min_tokens: int, + phase: str, +) -> torch.Tensor: + """Masked/penalized/shaped scores for one row (float32 path). + + Accepts either one row ``[vocab]`` (what the model's sampler sees) or a + batch ``[rows, vocab]`` and returns the same shape it was given. + """ + single = logits.dim() == 1 + logits = logits.unsqueeze(0) if single else logits + scores = logits.float().clone() + end = ABC_END if phase == "abc" else MUSIC_END + allowed = torch.full_like(scores, float("-inf")) + if phase == "abc": + allowed[..., :EOD] = 0 + else: + allowed[..., CODEC_OFFSET : CODEC_OFFSET + CODEC_SIZE] = 0 + allowed[..., end] = 0 + scores = scores + allowed + if step < min_tokens: + scores[..., end] = -torch.inf + scores = window_penalty(scores, history[-penalty_window:], repetition_penalty) + if temperature == 0: + return scores.squeeze(0) if single else scores + if temperature != 1: + scores = scores / temperature + threshold = scores.topk(min(top_k, scores.shape[-1])).values[..., -1, None] + scores = scores.masked_fill(scores < threshold, -torch.inf) + if top_p < 1: + values, indices = scores.sort(descending=True) + probabilities = values.softmax(-1) + removed = probabilities.cumsum(-1) - probabilities > top_p + removed[..., :1] = False + values = values.masked_fill(removed, -torch.inf) + scores = values.scatter(-1, indices, values) + return scores.squeeze(0) if single else scores + + +def sample_row(scores: torch.Tensor, generator: torch.Generator, *, greedy: bool = False) -> int: + """Draw one token id from prepared scores with the request's generator.""" + if not torch.isfinite(scores).any(): + raise RuntimeError("sampling scores are all -inf; check phase masking") + if greedy: + return int(scores.argmax().item()) + probabilities = scores.softmax(-1) + next_id = torch.multinomial(probabilities, 1, generator=generator) + return int(next_id.item()) + + +# -------------------- NAR flow-matching pass -------------------- + + +"""Acoustic flow matching for YuE2-3B, ported from upstream ``nar.py``. + +One AR prefill per original chunk, then 32 midpoint ODE steps that walk only +the NAR path. The AR prefill re-uses the vLLM backbone's own projections +(fused ``qkv_proj`` split by head counts, ``q_norm``/``k_norm``, the layer's +``o_proj``/``mlp``), with RoPE and SDPA computed here exactly like the +reference: NAR positions attend the full AR prefix plus all NAR positions +bidirectionally, which a per-chunk prefill + concatenated K/V reproduces. +""" + + +def chunk_ranges(frames: int, prefix_tokens: int, context: int = CONTEXT): + size = min((context - prefix_tokens - 3) // 2, CONTEXT) + if frames < 1 or size < 1: + raise ValueError("Empty codec or prefix leaves no acoustic context") + return [(a, min(a + size, frames)) for a in range(0, frames, size)] + + +@dataclass +class Chunk: + ar_tokens: list[int] + noise: torch.Tensor + + +def song_chunks(prefix: Sequence[int], codec: Sequence[int], seed: int, context: int = CONTEXT): + """Draw the whole-song noise once, then cut it at the historical boundaries.""" + prefix = [int(v) for v in prefix] + codec = [int(v) for v in codec] + if min(prefix) < 0 or not codec or min(codec) < 0 or max(codec) >= CODEC_SIZE: + raise ValueError("Token IDs are outside their allowed vocabulary") + ranges = chunk_ranges(len(codec), len(prefix), int(context)) + generator = torch.Generator(device="cpu").manual_seed(int(seed)) + noise = torch.randn((len(codec), 64), dtype=torch.float32, device="cpu", generator=generator) + return [ + Chunk( + prefix + [value + CODEC_OFFSET for value in codec[a:b]] + [MUSIC_END], + noise[a:b], + ) + for a, b in ranges + ] + + +def _linear(module, x): + out = module(x) + return out[0] if isinstance(out, tuple) else out + + +def _attention(q, k, v, *, causal: bool, query_chunk_size: int | None = None): + """SDPA over [tokens, heads, dim]; query tiling only bounds temporaries.""" + if q.ndim != 3 or k.ndim != 3 or v.shape != k.shape or q.shape[-1] != k.shape[-1]: + raise ValueError("Expected Q/K/V [tokens, heads, dim] with matching K/V") + if min(q.shape) < 1 or min(k.shape) < 1 or q.shape[1] % k.shape[1]: + raise ValueError("Invalid attention lengths or grouped-query head count") + if causal and len(q) != len(k): + raise ValueError("Causal prefill requires matching Q/K sequence lengths") + block = query_chunk_size or (len(q) if q.device.type == "cuda" else 256) + query = q.transpose(0, 1).unsqueeze(0) + key = k.transpose(0, 1).unsqueeze(0) + value = v.transpose(0, 1).unsqueeze(0) + grouped = query.shape[1] != key.shape[1] + outputs = [] + for start in range(0, len(q), block): + end = min(start + block, len(q)) + used_key = key[..., :end, :] if causal else key + used_value = value[..., :end, :] if causal else value + mask = None + if causal and start: + mask = torch.arange(end, device=q.device)[None, :] <= torch.arange(start, end, device=q.device)[:, None] + outputs.append( + F.scaled_dot_product_attention( + query[..., start:end, :], + used_key, + used_value, + attn_mask=mask, + is_causal=causal and start == 0, + enable_gqa=grouped, + ) + ) + return torch.cat(outputs, dim=-2)[0].transpose(0, 1) + + +class CachedNAR: + """One acoustic chunk: AR prefix K/V prefilled once, ODE walks the NAR path.""" + + def __init__(self, model, chunk: Chunk, query_chunk_size: int | None = None): + self.model = model + self.chunk = chunk + self.query_chunk_size = query_chunk_size + weight = next(model.vae2llm.parameters()) + self.device, self.dtype = weight.device, weight.dtype + if chunk.noise.ndim != 2 or chunk.noise.shape[1] != 64 or len(chunk.noise) < 1: + raise ValueError("Expected nonempty acoustic noise [frames,64]") + self.ar_length, self.nar_length = len(chunk.ar_tokens), len(chunk.noise) + 2 + if self.ar_length < 1 or max(chunk.ar_tokens) >= model.vocab_size: + raise ValueError("AR prefix is empty or outside the model vocabulary") + if self.ar_length + self.nar_length > model.max_position_embeddings: + raise ValueError("Acoustic chunk exceeds the model context") + positions = torch.arange(self.ar_length, self.ar_length + self.nar_length, device=self.device)[None] + self.cos, self.sin = model.rotary(positions) + local = torch.arange(self.nar_length, device=self.device).clamp(max=model.max_latent_frames - 1) + self.pos_emb = model.latent_pos_embed(local)[None] + self.cache: list[tuple[torch.Tensor, torch.Tensor]] = [] + self._prefill() + + def _prefill(self): + backbone = self.model.model + ids = torch.tensor([self.chunk.ar_tokens], dtype=torch.long, device=self.device) + positions = torch.arange(self.ar_length, device=self.device)[None] + cos, sin = self.model.rotary(positions) + x = backbone.embed_tokens(ids)[0] # [T, H] + for layer in backbone.layers: + # Upstream: layer.self_attn.project_qkv(layer.input_layernorm(x), ...) + # The layernorm is load-bearing: without it the cached AR K/V + # conditioning is garbage and the ODE latents blow up (clipped audio). + q, k, v = self.model.ar_project(layer, layer.input_layernorm(x), cos, sin) + self.cache.append((k, v)) + h = _attention(q, k, v, causal=True, query_chunk_size=self.query_chunk_size) + x = x + _linear(layer.self_attn.o_proj, h.flatten(1)) + x = x + _linear(layer.mlp, layer.post_attention_layernorm(x)) + + @torch.inference_mode() + def velocity(self, state: torch.Tensor, raw_t: float) -> torch.Tensor: + model = self.model + x_nar = F.pad(state, (0, 0, 1, 1)) + shifted = model.shift_t(raw_t, self.device, self.dtype) + x = model.vae2llm(x_nar[None]) + x = x + model.time_embedder(shifted.expand(self.nar_length))[None] + x = x + self.pos_emb + for layer, (ar_k, ar_v) in zip(model.nar_layers, self.cache): + q, k, v = layer.self_attn.project_qkv(layer.input_layernorm(x), self.cos, self.sin) + k, v = torch.cat((ar_k, k[0])), torch.cat((ar_v, v[0])) + h = _attention(q[0], k, v, causal=False, query_chunk_size=self.query_chunk_size) + x = x + _linear(layer.self_attn.o_proj, h.flatten(1)[None]) + x = x + layer.mlp(layer.pre_mlp_layernorm(x)) + return model.llm2vae(model.model.norm(x))[0, 1:-1] + + @torch.inference_mode() + def solve(self, steps: int = 32) -> torch.Tensor: + if isinstance(steps, bool) or not isinstance(steps, Integral) or steps < 1: + raise ValueError("steps must be a positive integer") + state = self.chunk.noise.to(device=self.device, dtype=self.dtype) + dt = 1.0 / steps + for step in range(steps): + t = 1.0 - step * dt + raw = torch.logit(torch.tensor(t, dtype=torch.float64, device="cpu")).clamp(-20, 20).item() + first = self.velocity(state, raw) + mid = state - first * (dt / 2) + raw_mid = torch.logit(torch.tensor(t - dt / 2, dtype=torch.float64, device="cpu")).clamp(-20, 20).item() + state = state - self.velocity(mid, raw_mid) * dt + result = state.float().cpu() + if not torch.isfinite(result).all(): + raise FloatingPointError("Acoustic flow matching produced non-finite latents") + return result + + def close(self): + self.cache.clear() + self.cos = self.sin = self.pos_emb = None + + +@torch.inference_mode() +def synthesize(model, prefix: Sequence[int], codec: Sequence[int], seed: int, steps: int = 32) -> torch.Tensor: + """CPU FP32 latents [frames,64] for the whole song, chunk by chunk.""" + if model.training: + raise ValueError("synthesize requires model.eval()") + output = [] + for chunk in song_chunks(prefix, codec, seed): + engine = CachedNAR(model, chunk) + try: + output.append(engine.solve(steps)) + finally: + engine.close() + del engine + return torch.cat(output, dim=0) + + +# -------------------- tiled VAE decoder -------------------- + + +"""Portable FP32 YuE2 Oobleck VAE with exact-boundary tiled decoding. + +Oobleck and SnakeBeta derived from stable-audio-tools a6ae0cdf8b2eb1567a4b42ceadddec3712d99d45. +Copyright (c) 2023 Stability AI; Copyright (c) 2022 NVIDIA CORPORATION. +MIT: see THIRD_PARTY_NOTICES.md shipped with this module/model repository. +""" + + +def checkpoint(function, *args, **kwargs): + from torch.utils.checkpoint import checkpoint as torch_checkpoint + + kwargs.setdefault("use_reentrant", False) + return torch_checkpoint(function, *args, **kwargs) + + +def WNConv1d(*args, **kwargs): + return weight_norm(nn.Conv1d(*args, **kwargs)) + + +def WNConvTranspose1d(*args, **kwargs): + return weight_norm(nn.ConvTranspose1d(*args, **kwargs)) + + +def snake_beta(x, alpha, beta): + return x + (1.0 / (beta + 0.000000001)) * torch.pow(torch.sin(x * alpha), 2) + + +class SnakeBeta(nn.Module): + def __init__( + self, + in_features, + alpha=1.0, + alpha_trainable=True, + alpha_logscale=True, + ): + super().__init__() + self.in_features = in_features + self.alpha_logscale = alpha_logscale + if self.alpha_logscale: + self.alpha = nn.Parameter(torch.zeros(in_features) * alpha) + self.beta = nn.Parameter(torch.zeros(in_features) * alpha) + else: + self.alpha = nn.Parameter(torch.ones(in_features) * alpha) + self.beta = nn.Parameter(torch.ones(in_features) * alpha) + self.alpha.requires_grad = alpha_trainable + self.beta.requires_grad = alpha_trainable + self.no_div_by_zero = 0.000000001 + + def forward(self, x): + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) + beta = self.beta.unsqueeze(0).unsqueeze(-1) + if self.alpha_logscale: + alpha = torch.exp(alpha) + beta = torch.exp(beta) + return snake_beta(x, alpha, beta) + + +def get_activation(activation: Literal["elu", "snake", "none"], channels=None) -> nn.Module: + if activation == "elu": + return nn.ELU() + if activation == "snake": + return SnakeBeta(channels) + if activation == "none": + return nn.Identity() + raise ValueError(f"Unknown activation {activation}") + + +class ResidualUnit(nn.Module): + def __init__(self, in_channels, out_channels, dilation, act_type): + super().__init__() + self.dilation = dilation + padding = (dilation * (7 - 1)) // 2 + self.layers = nn.Sequential( + get_activation(act_type, channels=out_channels), + WNConv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=7, + dilation=dilation, + padding=padding, + ), + get_activation(act_type, channels=out_channels), + WNConv1d( + in_channels=out_channels, + out_channels=out_channels, + kernel_size=1, + ), + ) + + def forward(self, x): + residual = x + if self.training: + x = checkpoint(self.layers, x) + else: + x = self.layers(x) + return x + residual + + +class EncoderBlock(nn.Module): + def __init__(self, in_channels, out_channels, stride, act_type): + super().__init__() + self.layers = nn.Sequential( + ResidualUnit(in_channels, in_channels, 1, act_type), + ResidualUnit(in_channels, in_channels, 3, act_type), + ResidualUnit(in_channels, in_channels, 9, act_type), + get_activation(act_type, channels=in_channels), + WNConv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + ), + ) + + def forward(self, x): + return self.layers(x) + + +class DecoderBlock(nn.Module): + def __init__(self, in_channels, out_channels, stride, act_type): + super().__init__() + upsample_layer = WNConvTranspose1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + ) + self.layers = nn.Sequential( + get_activation(act_type, channels=in_channels), + upsample_layer, + ResidualUnit(out_channels, out_channels, 1, act_type), + ResidualUnit(out_channels, out_channels, 3, act_type), + ResidualUnit(out_channels, out_channels, 9, act_type), + ) + + def forward(self, x): + return self.layers(x) + + +class OobleckEncoder(nn.Module): + def __init__( + self, + in_channels=2, + channels=128, + latent_dim=32, + c_mults=(1, 2, 4, 8), + strides=(2, 4, 8, 8), + use_snake=False, + antialias_activation=False, + ): + super().__init__() + if antialias_activation: + raise ValueError("The released encoder does not use antialias_activation") + self.in_channels = in_channels + c_mults = [1] + list(c_mults) + self.depth = len(c_mults) + layers = [ + WNConv1d( + in_channels=in_channels, + out_channels=c_mults[0] * channels, + kernel_size=7, + padding=3, + ) + ] + act_type = "snake" if use_snake else "elu" + for i in range(self.depth - 1): + layers.append( + EncoderBlock( + in_channels=c_mults[i] * channels, + out_channels=c_mults[i + 1] * channels, + stride=strides[i], + act_type=act_type, + ) + ) + layers.extend( + [ + get_activation(act_type, channels=c_mults[-1] * channels), + WNConv1d( + in_channels=c_mults[-1] * channels, + out_channels=latent_dim, + kernel_size=3, + padding=1, + ), + ] + ) + self.layers = nn.Sequential(*layers) + + def forward(self, x): + return self.layers(x) + + +class OobleckDecoder(nn.Module): + def __init__( + self, + out_channels=2, + channels=128, + latent_dim=32, + c_mults=(1, 2, 4, 8), + strides=(2, 4, 8, 8), + use_snake=False, + snake_type="vanilla", + antialias_activation=False, + use_nearest_upsample=False, + use_filter=False, + final_tanh=True, + ): + super().__init__() + if antialias_activation or use_nearest_upsample or use_filter: + raise ValueError("Unsupported option for the released decoder") + if use_snake and snake_type != "vanilla": + raise ValueError("The released decoder uses vanilla SnakeBeta") + self.out_channels = out_channels + c_mults = [1] + list(c_mults) + self.depth = len(c_mults) + layers = [ + WNConv1d( + in_channels=latent_dim, + out_channels=c_mults[-1] * channels, + kernel_size=7, + padding=3, + ) + ] + act_type = "snake" if use_snake else "elu" + for i in range(self.depth - 1, 0, -1): + layers.append( + DecoderBlock( + in_channels=c_mults[i] * channels, + out_channels=c_mults[i - 1] * channels, + stride=strides[i - 1], + act_type=act_type, + ) + ) + layers.extend( + [ + get_activation(act_type, channels=c_mults[0] * channels), + WNConv1d( + in_channels=c_mults[0] * channels, + out_channels=out_channels, + kernel_size=7, + padding=3, + bias=False, + ), + nn.Tanh() if final_tanh else nn.Identity(), + ] + ) + self.layers = nn.Sequential(*layers) + + def forward(self, x): + return self.layers(x) + + +class YuE2VAEConfig(PretrainedConfig): + """Configuration shared by YuE2-Vae and YuE2-Vae-legacy. + + Use ``standard`` for listening and ``legacy`` for the paper metric baseline. + """ + + model_type = "yue2_vae" + + _hf_fields = frozenset( + { + "model_type", + "architectures", + "auto_map", + "transformers_version", + "dtype", + "torch_dtype", + "return_dict", + "output_hidden_states", + "output_attentions", + "use_cache", + "tie_word_embeddings", + "torchscript", + "is_decoder", + "is_encoder_decoder", + "add_cross_attention", + "bos_token_id", + "eos_token_id", + "pad_token_id", + "decoder_start_token_id", + "attn_implementation", + } + ) + + def to_dict(self): + return { + key: value + for key, value in super().to_dict().items() + if key in self._hf_fields or key in self._inference_fields + } + + _inference_fields = frozenset( + [ + "encoder_config", + "decoder_config", + "sample_rate", + "latent_dim", + "downsampling_ratio", + "audio_channels", + "release_variant", + "decode_core_frames", + "decode_halo_frames", + ] + ) + + def __init__( + self, + encoder_config=None, + decoder_config=None, + sample_rate=48000, + latent_dim=64, + downsampling_ratio=1920, + audio_channels=2, + release_variant="standard", + decode_core_frames=1024, + decode_halo_frames=16, + **kwargs, + ): + kwargs = {key: value for key, value in kwargs.items() if key in self._hf_fields} + kwargs.setdefault("architectures", ["YuE2VAE"]) + kwargs.setdefault( + "auto_map", + { + "AutoConfig": "modeling_vae.YuE2VAEConfig", + "AutoModel": "modeling_vae.YuE2VAE", + }, + ) + super().__init__(**kwargs) + self.encoder_config = encoder_config or dict( + in_channels=2, + channels=64, + c_mults=[1, 2, 4, 8, 16, 32], + strides=[2, 2, 4, 4, 5, 6], + latent_dim=128, + use_snake=True, + ) + self.decoder_config = decoder_config or dict( + out_channels=2, + channels=64, + c_mults=[1, 2, 4, 8, 16, 32], + strides=[2, 2, 4, 4, 5, 6], + latent_dim=64, + use_snake=True, + snake_type="vanilla", + use_filter=False, + final_tanh=False, + ) + encoder_fields = { + "in_channels", + "channels", + "latent_dim", + "c_mults", + "strides", + "use_snake", + "antialias_activation", + } + decoder_fields = { + "out_channels", + "channels", + "latent_dim", + "c_mults", + "strides", + "use_snake", + "snake_type", + "antialias_activation", + "use_nearest_upsample", + "use_filter", + "final_tanh", + } + self.encoder_config = {key: value for key, value in self.encoder_config.items() if key in encoder_fields} + self.decoder_config = {key: value for key, value in self.decoder_config.items() if key in decoder_fields} + self.sample_rate = int(sample_rate) + self.latent_dim = int(latent_dim) + self.downsampling_ratio = int(downsampling_ratio) + self.audio_channels = int(audio_channels) + self.release_variant = release_variant + self.decode_core_frames = int(decode_core_frames) + self.decode_halo_frames = int(decode_halo_frames) + if self.decode_core_frames < 1 or self.decode_halo_frames < 0: + raise ValueError("Invalid VAE core/halo configuration") + if math.prod(self.decoder_config["strides"]) != self.downsampling_ratio: + raise ValueError("Decoder strides do not match downsampling_ratio") + if self.decoder_config["latent_dim"] != self.latent_dim: + raise ValueError("Decoder input channels do not match latent_dim") + + +def _dependency_interval(module, low, high): + """Inclusive input support of an output interval; no waveform blending.""" + if isinstance(module, (nn.Sequential, OobleckDecoder, DecoderBlock)): + layers = module if isinstance(module, nn.Sequential) else module.layers + for child in reversed(list(layers)): + low, high = _dependency_interval(child, low, high) + return low, high + if isinstance(module, ResidualUnit): + a, b = _dependency_interval(module.layers, low, high) + return min(a, low), max(b, high) + if isinstance(module, nn.ConvTranspose1d): + s, p, d, k = (module.stride[0], module.padding[0], module.dilation[0], module.kernel_size[0]) + return -(-(low + p - d * (k - 1)) // s), (high + p) // s + if isinstance(module, nn.Conv1d): + s, p, d, k = (module.stride[0], module.padding[0], module.dilation[0], module.kernel_size[0]) + return low * s - p, high * s - p + d * (k - 1) + if isinstance(module, (SnakeBeta, nn.ELU, nn.Identity, nn.Tanh)): + return low, high + raise TypeError(f"No audited support rule for {type(module).__name__}") + + +def _output_length(module, length): + if isinstance(module, (nn.Sequential, OobleckDecoder, DecoderBlock)): + layers = module if isinstance(module, nn.Sequential) else module.layers + for child in layers: + length = _output_length(child, length) + return length + if isinstance(module, nn.ConvTranspose1d): + return ( + (length - 1) * module.stride[0] + - 2 * module.padding[0] + + module.dilation[0] * (module.kernel_size[0] - 1) + + module.output_padding[0] + + 1 + ) + if isinstance(module, nn.Conv1d): + return (length + 2 * module.padding[0] - module.dilation[0] * (module.kernel_size[0] - 1) - 1) // module.stride[ + 0 + ] + 1 + if isinstance(module, (ResidualUnit, SnakeBeta, nn.ELU, nn.Identity, nn.Tanh)): + return length + raise TypeError(f"No audited length rule for {type(module).__name__}") + + +class YuE2VAE(PreTrainedModel): + """Strict EMA model; only the decoder needs to reside on an accelerator. + + ``decode_tiled`` preserves finite receptive-field context and writes cropped + cores to CPU. The mathematical waveform is the full decoder's waveform; + convolution kernel choices may cause small FP32 rounding differences. + """ + + config_class = YuE2VAEConfig + base_model_prefix = "" + main_input_name = "audio" + + def __init__(self, config, decoder_only=False): + super().__init__(config) + self.decoder_only = bool(decoder_only) + if not self.decoder_only: + self.encoder = OobleckEncoder(**config.encoder_config) + self.decoder = OobleckDecoder(**config.decoder_config) + self.eval().requires_grad_(False) + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path, + *model_args, + config=None, + decoder_only=False, + device="cpu", + torch_dtype=None, + dtype=None, + revision=None, + token=None, + cache_dir=None, + local_files_only=False, + force_download=False, + subfolder="", + **kwargs, + ): + """Load a local export or Hub repository, selecting decoder tensors. + + Complete exports use unprefixed EMA keys. ``decoder_only=True`` avoids + constructing the encoder or loading encoder tensors into RAM/GPU. + """ + from safetensors import safe_open + + requested_dtype = dtype if dtype is not None else torch_dtype + if requested_dtype not in (None, "auto", "float32", torch.float32): + raise ValueError("The validated VAE requires FP32; quantize the LM separately") + device_map = kwargs.pop("device_map", None) + if device_map is not None: + if device_map == "auto": + device = "cuda" if torch.cuda.is_available() else "cpu" + elif isinstance(device_map, str): + device = device_map + elif isinstance(device_map, dict) and set(device_map) == {""}: + device = device_map[""] + else: + raise ValueError("Use decoder_only=True and a single device for the VAE") + for name in ( + "trust_remote_code", + "low_cpu_mem_usage", + "_from_auto", + "_from_pipeline", + "_commit_hash", + "adapter_kwargs", + "_fast_init", + "weights_only", + "use_safetensors", + ): + kwargs.pop(name, None) + output_loading_info = kwargs.pop("output_loading_info", False) + if kwargs or model_args: + raise TypeError(f"Unsupported VAE loading options: {sorted(kwargs)}") + path = Path(pretrained_model_name_or_path).expanduser() + if not path.is_dir(): + from vllm_omni.transformers_utils.repo_utils import hf_api + + path = Path( + hf_api().snapshot_download( + str(pretrained_model_name_or_path), + revision=revision, + token=token, + cache_dir=cache_dir, + local_files_only=local_files_only, + force_download=force_download, + allow_patterns=[ + f"{subfolder + '/' if subfolder else ''}{pattern}" + for pattern in ("config.json", "*.safetensors", "*.safetensors.index.json") + ], + ) + ) + path = path / subfolder + if config is None: + config = YuE2VAEConfig.from_pretrained(path, local_files_only=True) + # vLLM may construct us under a BF16 default dtype. Widen before copying + # the FP32 checkpoint so load_state_dict does not round its values. + model = cls(config, decoder_only=decoder_only).float() + index = path / "model.safetensors.index.json" + if index.exists(): + mapping = json.loads(index.read_text())["weight_map"] + files = sorted({name for key, name in mapping.items() if not decoder_only or key.startswith("decoder.")}) + else: + files = ["model.safetensors"] + state = {} + for name in files: + with safe_open(path / name, framework="pt", device="cpu") as handle: + for key in handle.keys(): + if not decoder_only or key.startswith("decoder."): + if key in state: + raise ValueError(f"Duplicate VAE tensor: {key}") + state[key] = handle.get_tensor(key) + expected = set(model.state_dict()) + if set(state) != expected: + raise ValueError( + f"VAE tensor mismatch: missing={sorted(expected - set(state))}, " + f"unexpected={sorted(set(state) - expected)}" + ) + if any(value.dtype != torch.float32 for value in state.values()): + raise ValueError("VAE export contains tensors that are not FP32") + model.load_state_dict(state, strict=True) + model.to(device=device, dtype=torch.float32).eval().requires_grad_(False) + if output_loading_info: + return model, dict(missing_keys=[], unexpected_keys=[], mismatched_keys=[], error_msgs=[]) + return model + + def save_pretrained(self, save_directory, *args, **kwargs): + if self.decoder_only: + raise ValueError("Reload decoder_only=False to save a complete VAE repository") + if kwargs.get("safe_serialization", True) is False: + raise ValueError("YuE2 VAE release exports require safetensors") + return super().save_pretrained(save_directory, *args, **kwargs) + + @property + def decoder_device(self): + return next(self.decoder.parameters()).device + + def _latent(self, latent): + latent = torch.as_tensor(latent) + if latent.ndim != 3 or latent.shape[1] != self.config.latent_dim or latent.shape[0] < 1 or latent.shape[-1] < 1: + raise ValueError(f"Expected nonempty [B,{self.config.latent_dim},T] latents") + if not torch.isfinite(latent).all(): + raise ValueError("VAE latents contain non-finite values") + if next(self.decoder.parameters()).dtype != torch.float32: + raise ValueError("VAE decoder weights must remain FP32") + return latent + + @torch.inference_mode() + def encode(self, audio, sample=False, generator=None, return_info=False): + """Encode FP32 stereo audio; posterior mean by default. + + Set ``sample=True`` with a per-request ``torch.Generator`` to reproduce + stochastic posterior sampling. This audio VAE is not the unreleased + semantic audio tokenizer semantic tokenizer. + """ + if self.decoder_only: + raise RuntimeError("Encoder not loaded; reload with decoder_only=False") + audio = torch.as_tensor(audio) + if ( + audio.ndim != 3 + or audio.shape[1] != self.config.audio_channels + or audio.shape[-1] < self.config.downsampling_ratio + ): + raise ValueError("Expected audio [B,2,S] with at least one latent frame") + if not torch.isfinite(audio).all(): + raise ValueError("Audio contains non-finite values") + device = next(self.encoder.parameters()).device + pre = self.encoder(audio.to(device=device, dtype=torch.float32)) + mean, scale = pre.chunk(2, dim=1) + stdev = torch.nn.functional.softplus(scale) + 1e-4 + if sample: + noise = torch.randn(mean.shape, dtype=mean.dtype, device=device, generator=generator) + latent = noise * stdev + mean + else: + latent = mean + if return_info: + return latent, dict(mean=mean, scale=scale, stdev=stdev) + return latent + + @torch.inference_mode() + def decode(self, latent): + """Full waveform, FP32 [B,2,1920*T-64], without clipping.""" + latent = self._latent(latent) + with torch.autocast(device_type=self.decoder_device.type, enabled=False): + return self.decoder(latent.to(device=self.decoder_device, dtype=torch.float32)) + + def natural_output_length(self, frames): + if int(frames) < 1: + raise ValueError("frames must be positive") + return _output_length(self.decoder, int(frames)) + + def required_halo(self, core_frames=None): + core_frames = self.config.decode_core_frames if core_frames is None else core_frames + ratio = self.config.downsampling_ratio + low, high = _dependency_interval(self.decoder, 0, core_frames * ratio - 1) + return max(0, -low, high - core_frames + 1) + + @torch.inference_mode() + def decode_tiled( + self, + latent, + core_frames=None, + halo_frames=None, + output_device="cpu", + on_progress: Callable[[int, int], None] | None = None, + ): + """Decode bounded tiles, retaining exact cores with natural end length. + + Each crop has enough left/right context for every dependency. There is + no crossfade or boundary smoothing, and no zero padding of final audio. + CPU output prevents an entire song from accumulating on the GPU. + ``on_progress(completed, total)`` runs after each existing crop copy; + no extra synchronization is added. With a CUDA output device, queued + work may still be executing. Callback exceptions propagate. + """ + latent = self._latent(latent) + core_frames = self.config.decode_core_frames if core_frames is None else core_frames + halo_frames = self.config.decode_halo_frames if halo_frames is None else halo_frames + if not isinstance(core_frames, int) or core_frames < 1: + raise ValueError("core_frames must be a positive integer") + required = self.required_halo(core_frames) + if not isinstance(halo_frames, int) or halo_frames < required: + raise ValueError(f"halo_frames must be at least {required} for this decoder") + frames = latent.shape[-1] + ratio = self.config.downsampling_ratio + total = self.natural_output_length(frames) + audio = torch.empty( + (latent.shape[0], self.config.audio_channels, total), dtype=torch.float32, device=output_device + ) + tiles = (frames + core_frames - 1) // core_frames + for tile_index, start in enumerate(range(0, frames, core_frames)): + end = min(frames, start + core_frames) + left = max(0, start - halo_frames) + right = min(frames, end + halo_frames) + tile = self.decode(latent[..., left:right]) + out_start, out_end = start * ratio, min(end * ratio, total) + crop_start = (start - left) * ratio + crop = tile[..., crop_start : crop_start + out_end - out_start] + if crop.shape[-1] != out_end - out_start: + raise RuntimeError("VAE tile did not cover its requested output core") + audio[..., out_start:out_end].copy_(crop.to(output_device)) + del tile, crop + if on_progress is not None: + on_progress(tile_index + 1, tiles) + return audio + + def decode_audio(self, latent, chunked=True, **kwargs): + return self.decode_tiled(latent, **kwargs) if chunked else self.decode(latent) + + def forward(self, audio, sample=False, generator=None): + return self.decode(self.encode(audio, sample=sample, generator=generator)) + + +YuE2VAEConfig.register_for_auto_class("AutoConfig") +YuE2VAE.register_for_auto_class("AutoModel") + + +# -------------------- the model -------------------- + + +"""YuE2-3B text-to-music model for vllm-omni: single-stage native-AR. + +The checkpoint is one Mixture-of-Transformers: a Qwen3-shaped AR path (which +is exactly a Qwen3-1.7B-class backbone over an extended vocabulary) plus a +parallel NAR path (``nar_self_attn``/``nar_mlp`` per layer) that a 32-step +midpoint flow-matching solver walks over 64-dim VAE latents. Both paths and +the projection heads live in the same ``model.safetensors``, so a single +vLLM stage loads everything once and no weights are duplicated. + +Sampling is model-owned (``prefer_model_sampler``), reproducing the upstream +request-local arithmetic: per-phase vocabulary masking, windowed repetition +penalty, seeded multinomial, ``min_tokens`` end-token suppression. One song +is one or two engine requests — the abc phase (``cot=full|melody``) first, +then the semantic phase; the driver stitches them. When the semantic phase's +end token is drawn (or its frame budget is hit), the model solves the ODE and +decodes 48 kHz stereo audio inside the step and ships it as the request's +final multimodal payload, so the whole song arrives in one piece. +""" +logger = init_logger(__name__) + +__all__ = ["Yue2ForCausalLM"] + +DEFAULT_SEED = 831001 + + +class _RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps).to(x.dtype) * self.weight + + +class _RotaryEmbedding(nn.Module): + """Checkpoint-compatible RoPE: cos/sin over [B, T, head_dim/2].""" + + def __init__(self, head_dim: int, base: float = 1000000.0): + super().__init__() + self.head_dim = head_dim + self.base = base + self._inv_freq: torch.Tensor | None = None + + def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + device = position_ids.device + if self._inv_freq is None or self._inv_freq.device != device: + self._inv_freq = 1.0 / ( + self.base ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32, device=device) / self.head_dim) + ) + angles = position_ids.float().unsqueeze(-1) * self._inv_freq + return angles.cos(), angles.sin() + + +def _apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + half = x.shape[-1] // 2 + x1, x2 = x[..., :half], x[..., half:] + cos, sin = cos.to(x.dtype), sin.to(x.dtype) + return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1) + + +class _NARAttention(nn.Module): + """NAR-path attention with checkpoint-native separate projections.""" + + def __init__(self, config): + super().__init__() + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False) + self.q_norm = _RMSNorm(self.head_dim, config.rms_norm_eps) + self.k_norm = _RMSNorm(self.head_dim, config.rms_norm_eps) + + def project_qkv(self, x, cos, sin): + B, T, _ = x.shape + q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim) + k = self.k_proj(x).view(B, T, self.num_kv_heads, self.head_dim) + v = self.v_proj(x).view(B, T, self.num_kv_heads, self.head_dim) + q, k = self.q_norm(q), self.k_norm(k) + rc, rs = cos.unsqueeze(2), sin.unsqueeze(2) + return _apply_rotary(q, rc, rs), _apply_rotary(k, rc, rs), v + + +class _NARMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) + + def forward(self, x): + return self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class _NARLayer(nn.Module): + """One NAR path layer; checkpoint ``model.layers.N.nar_*`` remaps onto it.""" + + def __init__(self, config): + super().__init__() + self.input_layernorm = _RMSNorm(config.hidden_size, config.rms_norm_eps) + self.self_attn = _NARAttention(config) + self.pre_mlp_layernorm = _RMSNorm(config.hidden_size, config.rms_norm_eps) + self.mlp = _NARMLP(config) + + +class _TimestepEmbedder(nn.Module): + def __init__(self, hidden_size: int, frequency_embedding_size: int = 256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size), + ) + self.frequency_embedding_size = frequency_embedding_size + + def forward(self, t): + half = self.frequency_embedding_size // 2 + freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device, dtype=torch.float32) / half) + args = t.float().unsqueeze(-1) * freqs.unsqueeze(0) + emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + return self.mlp(emb.to(next(self.parameters()).dtype)) + + +class _AudioPositionEmbedding(nn.Module): + def __init__(self, max_frames: int, hidden_size: int): + super().__init__() + pe = torch.zeros(max_frames, hidden_size) + position = torch.arange(0, max_frames, dtype=torch.float32).unsqueeze(1) + div_term = torch.exp(torch.arange(0, hidden_size, 2, dtype=torch.float32) * (-math.log(10000.0) / hidden_size)) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + self.register_buffer("pe", pe, persistent=False) + + def forward(self, position_ids): + return self.pe[position_ids] + + +def _request_key(value: Any) -> str: + if isinstance(value, (list, tuple)): + return str(value[0]) if value else "" + return "" if value is None else str(value) + + +@dataclass +class _RowConstants: + """Per-request values fixed at admission (a prefill step, never replayed).""" + + request_id: str + phase: str + temperature: float + top_p: float + top_k: int + repetition_penalty: float + penalty_window: int + min_tokens: int + max_audio_frames: int + seed: int + skip_synthesis: bool + + +@dataclass +class _RequestState: + """Everything one song request accumulates across decode steps.""" + + request_id: str + constants: _RowConstants + prompt_tokens: int = 0 + prompt_len: int = 0 # full prompt length; a completing prefill row reaches it + prefix_ids: list[int] = field(default_factory=list) + history: list[int] = field(default_factory=list) + generator: torch.Generator | None = None + finished: bool = False + truncated: bool = False + + +class Yue2ForCausalLM(nn.Module): + """YuE2-3B: vLLM-native Qwen3 AR backbone + NAR flow-matching + VAE.""" + + have_multimodal_outputs = True + prefer_model_sampler = True + has_postprocess = False + # The song is decoded in-model and shipped via make_omni_output; per-step + # hidden states must NOT ride the pooler payload, or the output pipeline + # remaps the accumulated "hidden" rows to the audio modality key and the + # driver receives hidden states instead of the decoded waveform + # (single-stage precedent: minimax_music3 talker). + omni_pooler_payload_include_hidden: bool = False + # No downstream stage consumes prefix-cached hidden/mm tensors for YuE2; + # opting out keeps the merged prefix-cache view from rebuilding "hidden" + # payloads (KV-block prefix caching for cot=full stays on). + requires_full_prefix_cached_hidden_states: bool = False + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + # The NAR pass keeps every KV head on each rank (ar_project splits by + # the per-rank head count), so TP>1 builds a server whose every song + # fails after its full AR decode. Pipeline parallelism has the same + # problem: AR and NAR must live in one stage. Fail at startup instead. + tp = int(getattr(vllm_config.parallel_config, "tensor_parallel_size", 1) or 1) + pp = int(getattr(vllm_config.parallel_config, "pipeline_parallel_size", 1) or 1) + if tp > 1 or pp > 1: + raise ValueError( + f"Yue2 is single-GPU only: the NAR path is replicated, so " + f"tensor_parallel_size={tp} and pipeline_parallel_size={pp} are " + "unsupported; use 1/1." + ) + self.vllm_config = vllm_config + config = vllm_config.model_config.hf_config + self.config = config + hidden = int(config.hidden_size) + + self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")) + if getattr(config, "tie_word_embeddings", False): + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + int(config.vocab_size), + hidden, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(int(config.vocab_size)) + + self.nar_layers = nn.ModuleList([_NARLayer(config) for _ in range(int(config.num_hidden_layers))]) + self.vae2llm = nn.Linear(int(getattr(config, "latent_dim", LATENT_DIM)), hidden) + self.llm2vae = nn.Linear(hidden, int(getattr(config, "latent_dim", LATENT_DIM))) + self.time_embedder = _TimestepEmbedder(hidden) + self.latent_pos_embed = _AudioPositionEmbedding( + int(getattr(config, "max_latent_frames", config.max_position_embeddings)), + hidden, + ) + self.rotary_emb = _RotaryEmbedding(int(config.head_dim), float(config.rope_theta)) + + self.vocab_size = int(config.vocab_size) + self.max_position_embeddings = int(config.max_position_embeddings) + self.max_latent_frames = int(getattr(config, "max_latent_frames", self.max_position_embeddings)) + self._timestep_shift = float(getattr(config, "timestep_shift", 1.0)) + + self._states: dict[str, _RequestState] = {} + self._row_constants: dict[str, _RowConstants] = {} + self._audio_queue: list[tuple[str, torch.Tensor, bool, bool]] = [] + self._deferred_cleanup_ids: set[str] = set() + self._step_rows: list[tuple[str, int, int]] = [] # (req_id, computed, scheduled) + self._decode_t0: dict[str, float] = {} # first decode-step wall time, for AR tok/s + self._last_mm: dict[str, Any] | None = None # current step's make_omni_output payload + self._vae: YuE2VAE | None = None + self._vae_device: torch.device | None = None + + # ------------------------------------------------------------ sampling + + def shift_t(self, raw_t: float, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + t_sig = torch.sigmoid(torch.tensor(raw_t, dtype=dtype, device=device)) + shift = self._timestep_shift + return shift * t_sig / (1 + (shift - 1) * t_sig) + + def rotary(self, positions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return self.rotary_emb(positions) + + def ar_project(self, layer, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): + """AR-path Q/K/V from the vLLM backbone's fused projection.""" + attn = layer.self_attn + qkv = attn.qkv_proj(x) + if isinstance(qkv, tuple): + qkv = qkv[0] + num_heads, num_kv_heads, head_dim = attn.num_heads, attn.num_kv_heads, attn.head_dim + q, k, v = torch.split( + qkv, + [num_heads * head_dim, num_kv_heads * head_dim, num_kv_heads * head_dim], + dim=-1, + ) + T = x.shape[-2] + q = attn.q_norm(q.view(T, num_heads, head_dim)) + k = attn.k_norm(k.view(T, num_kv_heads, head_dim)) + rc, rs = cos.unsqueeze(2), sin.unsqueeze(2) + # q/k are [T, heads, head_dim] (3D); rc/rs are [1, T, 1, head_dim/2] + # (4D). Broadcasting them directly promotes q/k to 4D, which the NAR + # attention rejects. Add the batch axis explicitly, then drop it. + q = _apply_rotary(q.unsqueeze(0), rc, rs).squeeze(0) + k = _apply_rotary(k.unsqueeze(0), rc, rs).squeeze(0) + return q, k, v.view(T, num_kv_heads, head_dim) + + # ------------------------------------------------------------ weights + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + ar_pairs, remapped = partition_checkpoint_weights(weights) + missing, unexpected = self.load_state_dict(dict(remapped), strict=False) + missing = [k for k in missing if not k.startswith(("model.", "lm_head"))] + if missing or unexpected: + raise RuntimeError(f"YuE2 NAR/projection weight mismatch: missing={missing}, unexpected={unexpected}") + + loader = AutoWeightsLoader(self) + loaded = loader.load_weights(iter(ar_pairs)) + # Load the VAE now, not lazily at the first finishing request: a bad + # path or a failed download must surface at startup, and the decoder + # weights must sit on the GPU before vLLM's memory profiling sizes + # the KV cache. + device = f"cuda:{torch.accelerator.current_device_index()}" if torch.cuda.is_available() else "cpu" + self._vae_model(torch.device(device)) + # The NAR/projection modules were populated by hand above; + # DefaultModelLoader.track_weights_loading diffs every named + # parameter against the returned set, so the side keys must be + # reported too or the loader flags them as uninitialized. + return loaded | {name for name, _ in remapped} + + # ------------------------------------------------------------ hooks + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + return self.model.embed_tokens(input_ids) + + def prepare_runner_inputs( + self, + *, + req_ids: list[str], + num_computed_tokens: Any, + num_scheduled_tokens: Any, + input_ids: torch.Tensor, + positions: torch.Tensor, + **_: Any, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Bind rows to requests for the step; prompt tokens are captured in + ``forward``, where the state is created (see ``_capture_constants``). + Decode rows feed tokens the model itself sampled, so they never + contribute to the prompt prefix.""" + computed = [int(v) for v in num_computed_tokens] + scheduled = [int(v) for v in num_scheduled_tokens] + self._step_rows = [(str(rid), computed[i], scheduled[i]) for i, rid in enumerate(req_ids)] + # make_omni_output runs every step; without this reset the first step's + # payload would stay live forever and the _audio_queue fallback in + # _ship_audio could never trigger. + self._last_mm = None + return input_ids, positions + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: Any | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor: + del intermediate_tensors + self._capture_constants(kwargs, input_ids) + hidden = self.model( + input_ids=input_ids if inputs_embeds is None else None, + positions=positions, + inputs_embeds=inputs_embeds, + ) + if isinstance(hidden, tuple): + hidden = hidden[0] + self._flush_deferred_cleanup() + return hidden + + def _capture_constants(self, kwargs: dict[str, Any], input_ids: torch.Tensor | None) -> None: + """Create per-request state on the request's FIRST scheduled step. + + With KV prefix caching the first step may arrive with ``comp > 0`` + (a shared head was cached by an earlier request), so ``comp == 0`` + cannot be the trigger. The scheduled ``input_ids`` slice then holds + only the uncached tail, so the driver ships the full prompt ids in + the request's extra args (KEY_PREFIX_IDS); they are the source of + truth for the NAR conditioning prefix and for the prompt length the + sampler uses to tell a completing prefill row from a mid-chunk one. + """ + extra_args = kwargs.get("sampling_extra_args") + if extra_args is None or input_ids is None: + return + offset = 0 + for row, (req_id, comp, span) in enumerate(self._step_rows): + if row >= len(extra_args): + break + if req_id in self._states: + offset += span + continue + args = extra_args[row] or {} + prefix = args.get(KEY_PREFIX_IDS) + if isinstance(prefix, (list, tuple)) and prefix: + prefix_ids = [int(v) for v in prefix] + prompt_len = len(prefix_ids) + elif comp == 0 and span >= 1: + # Legacy path (no ids in extra args): only correct when the + # whole prompt is scheduled in one chunk with no cache hit. + prefix_ids = input_ids[offset : offset + span].tolist() + prompt_len = len(prefix_ids) + else: + logger.warning( + "YuE2 request %s arrived with %d cached tokens but no " + "yue2_prefix_ids in its extra args; cannot rebuild its " + "full prompt (prefix-cache hit on the first request?).", + req_id, + comp, + ) + offset += span + continue + phase = str(args.get(KEY_PHASE, "semantic")) + preset = ABC_SAMPLING if phase == "abc" else SEMANTIC_SAMPLING + seed = args.get(KEY_SEED, DEFAULT_SEED) + constants = _RowConstants( + request_id=req_id, + phase=phase, + temperature=float(args.get(KEY_TEMPERATURE, preset["temperature"])), + top_p=float(args.get(KEY_TOP_P, preset["top_p"])), + top_k=int(args.get(KEY_TOP_K, preset["top_k"])), + repetition_penalty=float(args.get(KEY_REPETITION_PENALTY, preset["repetition_penalty"])), + penalty_window=int(args.get(KEY_PENALTY_WINDOW, preset["penalty_window"])), + min_tokens=int(args.get(KEY_MIN_TOKENS, preset["min_tokens"])), + max_audio_frames=int(args.get(KEY_MAX_AUDIO_FRAMES, preset["max_tokens"])), + seed=int(seed), + skip_synthesis=bool(args.get(KEY_SKIP_SYNTHESIS, phase == "abc")), + ) + device = input_ids.device + generator = torch.Generator(device=device if device.type != "mps" else "cpu") + generator.manual_seed(constants.seed) + self._row_constants[req_id] = constants + self._states[req_id] = _RequestState( + request_id=req_id, + constants=constants, + prompt_tokens=prompt_len, + prompt_len=prompt_len, + prefix_ids=prefix_ids, + generator=generator, + ) + offset += span + + def compute_logits(self, hidden_states: torch.Tensor, sampling_metadata: Any = None) -> torch.Tensor: + return self.logits_processor(self.lm_head, hidden_states) + + # ------------------------------------------------------------ sampler + + def sample(self, logits: torch.Tensor, sampling_metadata: Any) -> SamplerOutput | None: + """Draw one token per decoding row with the request's own arithmetic.""" + del sampling_metadata + rows = int(logits.shape[0]) + # The runner's input_ids buffer is int32; long ids crash its scatter + # when concurrent requests re-index rows. + token_ids = torch.zeros((rows, 1), dtype=torch.int32, device=logits.device) + if not self._step_rows: + return SamplerOutput(sampled_token_ids=token_ids, logprobs_tensors=None) + + for row, (req_id, comp, span) in enumerate(self._step_rows): + if row >= rows: + break + state = self._states.get(req_id) + if state is None or state.finished: + continue + if comp + span < state.prompt_len: + # Mid-chunk prefill rows: their sampled token is discarded. + # (A completing prefill row reaches the full prompt length + # even when a prefix-cache hit left comp > 0.) + continue + constants = state.constants + if state.request_id not in self._decode_t0: + self._decode_t0[state.request_id] = time.perf_counter() + scores = distribution( + logits[row], + temperature=constants.temperature, + top_p=constants.top_p, + top_k=constants.top_k, + repetition_penalty=constants.repetition_penalty, + penalty_window=constants.penalty_window, + history=state.history, + step=len(state.history), + min_tokens=constants.min_tokens, + phase=constants.phase, + ) + token = sample_row(scores, state.generator, greedy=constants.temperature == 0) + end = ABC_END if constants.phase == "abc" else MUSIC_END + if token == end: + state.finished = True + state.truncated = False + self._finish_request_safely(state, hit_end=True) + token_ids[row, 0] = end + else: + state.history.append(token) + budget = constants.max_audio_frames + if constants.phase == "semantic" and len(state.history) >= budget: + state.finished = True + state.truncated = True + self._finish_request_safely(state, hit_end=False) + token_ids[row, 0] = end + else: + token_ids[row, 0] = token + return SamplerOutput(sampled_token_ids=token_ids, logprobs_tensors=None) + + # ------------------------------------------------------------ audio + + def _finish_request_safely(self, state: _RequestState, *, hit_end: bool) -> None: + """Keep an NAR/VAE failure from escaping ``sample``. + + The runner calls ``model.sample()`` with no error handling, so an + exception here (OOM in the ODE solve, non-finite VAE output) would + kill Stage-0 and take every live request down with it. Fail only + this request: ship an empty error-flagged clip that serving turns + into a 500. + """ + try: + self._finish_request(state, hit_end=hit_end) + except Exception: + logger.exception("YuE2 synthesis failed for req=%s; failing only this request", state.request_id) + self._ship_audio(state.request_id, torch.zeros((2, 0)), state.truncated, error=True) + + def _finish_request(self, state: _RequestState, *, hit_end: bool) -> None: + """Solve the ODE and decode the song; abc-phase requests skip this.""" + constants = state.constants + if constants.skip_synthesis or not state.history: + return + codec = [t - CODEC_OFFSET for t in state.history] + if min(codec) < 0 or max(codec) >= CODEC_SIZE: + raise RuntimeError("semantic history contains non-codec tokens") + t0 = self._decode_t0.pop(state.request_id, None) + t_nar0 = time.perf_counter() + latents = synthesize(self, state.prefix_ids, codec, constants.seed, steps=ODE_STEPS) + logger.info( + "YUE2_LATENTS frames=%d absmax=%.4f mean=%.5f std=%.4f", + latents.shape[0], + float(latents.abs().max()), + float(latents.mean()), + float(latents.std()), + ) + t_nar = time.perf_counter() - t_nar0 + t_vae0 = time.perf_counter() + audio = self._decode_latents(latents) + t_vae = time.perf_counter() - t_vae0 + t_ar = (t_nar0 - t0) if t0 is not None else float("nan") + logger.info( + "YUE2_TIMING req=%s phase=%s ar_tokens=%d ar_s=%.2f ar_tps=%.1f nar_s=%.2f vae_s=%.2f", + state.request_id, + constants.phase, + len(state.history), + t_ar, + len(state.history) / t_ar if t_ar else float("nan"), + t_nar, + t_vae, + ) + self._ship_audio(state.request_id, audio, state.truncated) + del hit_end + + def _ship_audio(self, req_id: str, audio: torch.Tensor, truncated: bool, *, error: bool = False) -> None: + """Append the finished song to the CURRENT step's mm payload in place. + + make_omni_output already ran for this step (it ships the sparse marker + with empty lists), and the runner assembles per-request payloads only + after sample(), so mutating the same dict here is picked up by the + sparse routing on the request's last step in the output batch. Leaving + the payload for a later step would drop it: the request is gone by + then (gepard_talker flushes at its last step for the same reason). + """ + mm = self._last_mm + if mm is None: + # No forward built a payload this step; fall back to the queue. + self._audio_queue.append((req_id, audio, truncated, error)) + return + mm["model_outputs"].append(audio) + mm["sr"].append(torch.tensor(SAMPLE_RATE, dtype=torch.int32)) + meta = mm["meta"] + meta["req_id"].append(req_id) + # Keep the flags ints, not strs: the wire payload is tensor-only + # (_ensure_tensor_values) and a string would be dropped before serving. + meta["truncated"].append(int(truncated)) + meta["error"].append(int(error)) + + def _vae_model(self, device: torch.device) -> YuE2VAE: + if self._vae is None: + vae_path = os.environ.get("YUE2_VAE", DEFAULT_VAE_ID) + logger.info("Loading YuE2 VAE decoder from %s", vae_path) + # Keep the VAE out of the module tree (__dict__, not an nn.Module + # attribute): the checkpoint has no VAE keys, so a registered + # submodule would trip DefaultModelLoader.track_weights_loading, + # and a model-wide dtype/device pass would change VAE numerics. + self.__dict__["_vae"] = YuE2VAE.from_pretrained(vae_path, decoder_only=True, device="cpu") + if self._vae_device != device: + self._vae.to(device) + self._vae_device = device + return self._vae + + def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor: + """[frames, 64] FP32 CPU latents -> channels-first stereo [2, samples].""" + device = f"cuda:{torch.accelerator.current_device_index()}" if torch.cuda.is_available() else "cpu" + model = self._vae_model(torch.device(device)) + z = latents.T.unsqueeze(0) # [1, 64, T] + with torch.inference_mode(): + audio = model.decode_tiled( + z, + core_frames=VAE_CORE_FRAMES, + halo_frames=VAE_HALO_FRAMES, + output_device="cpu", + ) + if not torch.isfinite(audio).all(): + raise RuntimeError("VAE produced non-finite audio") + logger.info("YUE2_AUDIO_PRECLAMP absmax=%.4f", float(audio.abs().max())) + # Channels-first [2, T], the same layout MiniMax Music 3 ships in + # model_outputs; serving's create_audio transposes it for soundfile. + return audio[0].float().clamp(-1, 1).contiguous() + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + return model_outputs + by_req: dict[str, torch.Tensor] = {} + truncated: dict[str, bool] = {} + failed: dict[str, bool] = {} + for req_id, audio, is_truncated, is_error in self._audio_queue: + by_req[req_id] = audio + truncated[req_id] = is_truncated + failed[req_id] = is_error + self._audio_queue.clear() + ready = list(by_req) + sr = torch.tensor(SAMPLE_RATE, dtype=torch.int32) + # The sparse marker rides EVERY step, even ones that decoded nothing + # (req_id=[] is the legal zero-audio shape): without it the runner + # falls back to the dense pooler payload, whose per-step hidden states + # land under this stage's "audio" key and the driver receives hidden + # states instead of the decoded waveform (gepard_talker precedent). + mm: dict[str, Any] = { + "model_outputs": [by_req[r] for r in ready], + "sr": [sr for _ in ready], + "meta": { + "req_id": ready, + "sparse_audio": ["1"], + "truncated": [int(truncated[r]) for r in ready], + "error": [int(failed[r]) for r in ready], + }, + } + # sample() runs after this each step and appends finished songs into + # this very dict (see _ship_audio); keep the handle for it. + self._last_mm = mm + return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs=mm) + + def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None: + # Fires before forward; defer the free so the in-flight step can read. + for rid in finished_req_ids: + self._deferred_cleanup_ids.add(_request_key(rid)) + + def _flush_deferred_cleanup(self) -> None: + for req_id in self._deferred_cleanup_ids: + state = self._states.pop(req_id, None) + if state is not None and not state.finished: + logger.warning( + "YuE2 request %s ended without an end token or its frame " + "budget (aborted or preempted); no audio was produced.", + req_id, + ) + self._row_constants.pop(req_id, None) + self._deferred_cleanup_ids.clear() diff --git a/vllm_omni/tokenizers/yue2_prompt.py b/vllm_omni/tokenizers/yue2_prompt.py new file mode 100644 index 000000000..97d7b7dea --- /dev/null +++ b/vllm_omni/tokenizers/yue2_prompt.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Prompt construction for YuE2-3B, mirroring upstream ``protocol.py``. + +The driver (offline example or, later, the serving adapter) builds token-id +prefixes here and submits them as ``prompt_token_ids``. Both generation phases +share one prefix head (instruction + tags + lyrics); the abc phase continues +from ``ABC_START`` and stops at ``ABC_END``, the semantic phase embeds the ABC +score (generated or user-supplied) and continues from ``MUSIC_START``. +""" + +from __future__ import annotations + +from vllm_omni.model_executor.models.yue2.yue2 import ( + ABC_END, + ABC_START, + CODEC_OFFSET, + EOD, + MUSIC_END, + MUSIC_START, +) + +INSTRUCTIONS = { + "off": "Generate music with codec tokens from the given conditions.", + "melody": ( + "Generate a melody-only ABC transcription without chord symbols, then " + "generate music with codec tokens from the given conditions." + ), + "full": ( + "Generate a chord-annotated ABC transcription, then generate music with codec tokens from the given conditions." + ), +} + + +def song_text(style: str, lyrics: str, cot: str) -> str: + return f"{INSTRUCTIONS[cot]}\n[Tags]\n{style}\n[Lyrics]\n{lyrics}\n" + + +def abc_prefix_ids(encode, style: str, lyrics: str, cot: str) -> list[int]: + """Prompt for the abc phase (cot=full/melody, no external ABC).""" + if cot == "off": + raise ValueError("cot=off skips the abc phase entirely") + return [EOD] + encode(song_text(style, lyrics, cot)) + [ABC_START] + + +def semantic_prefix_ids( + encode, + style: str, + lyrics: str, + cot: str, + abc_ids: list[int] | None = None, +) -> list[int]: + """Prompt for the semantic phase. + + ``abc_ids`` are the token ids of the (generated or user-supplied) ABC + score. ``cot=off`` passes ``abc_ids=None`` and skips the span, matching + upstream ``token_prefixes``. + """ + base = [EOD] + encode(song_text(style, lyrics, cot)) + if cot == "off": + return base + [ABC_START, ABC_END, MUSIC_START] + if abc_ids is None: + raise ValueError("cot=melody/full requires abc ids") + if any(not 0 <= t < EOD for t in abc_ids): + raise ValueError("ABC ids must stay inside the ordinary text vocabulary") + return base + [ABC_START] + list(abc_ids) + [ABC_END, MUSIC_START] + + +def abc_ids_from_generated(generated: list[int]) -> list[int]: + """ABC token ids generated by the abc phase (end token stripped).""" + ids = list(generated) + if ids and ids[-1] == ABC_END: + ids.pop() + if ABC_START in ids: + ids = ids[ids.index(ABC_START) + 1 :] + if any(t >= EOD for t in ids): + raise ValueError("abc phase produced tokens outside the text vocabulary") + return ids + + +def semantic_frames(generated: list[int]) -> list[int]: + """Codec frame values (0..CODEC_SIZE-1) from semantic-phase output.""" + + ids = list(generated) + if ids and ids[-1] == MUSIC_END: + ids.pop() + return [t - CODEC_OFFSET for t in ids] + + +__all__ = [ + "INSTRUCTIONS", + "abc_ids_from_generated", + "abc_prefix_ids", + "semantic_frames", + "semantic_prefix_ids", + "song_text", +] diff --git a/vllm_omni/tokenizers/yue2_tokenizer.py b/vllm_omni/tokenizers/yue2_tokenizer.py new file mode 100644 index 000000000..35c222821 --- /dev/null +++ b/vllm_omni/tokenizers/yue2_tokenizer.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Minimal YuE2 text tokenizer (frozen tiktoken BPE), vendored from upstream. + +Used by the offline driver to build token-id prompts; the audio codec tokens +never pass through it. +""" + +from __future__ import annotations + +import unicodedata +from pathlib import Path + +import pybase64 as base64 + + +class YuE2TextTokenizer: + def __init__(self, merge_file): + import tiktoken + + self.merge_file = Path(merge_file) + ranks = { + base64.b64decode(t): int(r) + for t, r in (line.split() for line in self.merge_file.read_bytes().splitlines() if line) + } + if len(ranks) != 151643: + raise ValueError("Expected checkpoint-native qwen.tiktoken (151643 ordinary tokens)") + specials = ["<|endoftext|>", "<|im_start|>", "<|im_end|>", "", "", "", "", ""] + specials += [f"" for i in range(200)] + specials[204:206] = ["", ""] + pattern = ( + r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}" + r"| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+" + ) + self._enc = tiktoken.Encoding( + "YuE2", + pat_str=pattern, + mergeable_ranks=ranks, + special_tokens={s: i + len(ranks) for i, s in enumerate(specials)}, + ) + + def encode(self, text): + return self._enc.encode_ordinary(unicodedata.normalize("NFC", text)) + + def decode(self, ids): + return self._enc.decode([int(i) for i in ids if 0 <= i < self._enc.n_vocab], errors="replace") + + +__all__ = ["YuE2TextTokenizer"] diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index c8a5c19f5..10c16b322 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -26,6 +26,7 @@ _CLASS_TO_MODULE: dict[str, str] = { "VoxtralTTSConfig": "vllm_omni.transformers_utils.configs.voxtral_tts", "CosyVoice3Config": "vllm_omni.transformers_utils.configs.cosyvoice3", "MiniMaxMusic3Config": "vllm_omni.transformers_utils.configs.minimax_music3", + "Yue2Config": "vllm_omni.transformers_utils.configs.yue2", "OmniVoiceConfig": "vllm_omni.transformers_utils.configs.omnivoice", "BailingMoeV2Config": "vllm_omni.transformers_utils.configs.ming_flash_omni", "BailingMM2Config": "vllm_omni.transformers_utils.configs.ming_flash_omni", @@ -56,6 +57,7 @@ __all__ = [ "VoxtralTTSConfig", "CosyVoice3Config", "MiniMaxMusic3Config", + "Yue2Config", "OmniVoiceConfig", "BailingMoeV2Config", "BailingMM2Config", @@ -96,3 +98,4 @@ from vllm_omni.transformers_utils.configs import omnivoice as _omnivoice # noqa from vllm_omni.transformers_utils.configs import sensenova_u1 as _sensenova_u1 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import voxcpm2 as _voxcpm2 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import voxtral_tts as _voxtral_tts # noqa: F401, E402 +from vllm_omni.transformers_utils.configs import yue2 as _yue2 # noqa: F401, E402 diff --git a/vllm_omni/transformers_utils/configs/yue2.py b/vllm_omni/transformers_utils/configs/yue2.py new file mode 100644 index 000000000..5063d5b15 --- /dev/null +++ b/vllm_omni/transformers_utils/configs/yue2.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""HuggingFace config for m-a-p/YuE2-3B. + +The checkpoint's own ``config.json`` already carries every field the AR +backbone needs (it is Qwen3-1.7B-shaped with q_norm/k_norm and an extended +184,704-entry vocabulary), plus the acoustic-side fields (``latent_dim``, +``max_latent_frames``, ``timestep_shift``). Registering this class lets vLLM +build a ``ModelConfig`` without ``trust_remote_code``: the package's own +``Yue2ForCausalLM`` replaces the repository's auto_map modeling code, and the +deploy YAML pins ``hf_overrides.architectures`` to it. +""" + +from __future__ import annotations + +from typing import Any + +from transformers import AutoConfig +from transformers.configuration_utils import PretrainedConfig + + +class Yue2Config(PretrainedConfig): + """Config for the YuE2-3B AR–NAR Mixture-of-Transformers checkpoint.""" + + model_type = "yue2" + + def __init__( + self, + hidden_size: int = 2048, + num_hidden_layers: int = 28, + num_attention_heads: int = 16, + num_key_value_heads: int = 8, + head_dim: int = 128, + intermediate_size: int = 6144, + vocab_size: int = 184704, + rms_norm_eps: float = 1e-6, + rope_theta: float = 1000000.0, + max_position_embeddings: int = 24576, + tie_word_embeddings: bool = False, + latent_dim: int = 64, + max_latent_frames: int = 24576, + timestep_shift: float = 1.0, + **kwargs: Any, + ) -> None: + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.intermediate_size = intermediate_size + self.vocab_size = vocab_size + self.rms_norm_eps = rms_norm_eps + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + self.latent_dim = latent_dim + self.max_latent_frames = max_latent_frames + self.timestep_shift = timestep_shift + # The checkpoint's config.json (like upstream's qwen_config derivative, + # fast.py) omits the standard Qwen3 fields the AR backbone expects; + # synthesize the same defaults upstream does so vLLM's Qwen3Model and + # MLP construct without AttributeError. + self.hidden_act = str(kwargs.pop("hidden_act", "silu")) + self.attention_bias = bool(kwargs.pop("attention_bias", False)) + self.attention_dropout = float(kwargs.pop("attention_dropout", 0.0)) + self.use_sliding_window = bool(kwargs.pop("use_sliding_window", False)) + self.sliding_window = kwargs.pop("sliding_window", None) + self.max_window_layers = int(kwargs.pop("max_window_layers", 0)) + + +AutoConfig.register("yue2", Yue2Config) + +__all__ = ["Yue2Config"] From 7b15d22f48a4af2ba83f9a209944d0496348fc91 Mon Sep 17 00:00:00 2001 From: Doug Smith Date: Thu, 1 Oct 2026 15:10:21 -0400 Subject: [PATCH 02/22] [CI/Build] Package runtime data explicitly (#7956) Signed-off-by: Doug Smith --- .github/workflows/build_wheel.yml | 29 ++++++++ pyproject.toml | 15 +++- tools/check_wheel_package_data.py | 118 ++++++++++++++++++++++++++++++ 3 files changed, 161 insertions(+), 1 deletion(-) create mode 100644 tools/check_wheel_package_data.py diff --git a/.github/workflows/build_wheel.yml b/.github/workflows/build_wheel.yml index f8087c5e9..bc970d7b7 100644 --- a/.github/workflows/build_wheel.yml +++ b/.github/workflows/build_wheel.yml @@ -56,6 +56,35 @@ jobs: run: | bash scripts/build_wheel.sh --python python + - name: Validate runtime package data from source archive + run: | + set -euo pipefail + archive_dir="$(mktemp -d)" + install_dir="$(mktemp -d)" + vcs_wheels=(dist/*.whl) + if [[ ${#vcs_wheels[@]} -ne 1 || ! -f "${vcs_wheels[0]}" ]]; then + echo "Expected exactly one VCS-built wheel" >&2 + exit 1 + fi + git archive HEAD | tar -x -C "${archive_dir}" + ( + cd "${archive_dir}" + VLLM_OMNI_TARGET_DEVICE=cuda \ + VLLM_OMNI_VERSION_OVERRIDE=0.0.0 \ + python -m build --wheel --outdir archive-dist + ) + wheels=("${archive_dir}"/archive-dist/*.whl) + if [[ ${#wheels[@]} -ne 1 || ! -f "${wheels[0]}" ]]; then + echo "Expected exactly one archive-built wheel" >&2 + exit 1 + fi + python -m pip install --no-deps --target "${install_dir}" "${wheels[0]}" + python tools/check_wheel_package_data.py \ + --source-root "${archive_dir}" \ + --wheel "${wheels[0]}" \ + --vcs-wheel "${vcs_wheels[0]}" \ + --installed-root "${install_dir}" + - name: Upload Wheel Artifact uses: actions/upload-artifact@v7 with: diff --git a/pyproject.toml b/pyproject.toml index 23148146e..c65255f99 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -161,12 +161,25 @@ vllm-omni = "vllm_omni.entrypoints.cli.main:main" vllm_omni_register_models = "vllm_omni.engine.arg_utils:register_omni_models_to_vllm" +[tool.setuptools] +include-package-data = false + [tool.setuptools.packages.find] where = ["."] include = ["vllm_omni*"] [tool.setuptools.package-data] -"vllm_omni" = ["_version.py", "py.typed", "transformers_utils/chat_templates/*.jinja"] +"vllm_omni" = [ + "_version.py", + "py.typed", + "deploy/*.md", + "deploy/*.yaml", + "diffusion/distributed/csrc/*.cu", + "diffusion/distributed/csrc/*.hpp", + "model_executor/models/covo_audio/speaker_prompt/*.npy", + "model_executor/models/step_audio2/assets/*.wav", + "transformers_utils/chat_templates/*.jinja", +] [tool.setuptools_scm] # no extra settings needed, presence enables setuptools-scm diff --git a/tools/check_wheel_package_data.py b/tools/check_wheel_package_data.py new file mode 100644 index 000000000..c4f8c1796 --- /dev/null +++ b/tools/check_wheel_package_data.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Verify runtime package data in a built vLLM-Omni wheel. + +The source tree may contain package data through setuptools-scm's Git file +finder even when a VCS-less source archive cannot rediscover it. This check +compares each runtime-owned source directory with both the wheel payload and, +optionally, an installed copy of that wheel. +""" + +import argparse +import sys +import zipfile +from pathlib import Path + +PACKAGE_NAME = "vllm_omni" +RUNTIME_DATA_DIRS = ( + Path("deploy"), + Path("diffusion/distributed/csrc"), + Path("model_executor/models/covo_audio/speaker_prompt"), + Path("model_executor/models/step_audio2/assets"), + Path("transformers_utils/chat_templates"), +) + + +def runtime_data_files(source_root: Path) -> list[Path]: + package_root = source_root / PACKAGE_NAME + files: list[Path] = [] + for relative_dir in RUNTIME_DATA_DIRS: + data_dir = package_root / relative_dir + if not data_dir.is_dir(): + raise FileNotFoundError(f"runtime data directory not found: {data_dir}") + files.extend( + path.relative_to(source_root) for path in data_dir.rglob("*") if path.is_file() and path.suffix != ".py" + ) + return sorted(files) + + +def missing_wheel_files(wheel: Path, expected: list[Path]) -> list[Path]: + with zipfile.ZipFile(wheel) as archive: + members = set(archive.namelist()) + return [path for path in expected if path.as_posix() not in members] + + +def wheel_package_files(wheel: Path) -> set[Path]: + package_prefix = f"{PACKAGE_NAME}/" + with zipfile.ZipFile(wheel) as archive: + return { + Path(member) + for member in archive.namelist() + if member.startswith(package_prefix) and not member.endswith("/") + } + + +def missing_installed_files(installed_root: Path, expected: list[Path]) -> list[Path]: + missing = [] + for path in expected: + installed_path = installed_root / path + if not installed_path.is_file(): + missing.append(path) + continue + installed_path.read_bytes() + return missing + + +def report_missing(label: str, missing: list[Path]) -> bool: + if not missing: + return False + print(f"{label} is missing {len(missing)} runtime package-data file(s):", file=sys.stderr) + for path in missing: + print(f" {path.as_posix()}", file=sys.stderr) + return True + + +def report_wheel_difference(vcs_wheel: Path, archive_wheel: Path) -> bool: + vcs_files = wheel_package_files(vcs_wheel) + archive_files = wheel_package_files(archive_wheel) + only_in_vcs = sorted(vcs_files - archive_files) + only_in_archive = sorted(archive_files - vcs_files) + if not only_in_vcs and not only_in_archive: + return False + + print("VCS-built and archive-built wheel package contents differ:", file=sys.stderr) + for label, paths in (("only in VCS-built wheel", only_in_vcs), ("only in archive-built wheel", only_in_archive)): + if paths: + print(f" {label} ({len(paths)}):", file=sys.stderr) + for path in paths: + print(f" {path.as_posix()}", file=sys.stderr) + return True + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source-root", type=Path, required=True) + parser.add_argument("--wheel", type=Path, required=True) + parser.add_argument("--vcs-wheel", type=Path) + parser.add_argument("--installed-root", type=Path) + args = parser.parse_args() + + expected = runtime_data_files(args.source_root.resolve()) + failed = report_missing("wheel", missing_wheel_files(args.wheel, expected)) + if args.vcs_wheel is not None: + failed |= report_wheel_difference(args.vcs_wheel.resolve(), args.wheel.resolve()) + if args.installed_root is not None: + failed |= report_missing( + "installed wheel", + missing_installed_files(args.installed_root.resolve(), expected), + ) + + if failed: + return 1 + print(f"validated {len(expected)} runtime package-data files in {args.wheel}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 2560f918592a8ee5d24b7d3fb2ac0bb7da35db1c Mon Sep 17 00:00:00 2001 From: Joshna-Medisetty Date: Thu, 1 Oct 2026 15:45:09 -0700 Subject: [PATCH 03/22] fix(ming): build ISTFT spectrogram in float32 (#7538) Signed-off-by: Joshna Medisetty Co-authored-by: Cursor --- vllm_omni/model_executor/models/common/ming/audio_dsp.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm_omni/model_executor/models/common/ming/audio_dsp.py b/vllm_omni/model_executor/models/common/ming/audio_dsp.py index d32415352..d94370bb8 100644 --- a/vllm_omni/model_executor/models/common/ming/audio_dsp.py +++ b/vllm_omni/model_executor/models/common/ming/audio_dsp.py @@ -219,7 +219,8 @@ class ISTFTHead(FourierHead): # phase = torch.atan2(y, x) # S = mag * torch.exp(phase * 1j) # better directly produce the complex value - S = mag * (x + 1j * y) + # float32 operands: bf16 promotes to bcomplex32, which torch.fft cannot handle + S = mag.float() * (x.float() + 1j * y.float()) audio, audio_buffer, window_buffer = self.istft( S, audio_buffer=audio_buffer, window_buffer=window_buffer, streaming=streaming, last_chunk=last_chunk ) From bbee488da3baf65817f4f8826b1f6e89cc28fda7 Mon Sep 17 00:00:00 2001 From: andyluo7 <43718156+andyluo7@users.noreply.github.com> Date: Thu, 1 Oct 2026 16:36:41 -0700 Subject: [PATCH 04/22] fix: repair ROCm test routing, CPU timeouts, and Breeze autotuning (#8342) Signed-off-by: andyluo7 --- .buildkite/amd/test-amd-merge.yml | 33 ++++++- .buildkite/amd/test-amd-ready.yml | 34 +++++++- .../models/auk/test_auk_vae_fastpath.py | 85 +++++++++++++++---- .../test_breeze_tts_2_first_code_sampler.py | 28 +++++- .../models/test_breeze_tts_2_graphs.py | 22 ++++- .../models/breeze_tts_2/depth_decoder.py | 10 ++- 6 files changed, 186 insertions(+), 26 deletions(-) diff --git a/.buildkite/amd/test-amd-merge.yml b/.buildkite/amd/test-amd-merge.yml index 1e3e852e9..e702e4d9c 100644 --- a/.buildkite/amd/test-amd-merge.yml +++ b/.buildkite/amd/test-amd-merge.yml @@ -41,7 +41,33 @@ steps: - export VLLM_ROCM_USE_AITER=0 # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM # that is not supported on CPU - - "pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB" + # Keep multi-GPU tests out of this one-GPU pool even when their module + # also carries the cpu marker for unit tests. + - "pytest -sv tests/diffusion -m 'core_model and cpu and not (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)' --ignore=tests/diffusion/cache/test_teacache_extractors.py --ignore=tests/diffusion/models/pi05/test_pi05_units.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB" + # This module takes about 22 minutes on MI300 and was consistently assigned to shard 4, + # pushing that shard to the 45-minute timeout. Keep its module-scoped fixture in one lane. + - label: "Simple · Pi0.5 CPU Test" + agent_pool: mi300_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + timeout_in_minutes: 45 + commands: + - export VLLM_ROCM_USE_AITER=0 + - "pytest -sv tests/diffusion/models/pi05/test_pi05_units.py -m 'core_model and cpu'" + + - label: "Diffusion · AuK VAE GPU Test" + agent_pool: mi300_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + timeout_in_minutes: 40 + commands: + - export VLLM_ROCM_USE_AITER=0 + - >- + timeout --signal=TERM --kill-after=2m 35m + pytest -sv tests/diffusion/models/auk/test_auk_vae_fastpath.py + -m 'core_model and rocm and cards_1' --run-level core_model - label: "Simple · Engine&Entrypoints Test" agent_pool: mi300_1 @@ -131,6 +157,11 @@ steps: grade: Blocking commands: - pytest -s -v tests/diffusion/distributed/test_tensor_parallel.py + - >- + timeout 20m pytest -s -v + tests/diffusion/attention/test_ulysses_uaa.py::test_ulysses_uaa_2d_mask_layout_matches_baseline + -m 'core_model and cards_2' + --run-level "core_model" - label: "Diffusion Model Test" agent_pool: mi300_1 diff --git a/.buildkite/amd/test-amd-ready.yml b/.buildkite/amd/test-amd-ready.yml index ae64e7c4b..3491d59c8 100644 --- a/.buildkite/amd/test-amd-ready.yml +++ b/.buildkite/amd/test-amd-ready.yml @@ -41,7 +41,33 @@ steps: - export VLLM_ROCM_USE_AITER=0 # ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM # that is not supported on CPU - - "pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB" + # Keep multi-GPU tests out of this one-GPU pool even when their module + # also carries the cpu marker for unit tests. + - "pytest -sv tests/diffusion -m 'core_model and cpu and not (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)' --ignore=tests/diffusion/cache/test_teacache_extractors.py --ignore=tests/diffusion/models/pi05/test_pi05_units.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB" + # This module takes about 22 minutes on MI300 and was consistently assigned to shard 4, + # pushing that shard to the 45-minute timeout. Keep its module-scoped fixture in one lane. + - label: "Simple · Pi0.5 CPU Test" + agent_pool: mi300_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + timeout_in_minutes: 45 + commands: + - export VLLM_ROCM_USE_AITER=0 + - "pytest -sv tests/diffusion/models/pi05/test_pi05_units.py -m 'core_model and cpu'" + + - label: "Diffusion · AuK VAE GPU Test" + agent_pool: mi300_1 + depends_on: amd-build + mirror_hardwares: [amdproduction] + grade: Blocking + timeout_in_minutes: 40 + commands: + - export VLLM_ROCM_USE_AITER=0 + - >- + timeout --signal=TERM --kill-after=2m 35m + pytest -sv tests/diffusion/models/auk/test_auk_vae_fastpath.py + -m 'core_model and rocm and cards_1' --run-level core_model - label: "Simple · Engine&Entrypoints Test" agent_pool: mi300_1 @@ -141,6 +167,11 @@ steps: - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model + - >- + timeout 20m pytest -s -v + tests/diffusion/attention/test_ulysses_uaa.py::test_ulysses_uaa_2d_mask_layout_matches_baseline + -m 'core_model and cards_2' + --run-level "core_model" - >- timeout 20m pytest -s -v tests/diffusion/models/ltx2/test_ltx2_transformer_ulysses.py @@ -165,6 +196,7 @@ steps: # registered multi-GPU marker from this one-GPU job. - >- timeout 15m pytest -sv tests/diffusion/models/ + --ignore=tests/diffusion/models/auk/test_auk_vae_fastpath.py -m 'core_model and cuda and not (cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)' --run-level "core_model" diff --git a/tests/diffusion/models/auk/test_auk_vae_fastpath.py b/tests/diffusion/models/auk/test_auk_vae_fastpath.py index b25381097..8a1b39aa1 100644 --- a/tests/diffusion/models/auk/test_auk_vae_fastpath.py +++ b/tests/diffusion/models/auk/test_auk_vae_fastpath.py @@ -11,10 +11,28 @@ eager off CUDA) and the tiled decode of long clips. import pytest import torch +from tests.helpers.mark import hardware_test from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE, LowPass, SnakeBeta, Upsample from vllm_omni.diffusion.models.auk.vae_cudagraph import AuKVAEDecodeGraph, plan_tiles +from vllm_omni.platforms import current_omni_platform -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +pytestmark = [pytest.mark.core_model] + + +@pytest.fixture(autouse=True) +def _bound_cpu_decode_threads(request): + if request.node.get_closest_marker("cpu") is None: + yield + return + # These narrow convolutions run in shared CI pods. A host-sized Torch + # pool spends more time coordinating workers than decoding small clips. + previous_threads = torch.get_num_threads() + torch.set_num_threads(1) + print(f"AuK CPU decode threads: {previous_threads} -> {torch.get_num_threads()}") + try: + yield + finally: + torch.set_num_threads(previous_threads) def _small_vae() -> AuKVAE: @@ -45,6 +63,23 @@ def _reference_snake(module: SnakeBeta, x: torch.Tensor) -> torch.Tensor: return x + (1.0 / (beta + 1e-9)) * torch.sin(x * alpha).pow(2) +def _assert_plain_graph_matches_eager(replay: torch.Tensor, eager: torch.Tensor) -> None: + assert replay.shape == eager.shape and replay.dtype == eager.dtype + assert torch.isfinite(replay).all() and torch.isfinite(eager).all() + error = (replay - eager).abs() + rms = error.square().mean().sqrt().item() + print(f"AuK plain graph: max_abs={error.max().item():.9g}, rms={rms:.9g}") + if current_omni_platform.is_rocm(): + # MI300 build 13040 measured max_abs <= 1.64e-7 and RMS <= 4.13e-8. + # Bound both peak and aggregate FP32 error; near-zero samples must + # not hide behind a relative tolerance. NVIDIA retains bit parity. + torch.testing.assert_close(replay, eager, atol=1e-6, rtol=0.0) + assert rms <= 1e-7, rms + else: + assert torch.equal(replay, eager) + + +@pytest.mark.cpu @torch.inference_mode() def test_window_padding_decodes_to_raw_zero() -> None: """Bucket padding is the normalized latent whose denormalized value is zero.""" @@ -57,6 +92,7 @@ def test_window_padding_decodes_to_raw_zero() -> None: torch.testing.assert_close(raw, torch.zeros_like(raw), atol=1e-6, rtol=0.0) +@pytest.mark.cpu @torch.inference_mode() def test_eager_snake_matches_the_reference_formula() -> None: module = SnakeBeta(6, alpha_logscale=True) @@ -69,6 +105,7 @@ def test_eager_snake_matches_the_reference_formula() -> None: assert torch.equal(module(x), _reference_snake(module, x)) +@pytest.mark.cpu @torch.inference_mode() def test_filter_cache_leaves_the_filters_and_output_unchanged() -> None: lowpass = LowPass(cutoff=0.25, half_width=0.3, stride=2, kernel_size=12, causal=True) @@ -88,6 +125,7 @@ def test_filter_cache_leaves_the_filters_and_output_unchanged() -> None: assert module._expanded.shape == (3, 1, 12) +@pytest.mark.cpu @torch.inference_mode() def test_decode_fast_paths_reproduce_the_plain_decode() -> None: vae = _small_vae() @@ -103,6 +141,7 @@ def test_decode_fast_paths_reproduce_the_plain_decode() -> None: assert all(module._cached for module in vae.modules() if isinstance(module, SnakeBeta)) +@pytest.mark.cpu @torch.inference_mode() def test_graph_wrapper_falls_back_to_eager_off_cuda(mocker) -> None: vae = _small_vae() @@ -119,6 +158,7 @@ def test_graph_wrapper_falls_back_to_eager_off_cuda(mocker) -> None: assert not wrapper._cache and not wrapper._compiled +@pytest.mark.cpu def test_compiled_bucket_is_the_smallest_captured_one_that_fits() -> None: wrapper = AuKVAEDecodeGraph(_small_vae(), compile_shapes=(16, 8, 8)) assert wrapper.compile_shapes == [8, 16] @@ -131,6 +171,7 @@ def test_compiled_bucket_is_the_smallest_captured_one_that_fits() -> None: assert wrapper.compiled_bucket(6) == 16 +@pytest.mark.cpu def test_plan_tiles_covers_the_clip_once_with_full_context() -> None: assert plan_tiles(40, 64, 10, 3) == [(0, 40, 0, 40)] windows = plan_tiles(1000, 512, 53, 9) @@ -149,10 +190,16 @@ def test_plan_tiles_covers_the_clip_once_with_full_context() -> None: # 512 and a 256 window instead of two 512s. assert plan_tiles(600, 512, 53, 9, sizes=(128, 256, 512)) == [(0, 512, 0, 503), (344, 256, 503, 600)] assert plan_tiles(1000, 512, 53, 9, sizes=(128, 256, 512))[-1] == (872, 128, 953, 1000) + narrow = plan_tiles(200, 64, 53, 9) + assert len(narrow) == 69 + assert [emit for _, _, emit, _ in narrow[:3]] == [0, 55, 57] + assert all(a[3] == b[2] for a, b in zip(narrow, narrow[1:])) + assert narrow[-1][3] == 200 with pytest.raises(ValueError, match="must exceed"): plan_tiles(100, 60, 53, 9) +@pytest.mark.cpu @torch.inference_mode() def test_decode_context_bounds_the_measured_receptive_field() -> None: vae = _small_vae() @@ -170,25 +217,26 @@ def test_decode_context_bounds_the_measured_receptive_field() -> None: assert differing and all(index < left or end - start - index <= right for index in differing) +@pytest.mark.cpu @torch.inference_mode() def test_tiled_decode_matches_the_whole_decode() -> None: vae = _small_vae() # Off the accelerator every window decodes eagerly, so this isolates the stitching. - wrapper = AuKVAEDecodeGraph(vae, compile_shapes=(64,), tile_frames=64) - assert wrapper.tile_frames == 64 and wrapper.context_frames == (53, 9) - for frames in (64, 65, 200, 331): + wrapper = AuKVAEDecodeGraph(vae, compile_shapes=(128,), tile_frames=128) + assert wrapper.tile_frames == 128 and wrapper.context_frames == (53, 9) + for frames in (128, 129, 200, 331): latents = torch.randn(1, frames, vae.latent_dim) whole = vae.decode(latents) tiled = wrapper(latents) assert tiled.shape == whole.shape torch.testing.assert_close(tiled, whole, atol=1e-5, rtol=0.0) - assert wrapper.last_mode == ("eager" if frames == 64 else "tiled") + assert wrapper.last_mode == ("eager" if frames == 128 else "tiled") # The streaming form yields the same audio in order. latents = torch.randn(1, 200, vae.latent_dim) pieces = [(start, chunk.clone()) for start, chunk in wrapper.decode_tiles(latents)] - # With a 64-frame tile and a 62-frame halo each interior tile adds two frames. - assert [start for start, _ in pieces] == [emit for _, _, emit, _ in plan_tiles(200, 64, 53, 9)] - assert [start for start, _ in pieces][:3] == [0, 55, 57] + # A 128-frame tile avoids repeatedly decoding the same 62-frame halo. + assert [start for start, _ in pieces] == [emit for _, _, emit, _ in plan_tiles(200, 128, 53, 9)] + assert [start for start, _ in pieces][:3] == [0, 119, 185] torch.testing.assert_close( torch.cat([chunk for _, chunk in pieces], dim=1), vae.decode(latents), atol=1e-5, rtol=0.0 ) @@ -198,20 +246,24 @@ def test_tiled_decode_matches_the_whole_decode() -> None: AuKVAEDecodeGraph(vae, tile_frames=60) -@pytest.mark.cuda +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA") @torch.inference_mode() def test_graph_replay_matches_eager_per_length() -> None: vae = _small_vae().to("cuda") wrapper = AuKVAEDecodeGraph(vae, max_graphs=2) pools = [] + saved_outputs: list[tuple[torch.Tensor, torch.Tensor]] = [] for frames in (6, 9, 6, 12, 6): latents = torch.randn(1, frames, vae.latent_dim, device="cuda") eager = vae.decode(latents) replay = wrapper(latents) - # Exact-length graphs replay the very same kernels: bit-identical, - # including after the generation was retired. - assert torch.equal(replay, eager), frames + _assert_plain_graph_matches_eager(replay, eager) + # Returned waveforms own their storage and survive later replays + # and retirement of the generation that produced them. + for previous, snapshot in saved_outputs: + assert torch.equal(previous, snapshot) + saved_outputs.append((replay, replay.clone())) pools.append(wrapper._plain_pool) # The third distinct length found the cache full, so the whole generation # (6 and 9) was retired together with its pool rather than one graph at a @@ -220,7 +272,7 @@ def test_graph_replay_matches_eager_per_length() -> None: assert pools[0] is pools[2] and pools[3] is pools[4] and pools[2] is not pools[3] -@pytest.mark.cuda +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA") @torch.inference_mode() def test_bucketed_graph_only_disturbs_the_tail() -> None: @@ -240,7 +292,7 @@ def test_bucketed_graph_only_disturbs_the_tail() -> None: torch.testing.assert_close(replay, eager, atol=0.1, rtol=0.0) -@pytest.mark.cuda +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) @pytest.mark.skipif(not torch.cuda.is_available(), reason="torch.compile + CUDA graph capture requires CUDA") @torch.inference_mode() def test_compiled_buckets_replay_within_fusion_tolerance_and_leave_longer_clips_to_plain_graphs() -> None: @@ -266,11 +318,12 @@ def test_compiled_buckets_replay_within_fusion_tolerance_and_leave_longer_clips_ # a longer clip gets its own plain graph. assert wrapper.tile_frames == 0 long = torch.randn(1, 12, vae.latent_dim, device="cuda") - assert torch.equal(wrapper(long), vae.decode(long)) + replay, eager = wrapper(long), vae.decode(long) + _assert_plain_graph_matches_eager(replay, eager) assert wrapper.last_mode == "graph" and list(wrapper._cache) == [12] -@pytest.mark.cuda +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) @pytest.mark.skipif(not torch.cuda.is_available(), reason="torch.compile + CUDA graph capture requires CUDA") @torch.inference_mode() def test_tiles_replay_the_compiled_bucket_for_long_clips() -> None: diff --git a/tests/model_executor/models/test_breeze_tts_2_first_code_sampler.py b/tests/model_executor/models/test_breeze_tts_2_first_code_sampler.py index 660737765..9d9047df1 100644 --- a/tests/model_executor/models/test_breeze_tts_2_first_code_sampler.py +++ b/tests/model_executor/models/test_breeze_tts_2_first_code_sampler.py @@ -4,6 +4,7 @@ import pytest import torch +from vllm.platforms import current_platform from tests.helpers.mark import hardware_test from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import sample_logits @@ -16,6 +17,10 @@ VOCAB = 2052 EOS = 2051 +def _captured_generator_state_supported() -> bool: + return current_platform.is_cuda() and hasattr(torch.cuda.CUDAGraph, "register_generator_state") + + def _logits(count: int, *, device: str, seed: int = 100, vocab: int = VOCAB) -> list[torch.Tensor]: source = torch.Generator(device="cpu").manual_seed(seed) values = torch.randn((count, vocab), generator=source, dtype=torch.float32) * 3 @@ -164,6 +169,14 @@ def test_cuda_replay_preserves_tokens_rng_and_owned_outputs(monkeypatch) -> None right[codebook].exponential_(generator=live_expected[index]) torch.testing.assert_close(left, right, atol=0, rtol=0) _assert_states(actual_generators, expected_generators) + if not _captured_generator_state_supported(): + # ROCm deliberately uses the eager sampler because CUDA graph generator-state registration is an + # NVIDIA-only path. The loop above still verifies GPU tokens and every request-owned RNG stream. + assert not sampler._graphs + for value, expected in retained: + torch.testing.assert_close(value, expected, atol=0, rtol=0) + torch.testing.assert_close(torch.cuda.get_rng_state(), global_state, atol=0, rtol=0) + return # B1/3/16/17/31/32 must use exactly buckets 1/4/16/32. assert len(sampler._graphs) == 4 rows = _logits(3, device="cuda", seed=2026) @@ -189,6 +202,7 @@ def test_cuda_replay_preserves_tokens_rng_and_owned_outputs(monkeypatch) -> None @torch.inference_mode() def test_cuda_parameters_vocab_and_missing_graph_api(monkeypatch) -> None: sampler = BreezeFirstCodeSampler() + captures_generator_state = _captured_generator_state_supported() actual_generators = _generators([42, 43, 44], "cuda") expected_generators = _generators([42, 43, 44], "cuda") cases = ((0.9, 50, 1.0), (0.7, 100, 0.8), (1.2, 0, 0.95), (1.0, 1, 1.0), (1.0, 0, 0.5)) @@ -201,12 +215,12 @@ def test_cuda_parameters_vocab_and_missing_graph_api(monkeypatch) -> None: expected = _reference(rows, parameters, expected_generators) _assert_outputs(sampler.sample(rows, parameters, actual_generators), expected, rows) _assert_states(actual_generators, expected_generators) - assert len(sampler._graphs) == len(cases) + assert len(sampler._graphs) == (len(cases) if captures_generator_state else 0) rows = _logits(3, device="cuda", vocab=17) expected = _reference(rows, PARAMETERS, expected_generators) _assert_outputs(sampler.sample(rows, PARAMETERS, actual_generators), expected, rows) _assert_states(actual_generators, expected_generators) - assert len(sampler._graphs) == len(cases) + 1 + assert len(sampler._graphs) == (len(cases) + 1 if captures_generator_state else 0) class UnavailableGraph: def __init__(self, *_args, **_kwargs): @@ -243,7 +257,7 @@ def test_cuda_shared_generator_uses_sequential_eager_fallback(monkeypatch) -> No # The duplicate-generator guard must precede lookup of an existing graph. sampler.sample(rows, PARAMETERS, _generators([100, 101, 102], "cuda")) cached = dict(sampler._graphs) - assert len(cached) == 1 + assert len(cached) == (1 if _captured_generator_state_supported() else 0) actual = torch.Generator(device="cuda").manual_seed(2**64 - 23) reference = torch.Generator(device="cuda").manual_seed(2**64 - 23) multinomial = torch.multinomial @@ -270,6 +284,7 @@ def test_cuda_shared_generator_uses_sequential_eager_fallback(monkeypatch) -> No @torch.inference_mode() def test_cuda_cache_limit_falls_back_without_evicting_warmed_graphs(monkeypatch) -> None: sampler = BreezeFirstCodeSampler() + captures_generator_state = _captured_generator_state_supported() rows = _logits(1, device="cuda") actual_generators = _generators([42], "cuda") expected_generators = _generators([42], "cuda") @@ -279,7 +294,7 @@ def test_cuda_cache_limit_falls_back_without_evicting_warmed_graphs(monkeypatch) _assert_outputs(sampler.sample(rows, parameters, actual_generators), expected, rows) _assert_states(actual_generators, expected_generators) cached = dict(sampler._graphs) - assert len(cached) == 32 + assert len(cached) == (32 if captures_generator_state else 0) parameters = (0.9, 33, 1.0) expected = _reference(rows, parameters, expected_generators) multinomial = torch.multinomial @@ -298,6 +313,11 @@ def test_cuda_cache_limit_falls_back_without_evicting_warmed_graphs(monkeypatch) assert all(sampler._graphs[key] is value for key, value in cached.items()) expected = _reference(rows, (0.9, 1, 1.0), expected_generators) + if not captures_generator_state: + _assert_outputs(sampler.sample(rows, (0.9, 1, 1.0), actual_generators), expected, rows) + _assert_states(actual_generators, expected_generators) + return + def unexpected_multinomial(*_args, **_kwargs): raise AssertionError("Cache exhaustion must retain existing captured entries") diff --git a/tests/model_executor/models/test_breeze_tts_2_graphs.py b/tests/model_executor/models/test_breeze_tts_2_graphs.py index 72c2cfeea..b60c8bf46 100644 --- a/tests/model_executor/models/test_breeze_tts_2_graphs.py +++ b/tests/model_executor/models/test_breeze_tts_2_graphs.py @@ -9,7 +9,11 @@ from transformers.models.t5gemma2.modeling_t5gemma2 import T5Gemma2TextEncoder from vllm.platforms import current_platform from tests.helpers.mark import hardware_test -from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import BreezeDepthDecoder, sample_logits +from vllm_omni.model_executor.models.breeze_tts_2.depth_decoder import ( + BreezeDepthDecoder, + _depth_layer_compile_options, + sample_logits, +) from vllm_omni.model_executor.models.breeze_tts_2.modeling_breeze import BreezeForConditionalGeneration from vllm_omni.model_executor.models.breeze_tts_2.text_encoder_graph import ( BreezeTextEncoderCompiled, @@ -19,6 +23,17 @@ from vllm_omni.model_executor.models.breeze_tts_2.text_encoder_graph import ( pytestmark = [pytest.mark.core_model] +@pytest.mark.cpu +@pytest.mark.parametrize(("is_rocm", "expected"), [(True, False), (False, True)]) +def test_depth_layer_compile_options(monkeypatch, is_rocm: bool, expected: bool) -> None: + monkeypatch.setattr(current_platform, "is_rocm", lambda: is_rocm) + + assert _depth_layer_compile_options() == { + "epilogue_fusion": False, + "max_autotune": expected, + } + + @pytest.fixture def full_precision_matmul(): # Different padded shapes can select different TF32 kernels. Compare @@ -241,6 +256,7 @@ def test_depth_graph_buckets_bound_shapes_and_preserve_requests(full_precision_m seeds[0], seeds[-1] = 2**63 + 17, 2**64 - 1 generators = [torch.Generator(device="cuda").manual_seed(seed) for seed in seeds] reference_generators = [torch.Generator(device="cuda").manual_seed(seed) for seed in seeds] + captures_generator_state = current_platform.is_cuda() and hasattr(torch.cuda.CUDAGraph, "register_generator_state") expected_greedy = torch.cat( [ @@ -314,7 +330,7 @@ def test_depth_graph_buckets_bound_shapes_and_preserve_requests(full_precision_m branches = 2 if guidance_scale != 1 else 1 graph_key = (request_bucket * branches, branches, temperature == 0) with monkeypatch.context() as rng_guard: - if temperature > 0 and graph_key in depth._graphs: + if temperature > 0 and graph_key in depth._graphs and captures_generator_state: # A warmed replay must generate noise within the CUDA # graph, without dispatching new RNG kernels from Python. rng_guard.setattr(torch.Tensor, "exponential_", reject_eager_noise) @@ -349,7 +365,7 @@ def test_depth_graph_buckets_bound_shapes_and_preserve_requests(full_precision_m expected_keys.update((batch, 2, greedy) for batch in (2, 8, 16, 32) for greedy in (False, True)) assert set(depth._graphs) == expected_keys for (batch_bucket, branches, greedy), entry in depth._graphs.items(): - if greedy: + if greedy or not captures_generator_state: assert entry.noise_generators is None else: assert entry.noise_generators is not None diff --git a/vllm_omni/model_executor/models/breeze_tts_2/depth_decoder.py b/vllm_omni/model_executor/models/breeze_tts_2/depth_decoder.py index 70b7262f1..0adc38d9b 100644 --- a/vllm_omni/model_executor/models/breeze_tts_2/depth_decoder.py +++ b/vllm_omni/model_executor/models/breeze_tts_2/depth_decoder.py @@ -16,6 +16,14 @@ from vllm.platforms import current_platform from vllm_omni.platforms import current_omni_platform +def _depth_layer_compile_options() -> dict[str, bool]: + """Keep ROCm away from the unstable Triton GEMM autotune path.""" + return { + "epilogue_fusion": False, + "max_autotune": not current_platform.is_rocm(), + } + + def sample_logits( logits: torch.Tensor, temperature: float, @@ -283,7 +291,7 @@ class BreezeDepthDecoder(nn.Module): BreezeDepthLayer.forward, fullgraph=True, dynamic=True, - options={"epilogue_fusion": False, "max_autotune": True}, + options=_depth_layer_compile_options(), ) self._compiled_sampler = torch.compile(sample_graph_logits, fullgraph=True, dynamic=False) entry = BreezeDepthGraph( From 68003cf6adff547791505e6d5e26f161cef4a78a Mon Sep 17 00:00:00 2001 From: wtz2333 <2955110911@qq.com> Date: Fri, 2 Oct 2026 11:57:50 +0800 Subject: [PATCH 05/22] [Bugfix] Fix washed-out LingBot Video outputs (#7037) Signed-off-by: wtz2333 <2955110911@qq.com> Co-authored-by: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> --- .../test_pipeline_lingbot_video.py | 43 ++++++++++++++++++- .../lingbot_video/pipeline_lingbot_video.py | 8 ++-- 2 files changed, 45 insertions(+), 6 deletions(-) diff --git a/tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py b/tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py index f3e0b1ce3..45218db5a 100644 --- a/tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py +++ b/tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py @@ -148,7 +148,8 @@ def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested(): assert array.shape == (2, 4, 4, 3) video = post({"video": frames}, SimpleNamespace(output_type="pt")) assert set(video) == {"video"} - assert video["video"] is frames + assert isinstance(video["video"], np.ndarray) + np.testing.assert_array_equal(video["video"], frames.numpy()) image = post( {"image": torch.ones(4, 4, 3)}, SimpleNamespace(output_type="pt"), @@ -162,6 +163,40 @@ def test_postprocess_keeps_torch_by_default_and_converts_np_when_requested(): assert latent_output["image"] is latent +@pytest.mark.parametrize("output_type", [None, "pt", "np"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_postprocess_preserves_video_range_for_serving(output_type: str | None, dtype: torch.dtype) -> None: + from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func + from vllm_omni.entrypoints.openai.video_api_utils import _coerce_video_to_frames + + # Include black and midtones: all-white frames cannot expose double normalization. + frames = torch.tensor([0.0, 0.25, 0.5, 0.75, 1.0], dtype=dtype) + frames = frames.reshape(5, 1, 1, 1).expand(5, 2, 4, 3).clone() + expected = frames.float().numpy().copy() + sampling = None if output_type is None else OmniDiffusionSamplingParams(output_type=output_type) + post = get_lingbot_video_post_process_func(None) + + processed = post({"video": frames}, sampling) + encoded_frames = np.stack(_coerce_video_to_frames(processed["video"])) + + assert set(processed) == {"video"} + assert encoded_frames.dtype == np.float32 + np.testing.assert_array_equal(encoded_frames, expected) + np.testing.assert_array_equal(frames.float().numpy(), expected) + + +@pytest.mark.parametrize("output_key", ["image", "video"]) +def test_postprocess_preserves_latent_payload(output_key: str) -> None: + from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func + + latents = torch.tensor([-2.0, 0.0, 2.0]).reshape(1, 3, 1, 1, 1) + sampling = OmniDiffusionSamplingParams(output_type="latent") + processed = get_lingbot_video_post_process_func(None)({output_key: latents}, sampling) + + assert set(processed) == {output_key} + assert processed[output_key] is latents + + def test_postprocess_preserves_image_contract_for_serving(): from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.models.lingbot_video import get_lingbot_video_post_process_func @@ -181,7 +216,9 @@ def test_postprocess_preserves_image_contract_for_serving(): sampling_params=sampling, request_id="lingbot-image-contract-test", ) - raw_output = {"image": torch.ones(8, 12, 3)} + image = torch.zeros(8, 12, 3) + image[4:] = 1.0 + raw_output = {"image": image} processed = get_lingbot_video_post_process_func(SimpleNamespace())(raw_output, sampling) normalized = normalize_diffusion_postprocess_output(processed) @@ -199,6 +236,8 @@ def test_postprocess_preserves_image_contract_for_serving(): endpoint_images = _extract_images_from_result(result) assert len(endpoint_images) == 1 assert endpoint_images[0].size == (12, 8) + assert endpoint_images[0].getpixel((0, 0)) == (0, 0, 0) + assert endpoint_images[0].getpixel((0, 7)) == (255, 255, 255) def test_forward_resolves_t2v_sampling_and_flow_shift_alias(): diff --git a/vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py b/vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py index 51687a241..50fb33a33 100644 --- a/vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py +++ b/vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py @@ -288,13 +288,13 @@ def get_lingbot_video_post_process_func(od_config: OmniDiffusionConfig): output_key = output_keys[0] frames = frames[output_key] output_type = getattr(sampling_params, "output_type", None) or "pt" - # The image serving path currently accepts PIL images or NumPy arrays, - # while LingBot decodes T2I outputs to tensors. Keep that compatibility - # conversion without discarding the model's image/video payload key. + # Decoded LingBot pixels are already in [0, 1]. The video API treats + # floating tensors as [-1, 1], so use its NumPy path to avoid normalizing + # twice. Images also need NumPy for serving; keep latent tensors intact. if ( isinstance(frames, torch.Tensor) and output_type != "latent" - and (output_type == "np" or output_key == "image") + and (output_type == "np" or output_key in {"image", "video"}) ): frames = frames.float().cpu().numpy() return {output_key: frames} if output_key is not None else frames From d93e1a8e09d2a0d386c4a58424470d685ee7743c Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Thu, 1 Oct 2026 22:22:54 -0700 Subject: [PATCH 06/22] [Model] Optimize Qwen3-Omni MRv2 handoff, code prediction and audio decoding (#8223) Signed-off-by: Sy03 <1370724210@qq.com> --- .../common/ci_source_file_dependencies.yml | 8 + docs/design/index.md | 1 + docs/design/qwen3_omni_mrv2_performance.md | 22 + tests/config/test_environment_variables.py | 2 +- .../models/common/test_fused_sampling.py | 90 ++++ .../test_short_kv_attention.py | 2 +- .../qwen3_omni/test_first_frame_cuda.py | 32 ++ .../test_qwen3_omni_code2wav_cudnn.py | 87 ++++ .../test_qwen3_omni_code2wav_snake.py | 88 ++++ .../qwen3_omni/test_qwen3_omni_first_audio.py | 208 +++++++++ .../models/qwen3_omni/test_qwen3_omni_mrv2.py | 116 +++++ .../qwen3_omni/test_qwen3_omni_pp_capture.py | 1 + .../qwen3_tts/test_code_predictor_dtype.py | 21 + .../models/test_talker_first_audio.py | 70 +++ .../test_qwen3_omni_chunk_ramp.py | 116 +++++ .../test_qwen3_omni_first_token_handoff.py | 146 ++++++ .../test_qwen3_omni_streaming_helpers.py | 77 ++++ tests/worker_v2/test_omni_model_state.py | 6 + .../config/environment_variable_inventory.py | 1 + vllm_omni/deploy/qwen3_omni_moe_mrv2.yaml | 50 +++ .../deploy/qwen3_omni_moe_mrv2_h200.yaml | 27 ++ .../models/common/fused_sampling.py | 97 ++++ .../models/common/qwen3_code_predictor.py | 251 ++++++++++- .../short_kv_attention.py | 0 .../models/common/talker_first_audio.py | 29 ++ .../models/qwen3_omni/first_frame_decoder.py | 99 +++++ .../models/qwen3_omni/pipeline.py | 5 + .../models/qwen3_omni/qwen3_omni.py | 418 +++++++++++++++--- .../models/qwen3_omni/qwen3_omni_code2wav.py | 218 ++++++++- .../models/qwen3_tts/cached_code_predictor.py | 3 +- .../models/qwen3_tts/first_audio.py | 18 +- .../qwen3_tts_code_predictor_vllm.py | 3 + .../chunk_size_utils.py | 17 +- .../stage_input_processors/qwen3_omni.py | 179 ++++++-- .../model_states/omni_model_state.py | 16 +- 35 files changed, 2363 insertions(+), 161 deletions(-) create mode 100644 docs/design/qwen3_omni_mrv2_performance.md create mode 100644 tests/model_executor/models/common/test_fused_sampling.py rename tests/model_executor/models/{qwen3_tts => common}/test_short_kv_attention.py (95%) create mode 100644 tests/model_executor/models/qwen3_omni/test_first_frame_cuda.py create mode 100644 tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_cudnn.py create mode 100644 tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_snake.py create mode 100644 tests/model_executor/models/qwen3_omni/test_qwen3_omni_first_audio.py create mode 100644 tests/model_executor/models/qwen3_omni/test_qwen3_omni_mrv2.py create mode 100644 tests/model_executor/models/test_talker_first_audio.py create mode 100644 tests/model_executor/stage_input_processors/test_qwen3_omni_chunk_ramp.py create mode 100644 tests/model_executor/stage_input_processors/test_qwen3_omni_first_token_handoff.py create mode 100644 vllm_omni/deploy/qwen3_omni_moe_mrv2.yaml create mode 100644 vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml create mode 100644 vllm_omni/model_executor/models/common/fused_sampling.py rename vllm_omni/model_executor/models/{qwen3_tts => common}/short_kv_attention.py (100%) create mode 100644 vllm_omni/model_executor/models/common/talker_first_audio.py create mode 100644 vllm_omni/model_executor/models/qwen3_omni/first_frame_decoder.py diff --git a/.buildkite/common/ci_source_file_dependencies.yml b/.buildkite/common/ci_source_file_dependencies.yml index 9594e62e5..b34e9d271 100644 --- a/.buildkite/common/ci_source_file_dependencies.yml +++ b/.buildkite/common/ci_source_file_dependencies.yml @@ -16,6 +16,9 @@ x-qwen3-tts: &qwen3_tts - vllm_omni/worker/gpu_model_runner.py - vllm_omni/model_executor/models/qwen3_tts/ - vllm_omni/model_executor/models/common/qwen3_code_predictor.py + - vllm_omni/model_executor/models/common/fused_sampling.py + - vllm_omni/model_executor/models/common/short_kv_attention.py + - vllm_omni/model_executor/models/common/talker_first_audio.py - vllm_omni/model_executor/models/common/snake_activation.py - vllm_omni/model_executor/models/common/alias_free_activation.py - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py @@ -85,11 +88,16 @@ x-qwen3-omni: &qwen3_omni - vllm_omni/entrypoints/openai/serving_chat.py - vllm_omni/model_executor/models/qwen3_omni/ - vllm_omni/model_executor/models/common/qwen3_code_predictor.py + - vllm_omni/model_executor/models/common/fused_sampling.py + - vllm_omni/model_executor/models/common/short_kv_attention.py + - vllm_omni/model_executor/models/common/talker_first_audio.py - vllm_omni/model_executor/models/common/snake_activation.py - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py - vllm_omni/deploy/qwen3_omni_moe.yaml + - vllm_omni/deploy/qwen3_omni_moe_mrv2.yaml + - vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml - vllm_omni/deploy/qwen3_omni_moe_thinking.yaml x-wan22: &wan22 diff --git a/docs/design/index.md b/docs/design/index.md index 748e1d9ca..af8dbf438 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -116,4 +116,5 @@ The pre-#5137 pages are preserved in the [legacy module archive](module/archive/README.md) for historical reference and are not active design contracts. +- [Qwen3-Omni MRv2 performance](qwen3_omni_mrv2_performance.md) - [MiniCPM-o 4.5 turn-mode MRv2 performance](minicpm_o45_mrv2_performance.md) diff --git a/docs/design/qwen3_omni_mrv2_performance.md b/docs/design/qwen3_omni_mrv2_performance.md new file mode 100644 index 000000000..c81b742dd --- /dev/null +++ b/docs/design/qwen3_omni_mrv2_performance.md @@ -0,0 +1,22 @@ +# Qwen3-Omni MRv2 performance + +This model adaptation builds on #8184's existing stage runtime and output +contracts. It adds Talker CPU metadata parsing, sampled-token embedding +handoff, an optional first-frame decoder, length-grouped Code2Wav graph +buckets, and opt-in incremental short-KV code prediction. The H200 profile +also enables conservative fused sampling and bounded cuDNN algorithm search. +The shared predictor keeps the upstream re-prefill helper and its default +sampling behavior; short KV is explicitly selected by the deployment profile. + +Use `vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml` for the experimental H200 +profile. Preserve a control with `code_predictor_fused_sampling=false` and +`codec_cudnn_benchmark=false` when C64 P99 is the primary requirement. +cuDNN search restores process settings on exit and has startup cost. Kernel +arithmetic is not bitwise-equivalent; the first-frame decoder uses additional +model memory. Full-duplex performance is not established by turn-mode tests. + +First-frame delivery requires `talker_first_audio: true` in connector extra, +a one-frame initial codec chunk, CUDA, a local single-rank Talker and prefix caching disabled. +Other deployments retain the regular Code2Wav path. +`codec_fused_snake` explicitly opts CUDA profiles into fused decoder blocks; +the base V1 profile retains its original blocks and indexed text cache. diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index 6f004b6b6..e1f6aa651 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -169,7 +169,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29, EnvironmentVariableCategory.MODEL_SPECIFIC: 75, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21, - EnvironmentVariableCategory.INTERNAL: 6, + EnvironmentVariableCategory.INTERNAL: 7, } disposition_counts = Counter( diff --git a/tests/model_executor/models/common/test_fused_sampling.py b/tests/model_executor/models/common/test_fused_sampling.py new file mode 100644 index 000000000..45ac3f81d --- /dev/null +++ b/tests/model_executor/models/common/test_fused_sampling.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cuda, + pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"), +] + + +def _reference(logits, uniforms, top_k, top_p): + if top_k > 0: + threshold = logits.topk(top_k, dim=-1).values[:, -1:] + logits = logits.masked_fill(logits < threshold, -torch.inf) + if top_p < 1: + values, indices = logits.sort(dim=-1, descending=True) + probs = values.softmax(dim=-1, dtype=torch.float32) + values[(probs.cumsum(-1) - probs) >= top_p] = -torch.inf + logits = values.scatter(1, indices, values) + return (logits.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True) + + +@pytest.mark.parametrize("batch,vocab", [(1, 128), (7, 2048), (64, 2048)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("top_k,top_p", [(-1, 0.8), (0, 1.0), (50, 0.8), (50, 0.01)]) +def test_fixed_uniform_sampling_and_graph(batch, vocab, dtype, top_k, top_p): + from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel + + torch.manual_seed(173) + logits = torch.randn(batch, vocab, device="cuda", dtype=dtype) + # The MTP caller supplies a noncontiguous row slice of [B, 15, vocab]. + uniforms = torch.rand(batch, 15, vocab, device="cuda")[:, 3].clamp_min_(1e-7) + expected = _reference(logits, uniforms, top_k, top_p) + actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p) + logits.add_(torch.randn_like(logits) * 0.1) + uniforms.uniform_(1e-7, 1 - 1e-7) + graph.replay() + torch.testing.assert_close(captured, _reference(logits, uniforms, top_k, top_p), rtol=0, atol=0) + + +@pytest.mark.parametrize("top_k,top_p", [(3, 0.5), (3, 1.0), (0, 0.8)]) +def test_threshold_ties_and_degenerate_rows(top_k, top_p): + from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel + + logits = torch.tensor( + [[2, 2, 2, 2, 2, 1, 0, -1], [1] * 8, [-torch.inf] * 8, [torch.nan, 1, 2, 3, 4, 5, 6, 7]], + device="cuda", + dtype=torch.float32, + ) + uniforms = torch.tensor([[0.4, 0.7, 0.8, 0.9, 0.2, 0.3, 0.6, 0.5]], device="cuda").expand(4, -1) + actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p) + torch.testing.assert_close(actual, _reference(logits, uniforms, top_k, top_p), rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_native_nucleus_boundaries(dtype): + from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel + + torch.manual_seed(1923) + logits = torch.randn(1, 2048, device="cuda", dtype=dtype) + values, indices = logits.sort(descending=True) + probabilities = values.softmax(-1, dtype=torch.float32) + before = probabilities.cumsum(-1) - probabilities + for rank in [1, 19, 49, 1000, 2000]: + boundary = before[0, rank] + uniforms = torch.full_like(logits, 0.001, dtype=torch.float32) + uniforms[0, indices[0, rank]] = 1 - 1e-7 + for direction in [-torch.inf, torch.inf]: + top_p = torch.nextafter(boundary, boundary.new_tensor(direction)).item() + actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=0, top_p=top_p) + torch.testing.assert_close(actual, _reference(logits, uniforms, 0, top_p), rtol=0, atol=0) + + +@pytest.mark.parametrize("top_p", [0.8, 1.0]) +def test_strided_logits_and_non_power_of_two_vocabulary(top_p): + from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel + + for seed in range(16): + torch.manual_seed(seed) + logits = torch.randn(128, 2050, device="cuda")[:, ::2] + uniforms = torch.rand(128, 3, 1025, device="cuda")[:, 1].clamp_min_(1e-7) + actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=-1, top_p=top_p) + torch.testing.assert_close(actual, _reference(logits, uniforms, -1, top_p), rtol=0, atol=0) diff --git a/tests/model_executor/models/qwen3_tts/test_short_kv_attention.py b/tests/model_executor/models/common/test_short_kv_attention.py similarity index 95% rename from tests/model_executor/models/qwen3_tts/test_short_kv_attention.py rename to tests/model_executor/models/common/test_short_kv_attention.py index 15dcecb14..000872617 100644 --- a/tests/model_executor/models/qwen3_tts/test_short_kv_attention.py +++ b/tests/model_executor/models/common/test_short_kv_attention.py @@ -16,7 +16,7 @@ pytestmark = [ @pytest.mark.parametrize("heads,kv_heads,dim", [(16, 8, 128), (8, 8, 64)]) @pytest.mark.parametrize("positions", [[0, 1], [2], [8], [15]]) def test_short_attention_causal_gqa_and_graph(batch, heads, kv_heads, dim, positions): - from vllm_omni.model_executor.models.qwen3_tts.short_kv_attention import short_kv_attention + from vllm_omni.model_executor.models.common.short_kv_attention import short_kv_attention with torch.inference_mode(): torch.manual_seed(42) diff --git a/tests/model_executor/models/qwen3_omni/test_first_frame_cuda.py b/tests/model_executor/models/qwen3_omni/test_first_frame_cuda.py new file mode 100644 index 000000000..0d75d68ee --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/test_first_frame_cuda.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import Qwen3OmniFirstFrameDecoder + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("rows", [16, 19]) +@torch.inference_mode() +def test_first_frame_graph_outputs_own_storage(dtype, rows): + class Decoder(torch.nn.Linear): + config = SimpleNamespace(num_quantizers=1) + + def forward(self, codes): + return codes.to(dtype) * self.weight + + model = Decoder(1, 1, bias=False, device="cuda", dtype=dtype) + model.weight.fill_(1) + decoder = Qwen3OmniFirstFrameDecoder(model, sample_rate=24000) + decoder.capture() + codes = torch.arange(rows, device="cuda").reshape(-1, 1) + first = decoder.decode(codes) + second = decoder.decode(codes + 32) + assert first.dtype == second.dtype == torch.float32 + torch.testing.assert_close(first, codes.float(), rtol=0, atol=0) + torch.testing.assert_close(second, (codes + 32).float(), rtol=0, atol=0) diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_cudnn.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_cudnn.py new file mode 100644 index 000000000..d345473c9 --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_cudnn.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch +from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig + +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import Qwen3OmniMoeCode2Wav + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cuda, + pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"), +] + + +def _codec(enabled): + config = Qwen3OmniMoeCode2WavConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=8, + decoder_dim=32, + upsample_rates=[2, 2], + upsampling_ratios=[2], + codebook_size=16, + num_quantizers=2, + ) + config._attn_implementation = "sdpa" + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + hf_config=config, stage_connector_config={"extra": {"codec_cudnn_benchmark": enabled}} + ) + ) + model = Qwen3OmniMoeCode2Wav(vllm_config=vllm_config).cuda().eval() + model.precompute_snake_caches() + return model + + +def _flags(): + c = torch.backends.cudnn + return c.enabled, c.benchmark, c.benchmark_limit, c.deterministic, c.allow_tf32 + + +@torch.inference_mode() +def test_autotuned_codec_graph_replay_restores_flags(): + torch.manual_seed(21) + model = _codec(True) + codes = torch.randint(0, 16, (2, 2, 8), device="cuda") + original = _flags() + for _ in range(3): + model(codes) + assert _flags() == original + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = model(codes) + assert _flags() == original + codes.random_(0, 16) + graph.replay() + expected = model(codes) + torch.testing.assert_close(captured, expected) + assert _flags() == original + + +@pytest.mark.parametrize("enabled", [False, True]) +@torch.inference_mode() +def test_codec_failure_restores_cudnn_settings(mocker, enabled): + model = _codec(enabled) + codes = torch.randint(0, 16, (1, 2, 8), device="cuda") + original = _flags() + + def fail(hidden): + if enabled: + assert torch.backends.cudnn.benchmark + assert torch.backends.cudnn.benchmark_limit == 10 + else: + assert _flags() == original + raise RuntimeError("decode failed") + + mocker.patch.object(model, "_decode_waveform", side_effect=fail) + with pytest.raises(RuntimeError, match="decode failed"): + model(codes) + assert _flags() == original diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_snake.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_snake.py new file mode 100644 index 000000000..e18aaa109 --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_snake.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import copy +from types import SimpleNamespace + +import pytest +import torch +from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig +from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( + Qwen3OmniMoeCode2WavDecoderBlock, + Qwen3OmniMoeSnakeBeta, +) + +from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import Qwen3OmniMoeCode2Wav, use_fused_snake + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _blocks(): + config = Qwen3OmniMoeCode2WavConfig(decoder_dim=32, upsample_rates=[2, 2]) + torch.manual_seed(0) + blocks = torch.nn.ModuleList(Qwen3OmniMoeCode2WavDecoderBlock(config, i) for i in range(2)).eval() + with torch.no_grad(): + for module in blocks.modules(): + if isinstance(module, Qwen3OmniMoeSnakeBeta): + module.alpha.uniform_(-0.5, 0.5) + module.beta.uniform_(-0.5, 0.5) + return blocks + + +def test_decoder_blocks_swap_every_hf_snake_and_keep_weight_names(): + blocks = _blocks() + hf_count = sum(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in blocks.modules()) + keys = set(blocks.state_dict()) + + assert use_fused_snake(blocks) == hf_count > 0 + assert not any(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in blocks.modules()) + assert sum(isinstance(m, SnakeBeta) for m in blocks.modules()) == hf_count + assert set(blocks.state_dict()) == keys # checkpoint weights load unchanged + + +def test_fused_snake_blocks_match_hf_blocks(): + reference = _blocks() + fused = copy.deepcopy(reference) + use_fused_snake(fused) + fused.load_state_dict(reference.state_dict()) + x = torch.randn(2, 32, 20) + with torch.no_grad(): + expected = x + actual = x + for ref_block, fused_block in zip(reference, fused): + expected = ref_block(expected) + actual = fused_block(actual) + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + + +@pytest.mark.parametrize( + ("option", "device", "cuda", "fused"), + [ + (False, "cuda", True, False), + (True, "cuda", True, True), + (True, "cpu", True, False), + (True, "cuda", False, False), + ], +) +def test_decoder_block_fusion_requires_cuda_opt_in(monkeypatch, option, device, cuda, fused): + from vllm_omni.platforms import current_omni_platform + + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: cuda) + config = Qwen3OmniMoeCode2WavConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + decoder_dim=32, + ) + model = Qwen3OmniMoeCode2Wav( + vllm_config=SimpleNamespace( + model_config=SimpleNamespace( + hf_config=config, stage_connector_config={"extra": {"codec_fused_snake": option}} + ), + device_config=SimpleNamespace(device=device), + ) + ) + assert any(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in model.decoder.modules()) is (not fused) diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_first_audio.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_first_audio.py new file mode 100644 index 000000000..5ca8b7856 --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_first_audio.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch +from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig + +from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import Qwen3OmniFirstFrameDecoder +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import Qwen3OmniMoeForConditionalGeneration +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import ( + Qwen3OmniMoeCode2Wav, + plan_decode_groups, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_Q = 4 +_CODEBOOK = 16 + + +def _code2wav(quantizers=_Q) -> Qwen3OmniMoeCode2Wav: + torch.manual_seed(0) + config = Qwen3OmniMoeCode2WavConfig( + codebook_size=_CODEBOOK, + num_quantizers=quantizers, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=2, + decoder_dim=32, + sliding_window=8, + ) + model = Qwen3OmniMoeCode2Wav(vllm_config=SimpleNamespace(model_config=SimpleNamespace(hf_config=config))) + return model.eval() + + +def test_first_frame_pcm_equals_code2wav_streaming_chunk0(): + code2wav = _code2wav() + decoder = Qwen3OmniFirstFrameDecoder(code2wav, sample_rate=24000) + codes = torch.randint(0, _CODEBOOK, (3, _Q)) + with torch.inference_mode(): + for row in range(codes.shape[0]): + pcm = decoder.decode(codes[row : row + 1]) + chunk0 = code2wav.chunked_decode_streaming( + codes[row].reshape(1, _Q, 1), left_context_size=[0], seq_token_counts=[_Q] + )[0] + assert pcm.dtype == torch.float32 + assert torch.equal(pcm[0], chunk0.reshape(-1).float()) + assert 0 < pcm.shape[-1] < int(code2wav.total_upsample) + batched = decoder.decode(codes) + assert batched.shape == (3, pcm.shape[-1]) + + +def test_first_frame_is_a_prefix_of_a_longer_first_chunk(): + code2wav = _code2wav() + decoder = Qwen3OmniFirstFrameDecoder(code2wav, sample_rate=24000) + codes = torch.randint(0, _CODEBOOK, (_Q, 4)) + with torch.inference_mode(): + first = decoder.decode(codes[:, :1].T)[0] + chunk0 = code2wav.chunked_decode_streaming( + codes.reshape(1, _Q, 4), left_context_size=[0], seq_token_counts=[4 * _Q] + )[0] + torch.testing.assert_close(chunk0.reshape(-1)[: first.numel()].float(), first) + + +def test_chunk_ramp_adds_exact_code2wav_graph_sizes(monkeypatch): + import vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper as wrapper_module + + real = wrapper_module.CUDAGraphDecoderWrapper + wrapper = MagicMock(compute_capture_sizes=real.compute_capture_sizes) + monkeypatch.setattr(wrapper_module, "CUDAGraphDecoderWrapper", wrapper) + extra = {"codec_chunk_frames": 25, "codec_left_context_frames": 25, "codec_chunk_ramp": [1, 2, 4, 8, 16, 25]} + code2wav = SimpleNamespace(config=SimpleNamespace(num_quantizers=_Q)) + code2wav.enable_cudagraph = lambda **kw: Qwen3OmniMoeCode2Wav.enable_cudagraph( + code2wav, device=torch.device("cuda"), **kw + ) + model = object.__new__(Qwen3OmniMoeForConditionalGeneration) + torch.nn.Module.__init__(model) + model.code2wav = code2wav + model.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(async_chunk=True, enforce_eager=False, stage_connector_config={"extra": extra}) + ) + for batches in ([], [2, 4, 8]): + extra["codec_graph_batch_sizes"] = batches + Qwen3OmniMoeForConditionalGeneration._maybe_enable_code2wav_cudagraph(model) + args = wrapper.call_args.kwargs + sizes = args["capture_sizes"] + assert {1, 3, 7, 15, 31, 50} <= set(sizes) + assert set(real.compute_capture_sizes(codec_chunk_frames=25, codec_left_context_frames=25)) <= set(sizes) + assert set(args["extra_capture_shapes"]) == { + (b, size) for b in batches for size in sizes if size <= 50 and b * size <= 325 + } + assert code2wav._streaming_batch_sizes[50] == [1, 2, 4] + extra.pop("codec_chunk_ramp") + extra.pop("codec_graph_batch_sizes") + Qwen3OmniMoeForConditionalGeneration._maybe_enable_code2wav_cudagraph(model) + assert wrapper.call_args.kwargs["capture_sizes"] is None + + +_SIZES = [1, 2, 3, 4, 7, 8, 15, 16, 25, 31, 32, 50, 64] + + +def _bucket(length): + return next((size for size in _SIZES if size >= length), None) + + +def test_decode_groups_keep_a_uniform_batch_whole_and_split_by_the_largest_graph(): + assert plan_decode_groups([50] * 6, _bucket, lambda size: [1, 2, 4, 8]) == [(list(range(6)), 50)] + calls = plan_decode_groups([25] * 20, _bucket, lambda size: [1, 2, 4, 8, 16]) + assert [len(rows) for rows, _length in calls] == [16, 4] + calls = plan_decode_groups([1] * 8 + [50] * 8, _bucket, lambda size: [1, 2, 4, 8] if size < 50 else [1, 2, 4]) + assert [(len(rows), length) for rows, length in calls] == [(8, 1), (4, 50), (4, 50)] + assert plan_decode_groups([], _bucket, lambda size: [1, 2]) == [] + + +@pytest.mark.parametrize("frames", [65, 129]) +def test_decode_groups_beyond_largest_graph_use_eager(monkeypatch, frames): + from vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper import CUDAGraphDecoderWrapper + + wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper) + wrapper.enabled = wrapper._warmed_up = True + wrapper._bucket_sizes = [8, 32, 64] + wrapper._compiled_shapes = {(1, 64), (8, 64)} + wrapper._compiled_graphs = wrapper.graphs = dict.fromkeys(wrapper._compiled_shapes, MagicMock()) + wrapper.decoder = MagicMock(side_effect=lambda codes: codes.float() * 2) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + groups = plan_decode_groups([frames, frames], wrapper._get_padded_size, lambda _: (1,)) + assert groups == [([0], frames), ([1], frames)] + codes = torch.arange(2 * frames).reshape(2, 1, frames) + for rows, width in groups: + actual = wrapper.decode(codes[rows, :, :width]) + torch.testing.assert_close(actual, codes[rows].float() * 2, rtol=0, atol=0) + assert wrapper.decoder.call_count == 2 + for graph in (*wrapper.graphs.values(), *wrapper._compiled_graphs.values()): + graph.replay.assert_not_called() + + +class _PaddingGraphs: + def __init__(self, model, batch_sizes): + self.model = model + self.batch_sizes = sorted(batch_sizes) + self.calls = [] + + def _get_padded_size(self, length): + return _bucket(length) + + def decode(self, codes): + rows, _q, frames = codes.shape + size = _bucket(frames) + batch = next(b for b in self.batch_sizes if b >= rows) + self.calls.append((rows, frames)) + padded = torch.zeros(batch, codes.shape[1], size, dtype=codes.dtype) + padded[:rows, :, :frames] = codes + out = self.model(padded) + return out[:rows, :, : out.shape[-1] - (size - frames) * int(self.model.total_upsample)].clone() + + +def test_grouped_streaming_decode_matches_one_padded_batch(): + code2wav = _code2wav() + lengths = [1, 3, 50, 7, 31] + lefts = [0, 1, 25, 3, 15] + codes = torch.zeros(len(lengths), _Q, max(lengths), dtype=torch.long) + for row, length in enumerate(lengths): + codes[row, :, :length] = torch.randint(0, _CODEBOOK, (_Q, length)) + counts = [length * _Q for length in lengths] + with torch.inference_mode(): + reference = code2wav.chunked_decode_streaming(codes, left_context_size=lefts, seq_token_counts=counts) + graphs = _PaddingGraphs(code2wav, [1, 2, 4]) + code2wav._cudagraph_wrapper, code2wav._cudagraph_enabled = graphs, True + code2wav._streaming_batch_sizes = {size: [1, 2, 4] for size in _SIZES} + grouped = code2wav.chunked_decode_streaming(codes, left_context_size=lefts, seq_token_counts=counts) + + for want, got in zip(reference, grouped, strict=True): + assert want.shape == got.shape + torch.testing.assert_close(got, want, rtol=1e-4, atol=1e-5) + + +@pytest.mark.parametrize(("frames", "quantizers"), [(1, 16), (4, 16), (4, 4)]) +def test_delivered_first_frame_is_removed_from_code2wav_output(frames, quantizers): + from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY + + codec = _code2wav(quantizers) + model = object.__new__(Qwen3OmniMoeForConditionalGeneration) + torch.nn.Module.__init__(model) + model.model_stage, model.code2wav, model.code2wav_config = "code2wav", codec, codec.config + model.generate_audio = lambda codes, left, counts: codec.chunked_decode_streaming( + codes, left_context_size=left, seq_token_counts=counts + ) + codes = torch.randint(0, _CODEBOOK, (2, quantizers, frames)) + with torch.inference_mode(): + reference = model.generate_audio(codes, [0, 0], [quantizers * frames] * 2) + first = Qwen3OmniFirstFrameDecoder(codec, sample_rate=24000).decode(codes[0, :, :1].T)[0] + out = model.forward( + input_ids=codes.flatten(), + positions=None, + seq_token_counts=[quantizers * frames] * 2, + runtime_additional_information=[ + {"meta": {"left_context_size": 0, "first_audio": True}}, + {"meta": {"left_context_size": 0, "first_audio": False}}, + ], + ).multimodal_outputs + torch.testing.assert_close(torch.cat([first, out["model_outputs"][0].flatten()]), reference[0].flatten()) + torch.testing.assert_close(out["model_outputs"][1], reference[1].reshape(1, -1)) + assert [bool(flag) for flag in out[FIRST_AUDIO_REQUIRED_KEY]] == [True, False] diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_mrv2.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_mrv2.py new file mode 100644 index 000000000..6f5d8f6c7 --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_mrv2.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import Qwen3OmniMoeForConditionalGeneration + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_H = 4 +_Q = 3 + + +class _Talker(nn.Module): + num_code_groups = _Q + + def __init__(self): + super().__init__() + self.text_projection = nn.Identity() + self.code_predictor = SimpleNamespace( + _top_k=50, + _top_p=0.8, + set_sampling_params=lambda **kwargs: setattr(self, "sampling", kwargs), + ) + + +def _talker(*, async_chunk: bool = True) -> Qwen3OmniMoeForConditionalGeneration: + model = object.__new__(Qwen3OmniMoeForConditionalGeneration) + nn.Module.__init__(model) + model.model_stage = "talker" + model.talker = _Talker() + model.vllm_config = SimpleNamespace(model_config=SimpleNamespace(async_chunk=async_chunk, use_v2_model_runner=True)) + model.tts_eos_embed = torch.full((_H,), -1.0) + model.tts_pad_embed = torch.full((_H,), -2.0) + model._codec_codebook_size = 2048 + return model + + +def _rows(*values: float) -> torch.Tensor: + return torch.tensor([[value] * _H for value in values]) + + +def _step(model, payload: dict) -> torch.Tensor: + update: dict = {} + text_step = model._thinker_decode_to_talker_decode(payload, torch.device("cpu"), update) + for key, value in update.items(): + payload.setdefault(key, {}).update(value) + return text_step.reshape(-1) + + +def test_mtp_frame_valid_rejects_codec_eos_and_special_ids(): + layer0 = torch.tensor([0, 2047, 2048, 2150, -1]) + assert _talker().mtp_frame_valid(layer0).tolist() == [True, True, False, False, False] + + +def test_decode_rows_are_consumed_one_per_step_in_arrival_order(): + model = _talker() + # Two rows were already pending (cached at prefill); two more arrive at once. + payload = {"embed": {"cached_decode": _rows(1, 2), "decode": _rows(3, 4)}, "meta": {}} + steps = [_step(model, payload)[0].item()] + for arrival in (None, _rows(5), None, None, None, None): + payload["embed"]["decode"] = arrival + steps.append(_step(model, payload)[0].item()) + assert steps == [1, 2, 3, 4, 5, -1, -2] + assert payload["embed"]["cached_decode"] is None + + +@pytest.mark.parametrize(("v2", "resumable"), [(True, True), (False, False), (False, True)]) +def test_v1_and_resumable_requests_keep_the_indexed_cache_path(v2, resumable): + model = _talker() + model.vllm_config.model_config.use_v2_model_runner = v2 + payload = {"embed": {"cached_decode": _rows(1, 2, 3), "decode": None}, "meta": {"resumable": resumable}} + payload["meta"]["num_processed_tokens"] = 2 + assert _step(model, payload)[0].item() == 3 + + +def test_make_omni_output_aligns_codes_to_request_spans(): + model = _talker() + hidden = torch.zeros((8, _H)) # 5 real tokens + graph padding + decode_codes = torch.tensor([[5, 6, 7]]) + buffers = [ + {"req_id": "_warmup_0_"}, # no preprocess, no codes + {"req_id": "a", "codes": {"audio": torch.zeros((3, _Q), dtype=torch.long)}}, + {"req_id": "b", "codes": {"audio": decode_codes}}, + ] + out = model.make_omni_output( + hidden, model_intermediate_buffer=buffers, request_token_spans=[(0, 1), (1, 4), (4, 5)] + ) + codes = out.multimodal_outputs["codes"]["audio"] + assert codes.shape == (5, _Q) + assert codes[4].tolist() == [5, 6, 7] + assert out.text_hidden_states.shape[0] == 5 + + +def test_make_omni_output_without_spans_keeps_v1_concat(): + model = _talker() + buffers = [{"codes": {"audio": torch.tensor([[1, 2, 3]])}}, {"codes": {"audio": torch.tensor([[4, 5, 6]])}}] + out = model.make_omni_output(torch.zeros((4, _H)), model_intermediate_buffer=buffers) + assert out.multimodal_outputs["codes"]["audio"].tolist() == [[1, 2, 3], [4, 5, 6]] + assert out.text_hidden_states.shape[0] == 2 + + +def test_eager_output_is_token_major_with_invalid_rows(): + model = _talker() + model.eager_frames_active = True + hidden = torch.zeros((6, _H)) + buffers = [{"req_id": "a"}, {"req_id": "b"}] + out = model.make_omni_output(hidden, model_intermediate_buffer=buffers, request_token_spans=[(0, 3), (3, 4)]) + mm = out.multimodal_outputs + assert mm["codes"]["audio"].shape == (4, _Q) + assert mm["meta"]["codec_frame_valid"].tolist() == [0, 0, 0, 0] + assert out.text_hidden_states is hidden diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_pp_capture.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_pp_capture.py index de8d02e3e..fffb796ad 100644 --- a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_pp_capture.py +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_pp_capture.py @@ -104,6 +104,7 @@ def test_factory_declares_only_captures_from_earlier_ranks(monkeypatch, start, e if staged: model_config.model_stage = "thinker" vllm_config = SimpleNamespace(model_config=model_config, with_hf_config=lambda *args, **kwargs: None) + vllm_config.speculative_config = None thinker = nn.Module() thinker.weight = nn.Parameter(torch.ones(1)) thinker.language_model = SimpleNamespace(model=SimpleNamespace(start_layer=start)) diff --git a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py index fe19b44a9..52712a29c 100644 --- a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py +++ b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py @@ -1432,3 +1432,24 @@ def test_shared_predictor_ignores_tts_fast_path_options(mocker, loaded_target_cl expected = (logits.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True) actual = predictor._sample_per_call(logits, 1.0, 0, None, uniforms) torch.testing.assert_close(actual, expected) + + +def test_tts_cached_predictor_keeps_its_own_step_dispatch(mocker: MockerFixture, loaded_target_classes) -> None: + """The TTS cache must not enter Omni's uninitialized incremental loop.""" + _, _, wrapper, _, _ = loaded_target_classes + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + mocker.patch.object(common_mod.current_omni_platform, "is_cuda", return_value=True) + mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False) + mocker.patch.object(wrapper, "_stage_connector_extra_config", return_value={"code_predictor_kv_cache": True}) + cp_config, talker_config = _make_tiny_config(loaded_target_classes) + predictor = wrapper(vllm_config=_make_vllm_config(mocker), config=cp_config, talker_config=talker_config) + assert predictor._kv_requested + predictor._model_dtype = torch.float32 + predictor._codec_embeds_list = list(predictor.model.codec_embedding) + predictor._bucket_sizes = [1] + mocker.patch.object(predictor, "_setup_compile") + step = mocker.patch.object(predictor, "_predict_step_logits", return_value=torch.ones(1, cp_config.vocab_size)) + codes = predictor(torch.zeros(1, dtype=torch.long), torch.zeros(1, 32), torch.zeros(1, 32), do_sample=False) + assert step.call_count == cp_config.num_code_groups - 1 + assert codes.shape == (1, cp_config.num_code_groups) + assert predictor._kv_buf is None diff --git a/tests/model_executor/models/test_talker_first_audio.py b/tests/model_executor/models/test_talker_first_audio.py new file mode 100644 index 000000000..8d054290d --- /dev/null +++ b/tests/model_executor/models/test_talker_first_audio.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest + +from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import talker_first_audio_enabled + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize( + ("env", "extra", "async_chunk", "v2", "expected"), + [ + (None, {}, True, True, False), + ("1", {}, True, True, False), + (None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, True, True), + ("0", {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, True, False), + (None, {"talker_first_audio": True, "codec_chunk_ramp": [4, 8]}, True, True, False), + (None, {"talker_first_audio": True, "initial_codec_chunk_frames": 1}, True, True, True), + (None, {"talker_first_audio": True, "initial_codec_chunk_frames": 4}, True, True, False), + (None, {"talker_first_audio": True, "codec_chunk_ramp": [1]}, True, True, False), + (None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, False, True, False), + (None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, False, False), + ], +) +def test_talker_first_audio_default_and_conditions(monkeypatch, env, extra, async_chunk, v2, expected): + if env is None: + monkeypatch.delenv("VLLM_OMNI_TALKER_FIRST_AUDIO", raising=False) + else: + monkeypatch.setenv("VLLM_OMNI_TALKER_FIRST_AUDIO", env) + from vllm_omni.platforms import current_omni_platform + + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True) + config = SimpleNamespace( + model_config=SimpleNamespace( + async_chunk=async_chunk, use_v2_model_runner=v2, stage_connector_config={"extra": extra} + ), + device_config=SimpleNamespace(device="cuda"), + parallel_config=SimpleNamespace( + tensor_parallel_size=1, pipeline_parallel_size=1, distributed_executor_backend=None + ), + cache_config=SimpleNamespace(enable_prefix_caching=False), + ) + assert talker_first_audio_enabled(config) is expected + for section, field, value in [ + ("device_config", "device", "cpu"), + ("parallel_config", "tensor_parallel_size", 2), + ("parallel_config", "pipeline_parallel_size", 2), + ("parallel_config", "distributed_executor_backend", "mp"), + ("cache_config", "enable_prefix_caching", True), + ]: + with monkeypatch.context() as patch: + patch.setattr(getattr(config, section), field, value) + assert not talker_first_audio_enabled(config) + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False) + assert not talker_first_audio_enabled(config) + + +@pytest.mark.parametrize("option", [False, True]) +def test_tts_requires_connector_opt_in(monkeypatch, option): + from vllm_omni.model_executor.models.qwen3_tts import first_audio + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import ( + Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor, + ) + + monkeypatch.setattr(Predictor, "_stage_connector_extra_config", lambda _: {"talker_first_audio": option}) + monkeypatch.setattr(first_audio, "supports_talker_first_audio", lambda _: True) + assert first_audio.talker_first_audio_enabled(None) is option diff --git a/tests/model_executor/stage_input_processors/test_qwen3_omni_chunk_ramp.py b/tests/model_executor/stage_input_processors/test_qwen3_omni_chunk_ramp.py new file mode 100644 index 000000000..22fb3c9f3 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_qwen3_omni_chunk_ramp.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from collections import defaultdict +from types import SimpleNamespace + +import pytest +import torch + +import vllm_omni.model_executor.stage_input_processors.qwen3_omni as q3 +from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import ramp_decode_windows + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_Q = 2 +_RAMP = [1, 2, 4, 8, 16, 25] + + +def _manager(**extra): + config = {"codec_chunk_frames": 25, "codec_left_context_frames": 25, "initial_codec_chunk_frames": 4, **extra} + return SimpleNamespace( + code_prompt_token_ids=defaultdict(list), + put_req_chunk=defaultdict(int), + ramp_chunk_count=defaultdict(int), + connector=SimpleNamespace(config={"extra": config}), + ) + + +def _request(request_id="req"): + return SimpleNamespace(external_req_id=request_id, sampling_params=SimpleNamespace(stop_token_ids=[])) + + +def _frame(index: int) -> dict: + return { + "codes": {"audio": torch.tensor([[index, 1000 + index]], dtype=torch.long)}, + "meta": {"codec_frame_valid": torch.ones(1, dtype=torch.int8)}, + } + + +def _stream(manager, num_frames, finish_after_last=True): + request = _request() + chunks = [] + + def emit(payload): + if payload is None: + return + codes = payload.codes.audio.reshape(_Q, -1) + chunks.append((codes[0].tolist(), int(payload.meta.left_context_size), bool(payload.meta.finished))) + # The transport counts every enqueued chunk. + manager.put_req_chunk[request.external_req_id] += 1 + manager.ramp_chunk_count[request.external_req_id] += 1 + + for index in range(num_frames): + last = finish_after_last and index == num_frames - 1 + emit(q3.talker2code2wav_async_chunk(manager, _frame(index), request, is_finished=last)) + return chunks + + +def test_ramp_decode_windows_include_left_context(): + assert ramp_decode_windows(_RAMP, 25) == [1, 3, 7, 15, 31, 50] + assert ramp_decode_windows([1, 2], 0) == [1, 2] + + +def test_ramp_chunks_grow_and_resend_left_context(): + chunks = _stream(_manager(codec_chunk_ramp=_RAMP), 56 + 25, finish_after_last=False) + + new_frames = [len(window) - left for window, left, _finished in chunks] + assert new_frames == [1, 2, 4, 8, 16, 25, 25] + assert [left for _window, left, _finished in chunks] == [0, 1, 3, 7, 15, 25, 25] + assert [len(window) for window, _left, _finished in chunks] == [1, 3, 7, 15, 31, 50, 50] + emitted = 0 + for window, left, _finished in chunks: + # Each window ends at the newest frame and re-sends `left` earlier ones. + assert window == list(range(emitted - left, emitted + len(window) - left)) + emitted += len(window) - left + # The first chunk is the stream's first frame alone (no context). + assert chunks[0][0] == [0] + + +def test_ramp_flushes_partial_tail_with_left_context(): + chunks = _stream(_manager(codec_chunk_ramp=_RAMP), 10) + + assert [len(window) - left for window, left, _finished in chunks] == [1, 2, 4, 3] + window, left, finished = chunks[-1] + assert finished and left == 7 and window == list(range(0, 10)) + + +def test_ramp_finish_on_chunk_boundary_leaves_the_marker_to_the_transport(): + manager = _manager(codec_chunk_ramp=_RAMP) + request = _request() + for index in range(3): + payload = q3.talker2code2wav_async_chunk(manager, _frame(index), request) + if payload is not None: + manager.ramp_chunk_count[request.external_req_id] += 1 + # Frames 0 and 1-2 went out as chunks 0 and 1; the finish step has no frame. + finish = {"codes": {"audio": torch.zeros((1, _Q), dtype=torch.long)}, "meta": {"codec_frame_valid": torch.zeros(1)}} + assert q3.talker2code2wav_async_chunk(manager, finish, request, is_finished=True) is None + + +def test_without_ramp_the_initial_chunk_setting_is_unchanged(): + chunks = _stream(_manager(), 29 + 3) + + assert [(len(window), left) for window, left, _finished in chunks] == [(4, 0), (29, 4), (28, 25)] + + +def test_first_audio_marker_only_reaches_initial_codec_chunk(): + manager = _manager(codec_chunk_ramp=[1, 1]) + for index in range(2): + frame = _frame(index) + frame["meta"]["first_audio"] = torch.tensor([True]) + payload = q3.talker2code2wav_async_chunk(manager, frame, _request()) + assert bool(payload.meta.first_audio) is (index == 0) + manager.put_req_chunk["req"] += 1 + manager.ramp_chunk_count["req"] += 1 diff --git a/tests/model_executor/stage_input_processors/test_qwen3_omni_first_token_handoff.py b/tests/model_executor/stage_input_processors/test_qwen3_omni_first_token_handoff.py new file mode 100644 index 000000000..7e32f2c15 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_qwen3_omni_first_token_handoff.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from collections import defaultdict +from types import SimpleNamespace + +import pytest +import torch + +import vllm_omni.model_executor.stage_input_processors.qwen3_omni as q3 + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +ACCEPT = 24 +HIDDEN = 8 +PROMPT = [151644, 872, 5, 6, 151645, 151644, 77091, 198] +EOS = 151645 + + +class _Manager: + def __init__(self) -> None: + self.put_req_chunk: defaultdict[str, int] = defaultdict(int) + self.request_payload: dict[str, dict] = {} + self.config = SimpleNamespace( + hf_config=SimpleNamespace(talker_config=SimpleNamespace(accept_hidden_layer=ACCEPT)) + ) + + def process(self, output, request, is_finished=False): + payload = q3.thinker2talker_async_chunk(self, output, request, is_finished=is_finished) + if payload is not None: + self.put_req_chunk[request.external_req_id] += 1 + return payload + + +def _request(output_token_ids, *, history=True, max_tokens=64): + # The native data plane keeps token history only while chunk 0 is pending. + return SimpleNamespace( + external_req_id="r", + prompt_token_ids=list(PROMPT) if history else [], + output_token_ids=list(output_token_ids) if history else [], + all_token_ids=list(PROMPT) + list(output_token_ids) if history else [], + output_token_count=len(output_token_ids), + last_output_token_id=output_token_ids[-1] if output_token_ids else None, + resumable=False, + additional_information=None, + sampling_params=SimpleNamespace(all_stop_token_ids={EOS}, max_tokens=max_tokens), + ) + + +def _embed(token_id: int) -> torch.Tensor: + return torch.full((1, HIDDEN), float(token_id), dtype=torch.bfloat16) + + +def _output(token_ids, sampled=None): + layers = { + 0: torch.cat([_embed(t) for t in token_ids]), + ACCEPT: torch.randn(len(token_ids), HIDDEN).to(torch.bfloat16), + } + embed = {"tts_bos": _embed(1), "tts_eos": _embed(2), "tts_pad": _embed(3)} + if sampled is not None: + embed["sampled"] = sampled + return {"hidden_states": {"layers": layers}, "embed": embed} + + +def _run_turn(tokens, *, publish, max_tokens=64, prefill_chunks=(len(PROMPT),)): + manager = _Manager() + rows, chunks, start = [], 0, 0 + for i, size in enumerate(prefill_chunks): + last = i == len(prefill_chunks) - 1 + sampled = _embed(tokens[0]) if last else torch.empty(0, dtype=torch.bfloat16) + request = _request(tokens[:1] if last else [], max_tokens=max_tokens) + payload = manager.process(_output(PROMPT[start : start + size], sampled if publish else None), request) + start += size + if payload is not None: + chunks += 1 + rows.extend(payload.embed.prefill[len(PROMPT) :]) + for step in range(1, len(tokens)): + history = manager.put_req_chunk["r"] == 0 + request = _request(tokens[: step + 1], history=history, max_tokens=max_tokens) + output = _output([tokens[step - 1]], _embed(tokens[step]) if publish else None) + payload = manager.process(output, request) + if payload is None: + continue + chunks += 1 + rows.extend(payload.embed.prefill[len(PROMPT) :] if payload.embed.prefill is not None else payload.embed.decode) + return [int(row[0]) for row in rows], chunks, manager + + +@pytest.mark.parametrize( + ("tokens", "max_tokens"), + [ + ([11, 12, 13, EOS], 64), # ends on a stop token + ([11, 12, 13], 3), # ends on the max_tokens-th token + ([11, EOS], 64), + ], +) +def test_sampled_stream_matches_captured_stream(tokens, max_tokens): + late, _chunks, _manager = _run_turn(tokens, publish=False, max_tokens=max_tokens) + early, _chunks, manager = _run_turn(tokens, publish=True, max_tokens=max_tokens) + # The final token is never processed, so neither path feeds it to the Talker. + assert early == late == tokens[:-1] + assert q3._SAMPLED_TEXT_STREAM in manager.request_payload["r"] + + +def test_sampled_stream_sends_each_token_one_step_earlier(): + manager = _Manager() + chunk0 = manager.process(_output(PROMPT, sampled=_embed(11)), _request([11])) + assert chunk0 is not None # at the prefill step, not after the next decode step + assert torch.equal(chunk0.embed.prefill, torch.cat([_embed(t) for t in [*PROMPT, 11]])) + assert chunk0.hidden_states.output.shape[0] == len(PROMPT) + assert chunk0.ids.all == [*PROMPT, 11] + # Decode step 1 processes 11 and samples 12: it carries 12. + chunk1 = manager.process(_output([11], sampled=_embed(12)), _request([11, 12], history=False)) + assert torch.equal(chunk1.embed.decode, _embed(12)) + + +def test_chunked_prefill_uses_sample_of_final_chunk_only(): + early, chunks, _manager = _run_turn([11, 12, EOS], publish=True, prefill_chunks=(5, 3)) + assert early == [11, 12] and chunks == 2 + + +def test_terminal_first_token_keeps_captured_stream(): + manager = _Manager() + assert manager.process(_output(PROMPT, sampled=_embed(EOS)), _request([EOS])) is None + assert q3._SAMPLED_TEXT_STREAM not in manager.request_payload["r"] + + +@pytest.mark.parametrize("prefill_rows", [0, 1, 3]) +def test_missing_sample_on_sampled_stream(prefill_rows): + manager = _Manager() + manager.process(_output(PROMPT, sampled=_embed(11)), _request([11])) + chunk = manager.process(_output([11], sampled=_embed(12)), _request([11, 12], history=False)) + assert torch.equal(chunk.embed.decode, _embed(12)) + if not prefill_rows: + with pytest.raises(RuntimeError, match="sample embedding missing"): + manager.process(_output([12]), _request([11, 12], history=False)) + else: + before = manager.put_req_chunk["r"] + for tokens in (PROMPT[:prefill_rows], [11]): + partial = _output(tokens, sampled=torch.empty(0)) + assert manager.process(partial, _request([11, 12], history=False)) is None + assert manager.put_req_chunk["r"] == before + chunk = manager.process(_output([12], sampled=_embed(13)), _request([11, 12, 13], history=False)) + assert torch.equal(chunk.embed.decode, _embed(13)) diff --git a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py index dfbb9d91d..f67fc7a58 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py @@ -1047,3 +1047,80 @@ def test_qwen2_5_omni_thinker2talker_full_payload_noop() -> None: payload = thinker2talker_full_payload(None, {"any": "thing"}, None) assert payload is None + + +def _codec_transfer_manager(request_id: str, *, initial_frames: int = 2, chunk_frames: int = 4): + return SimpleNamespace( + code_prompt_token_ids=defaultdict(list), + put_req_chunk=defaultdict(int, {request_id: 0}), + connector=SimpleNamespace( + config={ + "extra": { + "initial_codec_chunk_frames": initial_frames, + "codec_chunk_frames": chunk_frames, + "codec_left_context_frames": 4, + } + } + ), + ) + + +def _codec_request(request_id: str) -> SimpleNamespace: + return SimpleNamespace( + external_req_id=request_id, + sampling_params=SimpleNamespace(stop_token_ids=[2150], stop_token_id=None), + ) + + +def test_talker2code2wav_async_chunk_reads_last_valid_row_of_multi_row_span() -> None: + # MRv2 eager frames: a prefill span carries its first frame in its last row. + request_id = "eager_prefill" + manager = _codec_transfer_manager(request_id) + codes = torch.zeros((5, 3), dtype=torch.long) + codes[-1] = torch.tensor([11, 12, 13]) + valid = torch.tensor([0, 0, 0, 0, 1], dtype=torch.int8) + + payload = q3.talker2code2wav_async_chunk( + manager, {"codes": {"audio": codes}, "meta": {"codec_frame_valid": valid}}, _codec_request(request_id) + ) + + assert payload is None # below the initial chunk size + assert [frame.tolist() for frame in manager.code_prompt_token_ids[request_id]] == [[[11, 12, 13]]] + + +def test_talker2code2wav_async_chunk_skips_invalid_frame_and_flushes_on_finish() -> None: + request_id = "eager_eos" + manager = _codec_transfer_manager(request_id) + manager.code_prompt_token_ids[request_id] = [torch.tensor([[7, 8, 9]])] + # The step that sampled codec EOS carries an invalid (placeholder) frame. + output = { + "codes": {"audio": torch.tensor([[2150, 0, 0]])}, + "meta": {"codec_frame_valid": torch.tensor([0], dtype=torch.int8)}, + } + + payload = q3.talker2code2wav_async_chunk(manager, output, _codec_request(request_id), is_finished=True) + + assert len(manager.code_prompt_token_ids[request_id]) == 1 + assert payload is not None + assert payload.codes.audio.tolist() == [7, 8, 9] + assert payload.meta.finished.item() is True + + +def test_talker2code2wav_async_chunk_without_validity_keeps_one_row_contract() -> None: + # V1 / MRv2 deferred frames: one row per step, zero rows are prefill + # placeholders and a stop-token CB0 marks the post-EOS step. + request_id = "legacy_rows" + manager = _codec_transfer_manager(request_id) + request = _codec_request(request_id) + + assert ( + q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.zeros((4, 3), dtype=torch.long)}}, request) + is None + ) + assert q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.tensor([[2150, 1, 1]])}}, request) is None + assert manager.code_prompt_token_ids[request_id] == [] + q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.tensor([[5, 6, 7]])}}, request) + payload = q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.tensor([[8, 9, 10]])}}, request) + + assert payload is not None # two frames fill the initial chunk + assert payload.codes.audio.tolist() == [5, 8, 6, 9, 7, 10] diff --git a/tests/worker_v2/test_omni_model_state.py b/tests/worker_v2/test_omni_model_state.py index b0178e9a5..4edfed603 100644 --- a/tests/worker_v2/test_omni_model_state.py +++ b/tests/worker_v2/test_omni_model_state.py @@ -620,6 +620,12 @@ def test_publish_sampled_embeddings_for_rows_whose_sample_is_kept(monkeypatch) - assert [tuple(row.shape) for row in sampled] == [(1, 3), (1, 3), (0,), (1, 3)] assert [row.tolist() for row in sampled if row.numel()] == [[[v] * 3] for v in (11.0, 12.0, 14.0)] + batch.is_prefilling_np[:] = True + batch.num_computed_prefill_tokens_np[:] = 0 + batch.num_scheduled_tokens = [1] * 4 + extra, _done = state.publish_sampled_embeddings(batch, torch.tensor([[11], [12], [13], [14]])) + assert all(row.numel() == 0 for row in extra["embed"]["sampled"]) + # Speculative steps sample several tokens per row: not published. assert state.publish_sampled_embeddings(batch, torch.tensor([[11, 1], [12, 1], [13, 1], [14, 1]])) is None diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index 69d9306a7..ee7f599f2 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -260,6 +260,7 @@ _BENCHMARK_TRANSITIONAL = ( ) _INTERNAL = ( + "VLLM_OMNI_TALKER_FIRST_AUDIO", "VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS", "VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS_SLACK_S", "VLLM_OMNI_CONNECTOR_RECV_POLL_MS", diff --git a/vllm_omni/deploy/qwen3_omni_moe_mrv2.yaml b/vllm_omni/deploy/qwen3_omni_moe_mrv2.yaml new file mode 100644 index 000000000..c40a29209 --- /dev/null +++ b/vllm_omni/deploy/qwen3_omni_moe_mrv2.yaml @@ -0,0 +1,50 @@ +# Opt-in CUDA Model Runner V2 profile for Qwen3-Omni-MoE: the Thinker (stage 0), +# Talker (stage 1) and Code2Wav (stage 2) run on MRv2 with its native data +# plane. The base profile keeps every stage on V1. +# +# On MRv2 the Thinker hands each generated token's embedding to the Talker with +# the step that sampled it (one Thinker step earlier than its next forward would +# capture the same embedding) and replays decode in FULL CUDA graphs together +# with the layers it captures for the Talker. +# +# When the stages share one GPU, run them under NVIDIA MPS (see +# docs/configuration/stage_configs.md); startup warns when no MPS server is found. +# +# MRv2 downstream stages support turn-based sessions only: realtime (resumable) +# requests and duplex sessions need the V1 chunk adapter. +base_config: qwen3_omni_moe.yaml + +connectors: + connector_of_shared_memory: + extra: + # First Code2Wav chunk is the stream's first frame alone, then chunks + # double up to the steady codec_chunk_frames (replaces the 4-frame + # initial chunk). Code2Wav re-decodes up to codec_left_context_frames + # before each chunk, so windows are 1, 3, 7, 15, 31, then 50 frames. + # A one-frame first chunk is also what lets the Talker deliver it + # itself. Opt in explicitly; VLLM_OMNI_TALKER_FIRST_AUDIO=0 disables it. + talker_first_audio: true + codec_chunk_ramp: [1, 2, 4, 8, 16, 25] + # Code2Wav decodes a streaming step as a few length groups on multi-row + # CUDA graphs instead of one eager call padded to its longest window + # (with the ramp, over half of those padded frames were waste). + codec_graph_batch_sizes: [2, 3, 4, 6, 8, 12, 16] + codec_fused_snake: true + +stages: + - stage_id: 0 + model_runner: v2 + - stage_id: 1 + model_runner: v2 + - stage_id: 2 + model_runner: v2 + +platforms: + npu: + model_runner: v1 + xpu: + model_runner: v1 + rocm: + model_runner: v1 + musa: + model_runner: v1 diff --git a/vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml b/vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml new file mode 100644 index 000000000..a5835dc02 --- /dev/null +++ b/vllm_omni/deploy/qwen3_omni_moe_mrv2_h200.yaml @@ -0,0 +1,27 @@ +# Opt-in single-H200 MRv2 profile, measured at client concurrency 32/64. +# Run all stages under NVIDIA MPS and warm representative shapes first. +# Retains Triton MoE and native nucleus reductions, fusing only the final +# sampling mask/Gumbel selection. cuDNN searches convolution algorithms +# during warmup; short-KV attention and convolution selection change rounding. +# This is not a claim of bitwise audio or universal P99/throughput improvement. +# Evidence: docs/design/qwen3_omni_mrv2_performance.md +base_config: qwen3_omni_moe_mrv2.yaml + +connectors: + connector_of_shared_memory: + extra: + code_predictor_kv_cache: true + code_predictor_short_kv: true + code_predictor_fused_sampling: true + codec_cudnn_benchmark: true + +stages: + - stage_id: 0 + devices: "0" + gpu_memory_utilization: 0.60 + - stage_id: 1 + devices: "0" + gpu_memory_utilization: 0.22 + - stage_id: 2 + devices: "0" + gpu_memory_utilization: 0.10 diff --git a/vllm_omni/model_executor/models/common/fused_sampling.py b/vllm_omni/model_executor/models/common/fused_sampling.py new file mode 100644 index 000000000..a4bfce13e --- /dev/null +++ b/vllm_omni/model_executor/models/common/fused_sampling.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Fused Gumbel selection after native nucleus probability reductions. + +Keep ATen's top-k threshold, sort, softmax and cumulative sum, including +boundary ties. Random numbers remain indexed by original token id. Only the +final mask, Gumbel transform and argmax are fused; no vocabulary scatter is +needed. This avoids changing which tokens survive a nucleus boundary through +different reduction orders. +""" + +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def _sorted_nucleus_gumbel( + logits, + indices, + probabilities, + cumulative, + uniforms, + output, + vocab: tl.constexpr, + logits_stride: tl.constexpr, + token_stride: tl.constexpr, + uniform_stride: tl.constexpr, + top_p: tl.constexpr, + block: tl.constexpr, +): + row = tl.program_id(0) + offsets = tl.arange(0, block) + valid = offsets < vocab + values = tl.load(logits + row * logits_stride + offsets * token_stride, valid, other=-float("inf")).to(tl.float32) + if top_p < 1.0: + ids = tl.load(indices + row * vocab + offsets, valid, other=0) + probs = tl.load(probabilities + row * vocab + offsets, valid, other=0) + cdf = tl.load(cumulative + row * vocab + offsets, valid, other=0) + values = tl.where((cdf - probs) >= top_p, -float("inf"), values) + else: + ids = offsets + noise = tl.load(uniforms + row * uniform_stride + ids, valid, other=0.5) + score = values - tl.extra.cuda.libdevice.log(-tl.extra.cuda.libdevice.log(noise)) + # torch.argmax returns the first original token for a tie or NaN. + first_nan = tl.min(tl.where(valid & (score != score), ids, 2147483647), 0) + maximum = tl.max(tl.where(valid, score, -float("inf")), 0) + winner = tl.min(tl.where(valid & (score == maximum), ids, 2147483647), 0) + tl.store(output + row, tl.where(first_nan < 2147483647, first_nan, winner)) + + +def sample_top_k_top_p_gumbel( + logits: torch.Tensor, uniforms: torch.Tensor, *, top_k: int, top_p: float +) -> torch.Tensor: + """Sample CUDA rows without scattering sorted logits back to the vocabulary. + + Uniforms are caller-owned FP32 values with the same shape as logits; their + generation and generator advancement are deliberately outside this kernel. + Softmax/scan use the reference ATen operations. The Gumbel logarithms + still use libdevice, so this is not a universal bitwise-equivalence claim. + """ + if logits.ndim != 2 or not logits.is_cuda or uniforms.device != logits.device or torch.version.hip is not None: + raise ValueError("fused sampling requires two-dimensional CUDA logits and uniforms") + if uniforms.shape != logits.shape or uniforms.dtype != torch.float32 or uniforms.stride(1) != 1: + raise ValueError("uniforms must be FP32 with the logits shape and contiguous rows") + rows, vocab = logits.shape + # The stored-mode caller disables top-k for both 0 and negative values + # (vLLM commonly uses -1). Keep that behavior when fusion is enabled. + top_k = max(top_k, 0) + if vocab == 0 or top_k > vocab or not 0 < top_p <= 1: + raise ValueError("invalid top-k/top-p sampling parameters") + if top_k: + threshold = logits.topk(top_k, dim=-1).values[:, -1:] + logits = logits.masked_fill(logits < threshold, -float("inf")) + indices = probabilities = cumulative = None + if top_p < 1: + logits, indices = logits.sort(dim=-1, descending=True) + probabilities = logits.softmax(dim=-1, dtype=torch.float32) + cumulative = probabilities.cumsum(-1) + result = torch.empty((rows, 1), dtype=torch.long, device=logits.device) + if rows: + _sorted_nucleus_gumbel[(rows,)]( + logits, + indices, + probabilities, + cumulative, + uniforms, + result, + vocab, + logits.stride(0), + logits.stride(1), + uniforms.stride(0), + top_p, + triton.next_power_of_2(vocab), + num_warps=4, + enable_fp_fusion=False, + ) + return result diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index ce2cd1e5e..37a9c2388 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -18,6 +18,7 @@ Shared by Qwen3-Omni and Qwen3-TTS talker models. from __future__ import annotations import dataclasses +import time from collections.abc import Iterable, Sequence import torch @@ -158,8 +159,8 @@ class CodePredictorAttention(nn.Module): """Multi-head self-attention for code predictor. Uses ``F.scaled_dot_product_attention`` with HF-compatible RoPE and RMSNorm. - No KV cache -- the code predictor always re-prefills the full (short) - sequence each AR step. + The default path re-prefills the short sequence; the opt-in cached path + appends only new positions. Input : [B, seq_len, hidden_size] Output: [B, seq_len, hidden_size] @@ -316,6 +317,53 @@ class CodePredictorAttention(nn.Module): attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1) return self.o_proj(attn_out) + def forward_cached( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + k_cache: torch.Tensor, + v_cache: torch.Tensor, + write_positions: torch.Tensor, + key_mask: torch.Tensor, + ) -> torch.Tensor: + """Attend only the new positions, reusing earlier positions' keys and values. + + ``k_cache``/``v_cache`` are ``[B, num_kv_heads, max_seq, head_dim]``; + the new positions' keys and values are written at ``write_positions`` + (``[n]``). ``key_mask`` (``[1, 1, n, max_seq]``, bool) keeps key + positions up to each query's, so slots not yet written in this call are + never read and every single-position step has the same shapes. + """ + bsz, seq_len, _ = hidden_states.shape + qkv = self.qkv_proj(hidden_states) + q_raw, k_raw, v_raw = self._split_qkv(qkv) + q = self.q_norm(q_raw.view(bsz, seq_len, self.num_heads, self.head_dim)).transpose(1, 2) + k = self.k_norm(k_raw.view(bsz, seq_len, self.num_kv_heads, self.head_dim)).transpose(1, 2) + v = v_raw.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) + + cos, sin = position_embeddings + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + q = (q * cos) + (_rotate_half(q) * sin) + k = (k * cos) + (_rotate_half(k) * sin) + k_cache.index_copy_(2, write_positions, k.to(k_cache.dtype)) + v_cache.index_copy_(2, write_positions, v.to(v_cache.dtype)) + if getattr(self, "use_short_kv_attention", False): + from .short_kv_attention import short_kv_attention + + attn_out = short_kv_attention(q, k_cache, v_cache, write_positions, self.scaling) + else: + attn_out = F.scaled_dot_product_attention( + q, + k_cache, + v_cache, + attn_mask=key_mask, + scale=self.scaling, + enable_gqa=self.is_gqa, + ) + attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1) + return self.o_proj(attn_out) + # =================================================================== # MLP @@ -351,7 +399,7 @@ class CodePredictorMLP(nn.Module): class CodePredictorDecoderLayer(nn.Module): - """Transformer decoder layer (SDPA, no KV cache).""" + """Transformer decoder layer with default re-prefill and opt-in KV reuse.""" def __init__(self, config, *, prefix: str = "") -> None: super().__init__() @@ -383,9 +431,29 @@ class CodePredictorDecoderLayer(nn.Module): hidden_states = residual + hidden_states return hidden_states + def forward_cached( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + k_cache: torch.Tensor, + v_cache: torch.Tensor, + write_positions: torch.Tensor, + key_mask: torch.Tensor, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn.forward_cached( + hidden_states, position_embeddings, k_cache, v_cache, write_positions, key_mask + ) + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + return residual + hidden_states + # =================================================================== -# Base Transformer Model (re-prefill, no KV cache) +# Base Transformer Model (re-prefill, or incremental with a K/V buffer) # =================================================================== @@ -454,6 +522,29 @@ class CodePredictorBaseModel(nn.Module): hidden_states = self.norm(hidden_states) return hidden_states.to(input_dtype) + def forward_cached( + self, + inputs_embeds: torch.Tensor, + position_ids: torch.Tensor, + kv_caches: Sequence[tuple[torch.Tensor, torch.Tensor]], + write_positions: torch.Tensor, + key_mask: torch.Tensor, + ) -> torch.Tensor: + """``forward`` for the new positions only (``position_ids`` ``[B, n]``), reusing earlier K/V.""" + input_dtype = inputs_embeds.dtype + use_fp32 = input_dtype == torch.float16 and inputs_embeds.device.type not in ("cpu", "npu") + if use_fp32: + inputs_embeds = inputs_embeds.float() + hidden_states = inputs_embeds + with torch.amp.autocast(inputs_embeds.device.type, enabled=use_fp32, dtype=torch.float32): + position_embeddings = self.rotary_emb(hidden_states, position_ids) + for layer, (k_cache, v_cache) in zip(self.layers, kv_caches): + hidden_states = layer.forward_cached( + hidden_states, position_embeddings, k_cache, v_cache, write_positions, key_mask + ) + hidden_states = self.norm(hidden_states) + return hidden_states.to(input_dtype) + # Mapping from HF checkpoint shard name to the (fused parameter, # slice index) pair used to assemble the fused tensor. See # ``load_weights`` below. @@ -626,8 +717,14 @@ class CodePredictorWrapper(nn.Module): 4. torch.compile on inner transformer. 5. Inline sampling (top-k + top-p) -- no custom op overhead. 6. Optional manual CUDA graph capture per batch-size bucket. + + ``code_predictor_kv_cache`` (CUDA, opt-in) switches to incremental K/V + decoding: at large batches the re-prefill FLOPs are not negligible. """ + _kv_cache_enabled: bool = False + _fused_sampling_enabled: bool = False + def __init__( self, *, @@ -672,6 +769,7 @@ class CodePredictorWrapper(nn.Module): # Sampling defaults for "stored" mode self._top_k: int = 50 self._top_p: float = 0.8 + self._do_sample: bool = True # Lazily initialised state self._proj_buf: torch.Tensor | None = None @@ -686,6 +784,7 @@ class CodePredictorWrapper(nn.Module): # ``None`` keeps the legacy power-of-two bucket derivation. self._execution_batch_buckets: list[int] | None = None prefix_graph_cfg = self._stage_connector_extra_config(vllm_config) + self._fused_sampling_enabled = self._parse_bool_config(prefix_graph_cfg.get("code_predictor_fused_sampling")) prefix_graphs_requested = self._parse_bool_config(prefix_graph_cfg.get("code_predictor_prefix_graphs")) is_npu = current_omni_platform.is_npu() # MRv2 captures the whole Talker MTP call. It can therefore capture @@ -705,15 +804,28 @@ class CodePredictorWrapper(nn.Module): prefix_graph_cfg.get("code_predictor_prefix_graph_seq_lens") ) self._prefix_reprefill_seq_lens = tuple(self._prefix_seq_lens(self._num_groups + 1)) + # Incremental decoding (opt-in, CUDA): each AR step runs only its new + # position against a static K/V buffer instead of re-prefilling the + # whole prefix (at batch 64 the re-prefill runs 17x the positions). + self._kv_cache_enabled = ( + self._parse_bool_config(prefix_graph_cfg.get("code_predictor_kv_cache")) and current_omni_platform.is_cuda() + ) + self._short_kv_requested = self._parse_bool_config(prefix_graph_cfg.get("code_predictor_short_kv")) + self._short_kv_enabled = False + self._kv_buf: torch.Tensor | None = None + self._compiled_kv_fwd = None + # Per AR step: (new positions [n], key mask [1, 1, n, max_seq]). + self._kv_steps: dict[int, tuple[torch.Tensor, torch.Tensor]] = {} def get_input_embeddings(self) -> nn.ModuleList: return self.model.get_input_embeddings() - def set_sampling_params(self, top_k: int = 50, top_p: float = 0.8) -> None: - """Configure sampling parameters to maintain consistency with previous implementation.""" + def set_sampling_params(self, top_k: int = 50, top_p: float = 0.8, do_sample: bool = True) -> None: + """Configure "stored"-mode sampling; ``do_sample=False`` selects greedy argmax.""" self._top_k = top_k self._top_p = top_p - logger.debug("Sampling parameters updated: top_k=%d, top_p=%.2f", top_k, top_p) + self._do_sample = do_sample + logger.debug("Sampling parameters updated: top_k=%d, top_p=%.2f, do_sample=%s", top_k, top_p, do_sample) # ------------------------------------------------------------------ # Lazy-init helpers @@ -730,6 +842,35 @@ class CodePredictorWrapper(nn.Module): ): return self._proj_buf = torch.zeros(bsz, max_seq, self._cp_hidden, dtype=dtype, device=device) + if self._kv_cache_enabled: + cfg = self.config + head_dim = int(getattr(cfg, "head_dim", None) or cfg.hidden_size // cfg.num_attention_heads) + self._kv_buf = torch.zeros( + int(cfg.num_hidden_layers), + 2, + bsz, + int(cfg.num_key_value_heads), + max_seq, + head_dim, + dtype=dtype, + device=device, + ) + + def _kv_views(self, bsz: int) -> list[tuple[torch.Tensor, torch.Tensor]]: + kv_buf = self._kv_buf + assert kv_buf is not None + return [(kv_buf[i, 0, :bsz], kv_buf[i, 1, :bsz]) for i in range(kv_buf.shape[0])] + + def _ensure_kv_steps(self, device: torch.device) -> None: + if self._kv_steps: + return + max_seq = self._num_groups + 1 + keys = torch.arange(max_seq, device=device) + for step in range(1, self._num_groups): + # Step 1 runs positions 0 (talker hidden) and 1 (CB0); step s > 1 runs position s. + positions = torch.arange(0, 2, device=device) if step == 1 else torch.tensor([step], device=device) + key_mask = (keys[None, :] <= positions[:, None]).view(1, 1, positions.numel(), max_seq) + self._kv_steps[step] = (positions, key_mask) def _setup_compile(self) -> None: """Lazily set up torch.compile with optional device graph capture.""" @@ -742,6 +883,9 @@ class CodePredictorWrapper(nn.Module): self._model_dtype = next(self.model.parameters()).dtype self._lm_heads_list = list(self.lm_head) self._codec_embeds_list = list(self.model.codec_embedding) + if self._model_dtype == torch.float16 or not current_omni_platform.supports_torch_inductor(): + # The fp16 body runs in fp32 under autocast; keep re-prefill there. + self._kv_cache_enabled = False # Torch 2.13 XPU Dynamo can double-register built-in handlers when # spawned workers compile this predictor. Keep this narrow path eager @@ -768,10 +912,28 @@ class CodePredictorWrapper(nn.Module): dynamic=False, options={"epilogue_fusion": False}, ) + self._short_kv_enabled = ( + self._kv_cache_enabled + and self._short_kv_requested + and current_omni_platform.is_cuda() + and self._model_dtype == torch.bfloat16 + ) + for layer in self.model.layers: + layer.self_attn.use_short_kv_attention = self._short_kv_enabled + if self._kv_cache_enabled: + self._compiled_kv_fwd = torch.compile( + self.model.forward_cached, + dynamic=self._short_kv_enabled, + options={"epilogue_fusion": False}, + ) with torch._dynamo.config.patch(cache_size_limit=self._compile_cache_size_limit()): self._warmup_buckets() - if self._wrapper_config.use_cuda_graphs: + if self._kv_cache_enabled: + # Fifteen short steps per call: they rely on the runner capturing + # the whole MTP call (MRv2) rather than on per-step inner graphs. + logger.info("code_predictor: torch.compile + incremental K/V decoding") + elif self._wrapper_config.use_cuda_graphs: self._capture_cuda_graphs() logger.info("code_predictor: torch.compile (no epilogue fusion) + CUDA graphs") else: @@ -840,13 +1002,13 @@ class CodePredictorWrapper(nn.Module): return row_generators @classmethod - def _sample_codes_gumbel( + def _sampling_uniforms( cls, logits: torch.Tensor, generator: _GeneratorLike = None, uniforms: torch.Tensor | None = None, ) -> torch.Tensor: - """Sample ``logits`` via Gumbel-max with optional precomputed noise.""" + """Draw one FP32 uniform per original token, preserving generator state.""" if uniforms is None: row_generators = cls._normalize_generators(generator, int(logits.shape[0])) u = torch.empty_like(logits, dtype=torch.float32) @@ -862,6 +1024,17 @@ class CodePredictorWrapper(nn.Module): f"logits={tuple(logits.shape)}" ) u = uniforms + return u + + @classmethod + def _sample_codes_gumbel( + cls, + logits: torch.Tensor, + generator: _GeneratorLike = None, + uniforms: torch.Tensor | None = None, + ) -> torch.Tensor: + """Sample ``logits`` via Gumbel-max with optional precomputed noise.""" + u = cls._sampling_uniforms(logits, generator, uniforms) return (logits.float() - torch.log(-torch.log(u))).argmax(dim=-1, keepdim=True) def _prefix_seq_lens(self, max_seq: int) -> list[int]: @@ -924,6 +1097,8 @@ class CodePredictorWrapper(nn.Module): if needs_full_graph: full_graph_entries += len(prefix_buckets) required_entries = full_graph_entries + len(prefix_buckets) * len(prefix_seq_lens) + if self._kv_cache_enabled: + required_entries = 4 # {first step, later steps} x {batch 1, dynamic batch} return max(torch._dynamo.config.cache_size_limit, required_entries) def _warmup_buckets(self) -> None: @@ -938,7 +1113,29 @@ class CodePredictorWrapper(nn.Module): self._ensure_buffers(device, self._model_dtype, max(self._bucket_sizes)) proj_buf = self._proj_buf - if self._prefix_reprefill_enabled: + if self._kv_cache_enabled: + self._ensure_kv_steps(device) + started = time.perf_counter() + for bsz in self._bucket_sizes: + caches = self._kv_views(bsz) + for step in (1, 2): + positions, key_mask = self._kv_steps[step] + lo = 0 if step == 1 else step + hidden = proj_buf[:bsz, lo : step + 1, :] + pos_ids = positions.unsqueeze(0).expand(bsz, -1) + if bsz > 1: + # One graph for every batch >= 2 per step shape (4 compiles, not 2 per bucket): + # a cold per-bucket compile outlasted the stage handshake timeout. + for tensor in (hidden, pos_ids, *(t for kv in caches for t in kv)): + torch._dynamo.mark_dynamic(tensor, 0) + for _ in range(2): + self._compiled_kv_fwd(hidden, pos_ids, caches, positions, key_mask) + logger.info( + "code_predictor: incremental K/V decoding warmed for buckets %s in %.1f s", + self._bucket_sizes, + time.perf_counter() - started, + ) + elif self._prefix_reprefill_enabled: prefix_seq_lens = self._prefix_reprefill_seq_lens needs_full_graph = set(prefix_seq_lens) != set(range(2, max_seq)) for bsz in self._bucket_sizes: @@ -1206,9 +1403,29 @@ class CodePredictorWrapper(nn.Module): all_codes = torch.empty(bsz, num_groups, dtype=torch.long, device=device) all_codes[:, 0] = layer0_code.reshape(bsz) + kv_caches = None + if self._kv_cache_enabled: + # Masked slots still enter the attention matmul with zero weight, so + # clear them: a non-finite value left by an earlier call would leak. + if not self._short_kv_enabled: + self._kv_buf[:, :, :padded_bsz].zero_() + kv_caches = self._kv_views(padded_bsz) # Autoregressive loop: predict layers 1..G-1 for step in range(1, num_groups): - logits = self._predict_step_logits(proj_buf, bsz, padded_bsz, step, is_npu_capturing) + if kv_caches is not None: + # Step 1 runs positions 0-1; later steps append one position. + positions, key_mask = self._kv_steps[step] + lo = 0 if step == 1 else step + new_hidden = self._compiled_kv_fwd( + proj_buf[:padded_bsz, lo : step + 1, :], + positions.unsqueeze(0).expand(padded_bsz, -1), + kv_caches, + positions, + key_mask, + ) + logits = self._lm_heads_list[step - 1](new_hidden[:bsz, -1, :]) + else: + logits = self._predict_step_logits(proj_buf, bsz, padded_bsz, step, is_npu_capturing) # Sample next code via Gumbel-max. # @@ -1222,7 +1439,15 @@ class CodePredictorWrapper(nn.Module): # defensive than ``multinomial`` around fully-masked/NaN inputs), # and the helper below can honor either one batch generator or one # generator per seeded row. - if stored_mode: + if stored_mode and not self._do_sample: + code = logits.argmax(dim=-1, keepdim=True) + elif stored_mode and self._fused_sampling_enabled and logits.is_cuda and torch.version.hip is None: + from .fused_sampling import sample_top_k_top_p_gumbel + + step_uniforms = sample_uniforms[:, step - 1, :] if sample_uniforms is not None else None + uniforms = self._sampling_uniforms(logits, sample_generator, step_uniforms) + code = sample_top_k_top_p_gumbel(logits, uniforms, top_k=s_top_k, top_p=s_top_p) + elif stored_mode: # "stored" mode: top-k -> top-p -> Gumbel-max if s_top_k > 0: topk_vals, _ = logits.topk(s_top_k, dim=-1) diff --git a/vllm_omni/model_executor/models/qwen3_tts/short_kv_attention.py b/vllm_omni/model_executor/models/common/short_kv_attention.py similarity index 100% rename from vllm_omni/model_executor/models/qwen3_tts/short_kv_attention.py rename to vllm_omni/model_executor/models/common/short_kv_attention.py diff --git a/vllm_omni/model_executor/models/common/talker_first_audio.py b/vllm_omni/model_executor/models/common/talker_first_audio.py new file mode 100644 index 000000000..ad67aff50 --- /dev/null +++ b/vllm_omni/model_executor/models/common/talker_first_audio.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Shared eligibility; each model separately controls its first-audio opt-in.""" + +from typing import Any + +import torch + +from vllm_omni.platforms import current_omni_platform + + +def supports_talker_first_audio(vllm_config: Any) -> bool: + """Require CUDA MRv2 streaming, local single-rank execution and no prefix replay. + + Other configurations retain the regular codec path. These conditions do + not enable the feature: model-specific options must also accept it. + """ + model = vllm_config.model_config + if not (getattr(model, "use_v2_model_runner", False) and getattr(model, "async_chunk", False)): + return False + parallel = vllm_config.parallel_config + return ( + current_omni_platform.is_cuda() + and torch.device(vllm_config.device_config.device).type == "cuda" + and parallel.tensor_parallel_size == 1 + and parallel.pipeline_parallel_size == 1 + and parallel.distributed_executor_backend in (None, "uni") + and not vllm_config.cache_config.enable_prefix_caching + ) diff --git a/vllm_omni/model_executor/models/qwen3_omni/first_frame_decoder.py b/vllm_omni/model_executor/models/qwen3_omni/first_frame_decoder.py new file mode 100644 index 000000000..4720071fe --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_omni/first_frame_decoder.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""First codec frame -> PCM inside the Qwen3-Omni Talker process. + +Code2Wav decodes each streaming chunk statelessly from its codes and left +context. With a one-frame first chunk (``codec_chunk_ramp`` starting at 1), +chunk 0 is a context-free decode of the stream's first frame, which the Talker +can run as soon as eager MTP completes that frame, removing the connector hop +and the Code2Wav step from time to first audio. The orchestrator then trims +Code2Wav's copy of those samples. + +The decoder is a second ``Qwen3OmniMoeCode2Wav`` built and loaded exactly like +the Code2Wav stage's, and captured with the same one-frame input shape as that +stage's ``(batch=1, frames=1)`` graph. Optional cuDNN autotuning can select +different algorithms in the two processes and change floating-point rounding. +Each bucket gets its own graph pool because it replays on a side stream while +the Talker's own graphs may still run. +""" + +from __future__ import annotations + +import os +from collections.abc import Iterable, Sequence +from typing import Any + +import torch +import torch.nn as nn +from torch.cuda import CUDAGraph +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.common.qwen3_code_predictor import CodePredictorWrapper +from vllm_omni.model_executor.models.common.talker_first_audio import supports_talker_first_audio +from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp + +logger = init_logger(__name__) + + +def talker_first_audio_enabled(vllm_config: Any) -> bool: + """Opt in only when chunk 0 carries the first frame in the same step.""" + extra = CodePredictorWrapper._stage_connector_extra_config(vllm_config) + ramp = parse_chunk_ramp(extra) + initial_frames = ramp[0] if ramp is not None else int(extra.get("initial_codec_chunk_frames") or 0) + return ( + CodePredictorWrapper._parse_bool_config(extra.get("talker_first_audio")) + and initial_frames == 1 + and os.environ.get("VLLM_OMNI_TALKER_FIRST_AUDIO", "1") == "1" + and supports_talker_first_audio(vllm_config) + ) + + +class Qwen3OmniFirstFrameDecoder(nn.Module): + def __init__(self, code2wav: nn.Module, sample_rate: int) -> None: + super().__init__() + self.code2wav = code2wav + self.num_quantizers = int(code2wav.config.num_quantizers) + self.sample_rate = int(sample_rate) + self._graphs: dict[int, tuple[CUDAGraph, torch.Tensor, torch.Tensor]] = {} + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load the checkpoint's ``code2wav.*`` weights as the Code2Wav stage does.""" + loaded = self.code2wav.load_weights(weights) + if hasattr(self.code2wav, "precompute_snake_caches"): + self.code2wav.precompute_snake_caches() + return loaded + + @torch.inference_mode() + def capture(self, batch_sizes: Sequence[int] = (1, 2, 4, 8)) -> None: + device = next(self.code2wav.parameters()).device + if device.type != "cuda": + return + for batch_size in sorted(set(batch_sizes)): + static_input = torch.zeros(batch_size, self.num_quantizers, 1, dtype=torch.long, device=device) + self.code2wav(static_input) + torch.accelerator.synchronize(device) + graph = CUDAGraph() + with torch.cuda.graph(graph, pool=torch.cuda.graph_pool_handle()): + static_output = self.code2wav(static_input) + self._graphs[batch_size] = (graph, static_input, static_output) + logger.info("Captured Qwen3-Omni Talker first-frame decoder graphs for batch sizes %s", sorted(self._graphs)) + + @torch.inference_mode() + def decode(self, codes: torch.Tensor) -> torch.Tensor: + """``codes`` [n, num_quantizers] long on device -> float32 PCM [n, samples] (new tensor).""" + if not self._graphs: + return self.code2wav(codes.to(torch.long).unsqueeze(-1))[:, 0, :].float() + outputs = [] + largest = max(self._graphs) + for start in range(0, int(codes.shape[0]), largest): + chunk = codes[start : start + largest] + rows = int(chunk.shape[0]) + batch_size = min(size for size in self._graphs if size >= rows) + graph, static_input, static_output = self._graphs[batch_size] + static_input.zero_() + static_input[:rows, :, 0].copy_(chunk) + graph.replay() + # A replay overwrites graph-owned storage; .float() aliases FP32 outputs. + outputs.append(static_output[:rows, 0, :].to(dtype=torch.float32, copy=True)) + return torch.cat(outputs, dim=0) diff --git a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py index 9caff6d27..7a6c07326 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py @@ -45,6 +45,8 @@ QWEN3_OMNI_PIPELINE = PipelineConfig( engine_output_type="latent", custom_process_next_stage_input_func=(f"{_PROC}.thinker2talker_full_payload"), async_chunk_process_next_stage_input_func=(f"{_PROC}.thinker2talker_async_chunk"), + # Takes effect only when the deploy selects model_runner v2 for this stage. + supports_native_mrv2_data_plane=True, sampling_constraints={"detokenize": True}, ), StagePipelineConfig( @@ -57,6 +59,8 @@ QWEN3_OMNI_PIPELINE = PipelineConfig( sync_process_input_func=f"{_PROC}.thinker2talker_token_only", custom_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_full_payload"), async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"), + # Each stage can use MRv2 when its deploy selects model_runner v2. + supports_native_mrv2_data_plane=True, sampling_constraints={ "detokenize": False, "stop_token_ids": [2150], @@ -72,6 +76,7 @@ QWEN3_OMNI_PIPELINE = PipelineConfig( final_output_type="audio", hf_config_name="thinker_config", engine_output_type="audio", + supports_native_mrv2_data_plane=True, sampling_constraints={"detokenize": True}, requires_full_payload_input=True, ), diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index cfba92e98..50601d1fa 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -4,7 +4,7 @@ """Inference-only Qwen3-Omni-Moe unified model (thinker + talker + code2wav).""" import asyncio -from collections.abc import AsyncGenerator, Iterable +from collections.abc import AsyncGenerator, Iterable, Mapping from functools import cached_property from typing import Any @@ -50,6 +50,7 @@ from vllm_omni.data_entry_keys import Embeddings, HiddenStates, Ids, OmniPayload from vllm_omni.metrics import definitions as defs from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import talker_first_audio_enabled from vllm_omni.model_executor.models.qwen3_omni.quantization import ( apply_outer_quant_config_mapping, ) @@ -61,7 +62,9 @@ from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( Qwen3OmniMoeThinkerProcessingInfo, ) from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, safe_tensor_reshape +from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp, ramp_decode_windows from vllm_omni.platforms import current_omni_platform +from vllm_omni.utils.device_copy import to_device_nonblocking # Special token IDs for Qwen3 Omni MoE # Reference: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json @@ -88,6 +91,13 @@ TALKER_CODEC_THINK_EOS_ID = 4205 # Think mode end logger = init_logger(__name__) +def _cpu_long(values: Any) -> torch.Tensor: + """Token ids as a host tensor (payload ids arrive as lists or CPU tensors).""" + if isinstance(values, torch.Tensor): + return values.detach().to("cpu", torch.long) + return torch.as_tensor(values, dtype=torch.long) + + @MULTIMODAL_REGISTRY.register_processor( Qwen3OmniMoeThinkerMultiModalProcessor, info=Qwen3OmniMoeThinkerProcessingInfo, @@ -135,6 +145,11 @@ class Qwen3OmniMoeForConditionalGeneration( realtime_max_tokens = 64 + # Set by the MRv2 model state once eager Talker-MTP frames are enabled. + eager_frames_active = False + # Postprocess only stores the last hidden for the deferred decode MTP. + eager_frames_need_postprocess = False + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.have_multimodal_outputs = True @@ -180,9 +195,18 @@ class Qwen3OmniMoeForConditionalGeneration( # thinker layers and the forward must return what stock vLLM expects. self.is_staged_run = getattr(vllm_config.model_config, "model_stage", None) is not None self._returns_tuple = self.model_stage == "thinker" and self.is_staged_run + # The captured layers are a fixed-structure tensor dict on the token + # axis, so MRv2 may replay decode in FULL graphs and return them too. + self.supports_mrv2_full_graph_aux_outputs = self._returns_tuple if self.model_stage == "thinker": self.use_async_omni_output = True + # MRv2: the Talker's prompt ends with the first generated token's + # embedding; publishing it with the prefill output lets the first + # Thinker->Talker chunk leave one decode step earlier. + # Speculative sampling can change the token width between steps. + # Keep the existing capture-stream handoff for the entire request. + self.publishes_sampled_embeddings = self.is_staged_run and vllm_config.speculative_config is None # Initialize thinker model (multimodal processing + text generation) # Create a new vllm_config with thinker_config as the hf_config thinker_vllm_config = vllm_config.with_hf_config( @@ -233,6 +257,19 @@ class Qwen3OmniMoeForConditionalGeneration( ) self.model = self.talker self.code2wav = None + # Residual codebooks (MTP) output key, shared by the V1 runner + # (talker_mtp_*) and the MRv2 model state (mtp_*). + self.talker_mtp_output_key = ("codes", "audio") + # MRv2: the step that samples CB0 may complete the frame itself, + # since the residual codebooks depend only on CB0 and that step's + # last hidden, and the next input is their embedding sum plus the + # next text step (see ``talker_mtp``). + self.mtp_eager_frames = True + self._codec_codebook_size = int(talker_config.code_predictor_config.vocab_size) + self._apply_subtalker_sampling_params(getattr(vllm_config.model_config, "subtalker_sampling_params", None)) + + # Decode each stream's first frame here (see talker_first_audio_enabled). + self.first_frame_decoder = self._build_first_frame_decoder(vllm_config, code2wav_config, prefix) # for CI: Initialize special tokens embeddings early to avoid AttributeError when loading dummy weights self._init_special_tokens_embeddings() @@ -247,6 +284,9 @@ class Qwen3OmniMoeForConditionalGeneration( # talker MTP codec codes must stay on GPU to avoid a per-step D2H # sync stall; build_mm_cpu handles the eventual D2H at payload time. ("codes", "audio"), + # The growing thinker-embed cache is read and extended every + # decode step; on CPU it cost a blocking D2H per request per step. + ("embed", "cached_decode"), } elif self.model_stage == "code2wav": @@ -285,6 +325,74 @@ class Qwen3OmniMoeForConditionalGeneration( ["hidden_states", "residual", *incoming_captures], thinker_config.text_config.hidden_size ) + @staticmethod + def _build_first_frame_decoder( + vllm_config: VllmConfig, + code2wav_config: Qwen3OmniMoeCode2WavConfig, + prefix: str, + ) -> nn.Module | None: + """A Code2Wav copy for the Talker's first-frame audio, or None when disabled. + + Only a streaming (async-chunk) Talker hands frames to Code2Wav chunk by + chunk, so only there is a stream's first chunk a context-free decode. + The copy is built like the Code2Wav stage's module (same class, config, + dtype and device context). Optional cuDNN autotuning is also shared, + although per-process algorithm selection can change PCM rounding. + """ + if not talker_first_audio_enabled(vllm_config): + return None + from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import Qwen3OmniFirstFrameDecoder + + code2wav = init_vllm_registered_model( + vllm_config=vllm_config.with_hf_config(code2wav_config, architectures=["Qwen3OmniMoeCode2Wav"]), + prefix=maybe_prefix(prefix, "first_frame_decoder.code2wav"), + hf_config=code2wav_config, + architectures=["Qwen3OmniMoeCode2Wav"], + ) + logger.info("Qwen3-Omni Talker decodes first-frame audio in-stage") + return Qwen3OmniFirstFrameDecoder(code2wav, defs.resolve_audio_sample_rate(code2wav_config)) + + def capture_first_frame_graphs(self) -> None: + """Called by the MRv2 runner after its own graph capture.""" + decoder = getattr(self, "first_frame_decoder", None) + if decoder is not None: + decoder.capture() + + def _apply_subtalker_sampling_params(self, params: Any) -> None: + """Apply the deploy's ``subtalker_sampling_params`` to the residual-codebook predictor. + + The predictor samples in its "stored" mode (top-k 50, top-p 0.8 by + default) and ignores per-call sampling arguments; ``do_sample: false`` + makes it greedy. + """ + if not isinstance(params, Mapping) or not params: + return + predictor = self.talker.code_predictor + predictor.set_sampling_params( + top_k=int(params.get("top_k", predictor._top_k)), + top_p=float(params.get("top_p", predictor._top_p)), + do_sample=bool(params.get("do_sample", True)), + ) + + # MRv2 capability names read the V1 canonical values, so a platform patch + # of ``talker_mtp_graph_safe`` applies to both runners. + @property + def mtp(self): + """MRv2 name for ``talker_mtp``; only the Talker stage has an MTP.""" + return self.talker_mtp if self.model_stage == "talker" else None + + @property + def mtp_output_key(self) -> tuple[str, str]: + return getattr(self, "talker_mtp_output_key", ("codes", "audio")) + + @property + def mtp_graph_safe(self) -> bool: + return bool(getattr(self, "talker_mtp_graph_safe", False)) + + def mtp_frame_valid(self, layer0: torch.Tensor) -> torch.Tensor: + """Rows whose CB0 is a codec id rather than codec EOS (or another special id).""" + return (layer0 >= 0) & (layer0 < self._codec_codebook_size) + def _thinker_capture_layer_indices(self) -> list[int]: accept_layer = getattr(self.talker_config, "accept_hidden_layer", None) return [0, int(accept_layer)] if self.is_staged_run and accept_layer is not None else [] @@ -496,38 +604,47 @@ class Qwen3OmniMoeForConditionalGeneration( # ========== Stage 3: Code2Wav ========== elif self.model_stage == "code2wav": seq_token_counts: list[int] | None = kwargs.get("seq_token_counts") + num_quantizers = int(self.code2wav.config.num_quantizers) # Extract codec codes from input - if input_ids.shape[0] % 16 == 0: + if input_ids.shape[0] % num_quantizers == 0: if seq_token_counts is not None: - max_seq_len = max(seq_token_counts) // 16 + max_seq_len = max(seq_token_counts) // num_quantizers batch_size = len(seq_token_counts) split_codes = torch.split(input_ids, seq_token_counts, dim=0) - codes = torch.zeros((batch_size, 16, max_seq_len), device=input_ids.device, dtype=input_ids.dtype) + codes = torch.zeros( + (batch_size, num_quantizers, max_seq_len), device=input_ids.device, dtype=input_ids.dtype + ) for idx, code in enumerate(split_codes): - seq_len = code.shape[0] // 16 - codes[idx, :, :seq_len] = code.reshape(16, seq_len) + seq_len = code.shape[0] // num_quantizers + codes[idx, :, :seq_len] = code.reshape(num_quantizers, seq_len) else: - codes = input_ids.reshape(1, 16, -1) + codes = input_ids.reshape(1, num_quantizers, -1) else: if seq_token_counts is None: logger.debug( - "Code2Wav warmup input length %s is not divisible by 16; padding with zeros.", + "Code2Wav warmup input length %s is not divisible by %s; padding with zeros.", input_ids.shape[0], + num_quantizers, ) else: logger.warning_once( - "Code2Wav input length is not divisible by 16; padding with zeros. " - "This is expected only during cudagraph warmup." + "Code2Wav input length is not divisible by %s; padding with zeros. " + "This is expected only during cudagraph warmup.", + num_quantizers, ) input_ids_flatten = input_ids.reshape(-1) input_ids_flatten = torch.cat( [ input_ids_flatten, - torch.zeros(16 - input_ids.shape[0] % 16, dtype=torch.long, device=input_ids.device), + torch.zeros( + num_quantizers - input_ids.shape[0] % num_quantizers, + dtype=torch.long, + device=input_ids.device, + ), ] ) - codes = input_ids_flatten.reshape(1, 16, -1) + codes = input_ids_flatten.reshape(1, num_quantizers, -1) # Generate audio from codec codes # Get every request's left_context_size from runtime_additional_information (passed via kwargs) @@ -541,6 +658,38 @@ class Qwen3OmniMoeForConditionalGeneration( logger.debug("No additional_information provided to code2wav stage.") audio_tensors = self.generate_audio(codes, left_context_size, seq_token_counts) + # Keep the skip decision attached to this forward's output, not + # mutable model state: asynchronous materialization can overlap steps. + from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY + + flags = [ + bool(info.get("meta", {}).get("first_audio", False)) for info in (runtime_additional_information or []) + ] + if flags and any(flags): + if len(flags) != len(audio_tensors): + raise ValueError("First-audio flags must align with Code2Wav requests") + sample_rate = defs.resolve_audio_sample_rate(self.code2wav_config) + frame_counts = ( + [count // num_quantizers for count in seq_token_counts] if seq_token_counts else [codes.shape[-1]] + ) + # The causal decoder withholds a right-edge tail. The first + # frame has fewer samples than total_upsample; trim its actual + # prefix length, including when chunk 0 contains several frames. + prefix_lengths = [ + max(0, audio.numel() - (frames - 1) * int(self.code2wav.total_upsample)) if skip else 0 + for audio, frames, skip in zip(audio_tensors, frame_counts, flags, strict=True) + ] + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={ + "model_outputs": [ + audio.reshape(1, -1)[..., prefix:] + for audio, prefix in zip(audio_tensors, prefix_lengths, strict=True) + ], + "sr": [torch.tensor(sample_rate, dtype=torch.int32) for _ in audio_tensors], + FIRST_AUDIO_REQUIRED_KEY: [torch.tensor(skip) for skip in flags], + }, + ) return audio_tensors # Fallback (shouldn't reach here) @@ -605,7 +754,13 @@ class Qwen3OmniMoeForConditionalGeneration( if "runtime_additional_information" in kwargs and "model_intermediate_buffer" not in kwargs: logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer") + # MRv2 passes each request's token span. + spans = kwargs.get("request_token_spans") + if spans is not None and self.eager_frames_active and self.vllm_config.model_config.async_chunk: + return self._make_eager_talker_output(talker_hidden, spans) code_predictor_codes = [info.get("codes", {}).get("audio") for info in info_dicts] + if spans is not None: + code_predictor_codes = self._align_codes_to_spans(code_predictor_codes, spans, talker_hidden.device) audio_codes = torch.cat(code_predictor_codes, dim=0) multimodal_outputs: OmniPayload = {"codes": {"audio": audio_codes}} span_len = audio_codes.shape[0] @@ -627,6 +782,44 @@ class Qwen3OmniMoeForConditionalGeneration( return model_outputs + def _align_codes_to_spans( + self, + codes: list[Any], + spans: list[tuple[int, int]], + device: torch.device, + ) -> list[torch.Tensor]: + """One ``codes.audio`` block per request with exactly its span's rows. + + MRv2 also runs requests that never went through preprocess (vLLM's + ``_warmup_`` requests), whose buffers carry no codes; they get zero + rows, which the downstream processor skips like prefill rows. + """ + q = int(self.talker.num_code_groups) + aligned: list[torch.Tensor] = [] + for block, (start, end) in zip(codes, spans, strict=True): + if isinstance(block, torch.Tensor) and block.ndim == 2 and block.shape[0] == end - start: + aligned.append(block) + else: + aligned.append(torch.zeros((end - start, q), dtype=torch.long, device=device)) + return aligned + + def _make_eager_talker_output(self, hidden: torch.Tensor, spans: list[tuple[int, int]]) -> OmniOutput: + """Async-chunk output when the MRv2 runner completes frames after sampling. + + ``run_eager_mtp`` writes each sampled row's frame and validity into the + last token row of its request span; every other row carries no frame. + """ + num_tokens = spans[-1][1] if spans else 0 + q = int(self.talker.num_code_groups) + mm: OmniPayload = { + "codes": {"audio": torch.zeros((num_tokens, q), dtype=torch.long, device=hidden.device)}, + "meta": { + "codec_frame_valid": torch.zeros((num_tokens,), dtype=torch.int8, device=hidden.device), + "first_audio": torch.zeros((num_tokens,), dtype=torch.int8, device=hidden.device), + }, + } + return OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm) + # ==================== Audio Generation ==================== def generate_audio( @@ -849,7 +1042,7 @@ class Qwen3OmniMoeForConditionalGeneration( def _proj_from_thinker(x_opt: torch.Tensor | None) -> torch.Tensor: if isinstance(x_opt, torch.Tensor) and x_opt.numel() > 0: - xin = _ensure_1x1(x_opt).to(module_device) + xin = to_device_nonblocking(_ensure_1x1(x_opt), module_device) else: xin = torch.zeros( (1, thinker_embed.shape[-1]), @@ -882,27 +1075,26 @@ class Qwen3OmniMoeForConditionalGeneration( voice_type = str(voice_type).lower().strip() start_index = meta.get("num_processed_tokens", 0) end_index = start_index + input_embeds.shape[0] - # Read thinker outputs for prefill - thinker_sequence_embeds = embed["prefill"].to( - device=self._module_device(self.talker), dtype=torch.bfloat16 + # Read thinker outputs for prefill. Every copy here is H2D from the + # connector payload; pageable copies would stall the whole Talker batch. + talker_device = self._module_device(self.talker) + thinker_sequence_embeds = to_device_nonblocking(embed["prefill"], talker_device).to( + torch.bfloat16 ) # Tensor [P,H] - thinker_hidden_states = hs["output"].to( - device=self._module_device(self.talker), dtype=torch.bfloat16 - ) # Tensor [K,H] + thinker_hidden_states = to_device_nonblocking(hs["output"], talker_device).to(torch.bfloat16) # Tensor [K,H] + # Token ids stay on CPU too: prefill parsing reads them on the host. + thinker_sequences_cpu = None if ids.get("all") is None else _cpu_long(ids["all"]) + thinker_chatml_ids_cpu = None if ids.get("prompt") is None else _cpu_long(ids["prompt"]) thinker_sequences = ( - ids.get("all") - if ids.get("all") is None - else torch.as_tensor(ids["all"], device=self._module_device(self.talker)) + None if thinker_sequences_cpu is None else to_device_nonblocking(thinker_sequences_cpu, talker_device) ) thinker_chatml_ids = ( - ids.get("prompt") - if ids.get("prompt") is None - else torch.as_tensor(ids["prompt"], device=self._module_device(self.talker)) + None if thinker_chatml_ids_cpu is None else to_device_nonblocking(thinker_chatml_ids_cpu, talker_device) ) - tts_bos_thinker = embed["tts_bos"].to(device=self._module_device(self.talker), dtype=torch.bfloat16) - tts_eos_thinker = embed["tts_eos"].to(device=self._module_device(self.talker), dtype=torch.bfloat16) - tts_pad_thinker = embed["tts_pad"].to(device=self._module_device(self.talker), dtype=torch.bfloat16) + tts_bos_thinker = to_device_nonblocking(embed["tts_bos"], talker_device).to(torch.bfloat16) + tts_eos_thinker = to_device_nonblocking(embed["tts_eos"], talker_device).to(torch.bfloat16) + tts_pad_thinker = to_device_nonblocking(embed["tts_pad"], talker_device).to(torch.bfloat16) if thinker_sequence_embeds is None or thinker_hidden_states is None: raise ValueError( @@ -924,6 +1116,8 @@ class Qwen3OmniMoeForConditionalGeneration( ) if ids_chatml.ndim == 1: ids_chatml = ids_chatml.unsqueeze(0) + if thinker_chatml_ids_cpu is not None and thinker_chatml_ids_cpu.ndim == 1: + thinker_chatml_ids_cpu = thinker_chatml_ids_cpu.unsqueeze(0) else: # Fallback: create dummy ids if not provided ids_chatml = torch.zeros( @@ -932,6 +1126,7 @@ class Qwen3OmniMoeForConditionalGeneration( device=self._module_device(self.talker), ) thinker_sequences = ids_chatml + thinker_chatml_ids_cpu = thinker_sequences_cpu = None speaker_id = self._get_text_spk_token_id(voice_type) req_input_ids, req_embeds, trailing_text_hidden = self._thinker_to_talker_prefill( @@ -944,6 +1139,8 @@ class Qwen3OmniMoeForConditionalGeneration( tts_bos_thinker=tts_bos_thinker, tts_eos_thinker=tts_eos_thinker, tts_pad_thinker=tts_pad_thinker, + input_ids_cpu=thinker_chatml_ids_cpu, + thinker_result_ids_cpu=thinker_sequences_cpu, ) # Queue trailing_text_hidden for decode (drop first for next steps), @@ -993,12 +1190,11 @@ class Qwen3OmniMoeForConditionalGeneration( if cached_thinker_decode_embeds is None: update_dict.setdefault("embed", {})["cached_decode"] = thinker_decode_embeds else: - cached_thinker_decode_embeds = cached_thinker_decode_embeds.to( - device=self._module_device(self.talker), dtype=torch.bfloat16 - ) - thinker_decode_embeds = thinker_decode_embeds.to( - device=self._module_device(self.talker), dtype=torch.bfloat16 + talker_device = self._module_device(self.talker) + cached_thinker_decode_embeds = to_device_nonblocking(cached_thinker_decode_embeds, talker_device).to( + torch.bfloat16 ) + thinker_decode_embeds = to_device_nonblocking(thinker_decode_embeds, talker_device).to(torch.bfloat16) update_dict.setdefault("embed", {})["cached_decode"] = torch.cat( [cached_thinker_decode_embeds, thinker_decode_embeds], dim=0 ) @@ -1015,26 +1211,31 @@ class Qwen3OmniMoeForConditionalGeneration( tts_bos_thinker: torch.Tensor | None = None, tts_eos_thinker: torch.Tensor | None = None, tts_pad_thinker: torch.Tensor | None = None, + input_ids_cpu: torch.Tensor | None = None, + thinker_result_ids_cpu: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: """ Project thinker outputs to talker inputs during prefill stage. + ``input_ids_cpu`` / ``thinker_result_ids_cpu`` (host copies of the same + ids) move segment parsing off the GPU: each ``.item()`` / ``nonzero`` / + mask test on device ids is a host sync that stalls the Talker batch. + Returns: (input_ids, input_embeds) for talker """ target_len = thinker_result_ids.shape[-1] - im_start_indexes = torch.cat( - ( - torch.nonzero(input_ids[0] == self.config.im_start_token_id).squeeze(-1), - torch.tensor([target_len], device=input_ids.device, dtype=input_ids.dtype), - ), - dim=-1, - ) # Shape [n_starts + 1]; Take batch 0 since batched inference is not supported here. + # Take batch 0 since batched inference is not supported here. + parse_ids = input_ids_cpu if input_ids_cpu is not None else input_ids + result_ids = thinker_result_ids_cpu if thinker_result_ids_cpu is not None else thinker_result_ids + im_start_indexes = torch.nonzero(parse_ids[0] == self.config.im_start_token_id).flatten().tolist() + im_start_indexes.append(target_len) multimodal_mask = ( - (thinker_result_ids == self.thinker_config.audio_token_id) | - (thinker_result_ids == self.thinker_config.image_token_id) | - (thinker_result_ids == self.thinker_config.video_token_id) - ).to(input_ids.device) # [t] # fmt: skip + (result_ids == self.thinker_config.audio_token_id) | + (result_ids == self.thinker_config.image_token_id) | + (result_ids == self.thinker_config.video_token_id) + ) # [t] # fmt: skip + role_tokens = parse_ids[0].tolist() tts_bos_embed, tts_eos_embed, tts_pad_embed = self._get_tts_embed( thinker_embed, tts_bos_thinker, tts_eos_thinker, tts_pad_thinker @@ -1045,21 +1246,21 @@ class Qwen3OmniMoeForConditionalGeneration( trailing_text_hidden_all: torch.Tensor | None = None # For every chatml parts for i in range(len(im_start_indexes) - 1): - im_start_index = im_start_indexes[i].item() - segment_end_index = im_start_indexes[i + 1].item() - role_token = input_ids[0][im_start_index + 1] + im_start_index = im_start_indexes[i] + segment_end_index = im_start_indexes[i + 1] + role_token = role_tokens[im_start_index + 1] # Talker should ignore thinker system prompt - if (role_token == self.config.system_token_id).item(): + if role_token == self.config.system_token_id: continue # Talker takes word embeddings for tokens and hidden state from `accept_hidden_layer` for multimodal inputs - elif (role_token == self.config.user_token_id).item(): + elif role_token == self.config.user_token_id: talker_user_part = self._get_talker_user_parts( im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed ) talker_input_embeds.append(talker_user_part) talker_input_ids.append(thinker_result_ids[im_start_index:segment_end_index]) # Take assistant output (for now) - elif (role_token == self.config.assistant_token_id).item() and i == len(im_start_indexes) - 2: + elif role_token == self.config.assistant_token_id and i == len(im_start_indexes) - 2: talker_assistant_embeds, talker_assistant_ids, trailing_text_hidden = self._get_talker_assistant_parts( im_start_index, segment_end_index, @@ -1078,7 +1279,7 @@ class Qwen3OmniMoeForConditionalGeneration( except Exception: pass # History assistant output (ignore for now) - elif (role_token == self.config.assistant_token_id).item() and i != len(im_start_indexes) - 2: + elif role_token == self.config.assistant_token_id and i != len(im_start_indexes) - 2: continue else: raise AssertionError("Expect role id after <|im_start|> (assistant, user, system)") @@ -1103,20 +1304,22 @@ class Qwen3OmniMoeForConditionalGeneration( cached_thinker_decode_embeds = embed.get("cached_decode", None) thinker_decode_embed = embed.get("decode", None) + if getattr(self.vllm_config.model_config, "use_v2_model_runner", False) and not meta.get("resumable", False): + return self._next_thinker_decode_text_step( + cached_thinker_decode_embeds, thinker_decode_embed, meta, device, update_dict + ) start_index = meta.get("num_processed_tokens", 0) if cached_thinker_decode_embeds is not None and start_index < cached_thinker_decode_embeds.shape[0]: - cached_thinker_decode_embeds = cached_thinker_decode_embeds.to(device) + cached_thinker_decode_embeds = to_device_nonblocking(cached_thinker_decode_embeds, device) thinker_embed = cached_thinker_decode_embeds[start_index] if thinker_decode_embed is not None: - thinker_decode_embed = thinker_decode_embed.to(device) + thinker_decode_embed = to_device_nonblocking(thinker_decode_embed, device) cached_thinker_decode_embeds = torch.cat([cached_thinker_decode_embeds, thinker_decode_embed], dim=0) update_dict.setdefault("embed", {})["cached_decode"] = cached_thinker_decode_embeds elif thinker_decode_embed is not None: - thinker_embed = thinker_decode_embed - if thinker_embed.device != device: - thinker_embed = thinker_embed.to(device) + thinker_embed = to_device_nonblocking(thinker_decode_embed, device) else: # When the tokens output by the thinker are exhausted, an EOS token needs to be appended. @@ -1129,6 +1332,41 @@ class Qwen3OmniMoeForConditionalGeneration( update_dict.setdefault("embed", {})["decode"] = None return self.talker.text_projection(thinker_embed).to(device) + def _next_thinker_decode_text_step( + self, + pending: torch.Tensor | None, + incoming: torch.Tensor | None, + meta: OmniPayloadMeta, + device: torch.device, + update_dict: OmniPayload, + ) -> torch.Tensor: + """Take the oldest unconsumed Thinker decode row as this step's text. + + Rows arrive in order but not one per Talker step: the MRv2 receiver + hands over every row that arrived since the previous step (the V1 + adapter delivers one chunk per step). New rows are appended to the + pending queue ``embed.cached_decode`` and each step consumes one. Once + the queue is empty the Thinker's text is exhausted: EOS once, then pad. + """ + embed_update = update_dict.setdefault("embed", {}) + embed_update["decode"] = None + if isinstance(incoming, torch.Tensor) and incoming.numel() > 0: + incoming = to_device_nonblocking(incoming, device).reshape(-1, incoming.shape[-1]) + if isinstance(pending, torch.Tensor) and pending.shape[0] > 0: + pending = to_device_nonblocking(pending, device) + pending = torch.cat([pending, incoming.to(pending.dtype)], dim=0) + else: + pending = incoming + if not isinstance(pending, torch.Tensor) or pending.shape[0] == 0: + embed_update["cached_decode"] = None + if meta.get("eos_emitted", False): + return self.tts_pad_embed.to(device) + update_dict.setdefault("meta", {})["eos_emitted"] = True + return self.tts_eos_embed.to(device) + pending = to_device_nonblocking(pending, device) + embed_update["cached_decode"] = pending[1:] if pending.shape[0] > 1 else None + return self.talker.text_projection(pending[:1]).to(device) + def talker_preprocess_decode( self, input_ids: torch.Tensor, input_embeds: torch.Tensor, update_dict: OmniPayload, payload: OmniPayload ): @@ -1203,6 +1441,21 @@ class Qwen3OmniMoeForConditionalGeneration( ) user_mm_mask = multimodal_mask[im_start_index:segment_end_index] + if user_mm_mask.device.type == "cpu": + # Host mask: no device sync for the test, and a text-only segment + # (the common TTS case) needs no masked gather at all. + if not bool(user_mm_mask.any()): + user_thinker_embed = thinker_embed[im_start_index:segment_end_index] + return self.talker.text_projection(user_thinker_embed).to(thinker_hidden.device, torch.bfloat16) + mm_index = to_device_nonblocking(torch.nonzero(user_mm_mask).flatten(), thinker_hidden.device) + text_index = to_device_nonblocking(torch.nonzero(~user_mm_mask).flatten(), thinker_hidden.device) + segment_hidden = thinker_hidden[im_start_index:segment_end_index] + segment_embed = thinker_embed[im_start_index:segment_end_index] + mm_hidden = self.talker.hidden_projection(segment_hidden.index_select(0, mm_index)) + user_talker_part.index_copy_(0, mm_index, mm_hidden.to(user_talker_part.dtype)) + text_hidden = self.talker.text_projection(segment_embed.index_select(0, text_index)) + user_talker_part.index_copy_(0, text_index, text_hidden.to(user_talker_part.dtype)) + return user_talker_part # Multimodal data exists if user_mm_mask.any(): user_thinker_hidden_mm = thinker_hidden[im_start_index:segment_end_index][user_mm_mask] @@ -1236,17 +1489,19 @@ class Qwen3OmniMoeForConditionalGeneration( ), dim=0, ) - codec_special_tokens = torch.tensor( - [ - self.config.talker_config.codec_nothink_id, - self.config.talker_config.codec_think_bos_id, - self.config.talker_config.codec_think_eos_id, - speaker_id, - self.config.talker_config.codec_pad_id, - self.config.talker_config.codec_bos_id, - ], - device=tts_pad_embed.device, - dtype=torch.long, + codec_special_tokens = to_device_nonblocking( + torch.tensor( + [ + self.config.talker_config.codec_nothink_id, + self.config.talker_config.codec_think_bos_id, + self.config.talker_config.codec_think_eos_id, + speaker_id, + self.config.talker_config.codec_pad_id, + self.config.talker_config.codec_bos_id, + ], + dtype=torch.long, + ), + tts_pad_embed.device, ) embed_input_ids = self.talker.embed_input_ids(codec_special_tokens).to( device=tts_pad_embed.device, dtype=torch.bfloat16 @@ -1381,18 +1636,31 @@ class Qwen3OmniMoeForConditionalGeneration( # ==================== Weight Loading ==================== - def _get_codec_frame_config(self) -> tuple[int, int]: - """Extract codec_chunk_frames and codec_left_context_frames from stage connector config.""" + def _get_codec_connector_extra(self) -> dict[str, Any]: model_cfg = getattr(self.vllm_config, "model_config", None) connector_cfg = getattr(model_cfg, "stage_connector_config", None) if isinstance(connector_cfg, dict): extra = connector_cfg.get("extra", {}) else: extra = getattr(connector_cfg, "extra", None) or {} + return extra if isinstance(extra, Mapping) else {} + + def _get_codec_frame_config(self) -> tuple[int, int]: + """Extract codec_chunk_frames and codec_left_context_frames from stage connector config.""" + extra = self._get_codec_connector_extra() chunk_frames = int(extra.get("codec_chunk_frames", 0) or 0) left_frames = int(extra.get("codec_left_context_frames", 0) or 0) return chunk_frames, left_frames + def _get_codec_ramp_windows(self) -> list[int]: + """Decode windows of the connector's ``codec_chunk_ramp`` (empty without a ramp).""" + extra = self._get_codec_connector_extra() + chunk_frames = int(extra.get("codec_chunk_frames", 25) or 25) + ramp = parse_chunk_ramp(dict(extra), steady=chunk_frames) + if not ramp: + return [] + return ramp_decode_windows(ramp, int(extra.get("codec_left_context_frames", 25) or 0)) + def _maybe_enable_code2wav_cudagraph(self) -> None: """Enable the inner Code2Wav CUDA graph unless this stage runs in eager mode.""" if not self.code2wav or not hasattr(self.code2wav, "enable_cudagraph"): @@ -1404,9 +1672,15 @@ class Qwen3OmniMoeForConditionalGeneration( return chunk_frames, left_frames = self._get_codec_frame_config() + streaming = bool(self.vllm_config.model_config.async_chunk) + # Opt-in multi-row graphs for streaming windows (connector extra + # ``codec_graph_batch_sizes``, e.g. [2, 3, 4, 6, 8, 12, 16]). + batch_sizes = self._get_codec_connector_extra().get("codec_graph_batch_sizes") or () self.code2wav.enable_cudagraph( codec_chunk_frames=chunk_frames, codec_left_context_frames=left_frames, + extra_capture_sizes=self._get_codec_ramp_windows() if streaming else (), + streaming_batch_sizes=[int(b) for b in batch_sizes] if streaming else (), ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -1439,6 +1713,12 @@ class Qwen3OmniMoeForConditionalGeneration( loaded_weights.update(talker_loaded) loaded_weights.update(self._init_special_tokens_embeddings()) + # The Talker's first-frame decoder is a Code2Wav copy (same code2wav.* weights). + first_frame_decoder = getattr(self, "first_frame_decoder", None) + if first_frame_decoder is not None and code2wav_weights: + decoder_loaded = first_frame_decoder.load_weights(code2wav_weights) + loaded_weights.update(add_prefix_to_loaded_weights(decoder_loaded, "first_frame_decoder.code2wav")) + # Load code2wav weights if self.code2wav and code2wav_weights: code2wav_loaded = self.code2wav.load_weights(code2wav_weights) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py index bc8f6b3b4..9ae711870 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py @@ -5,7 +5,8 @@ from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Callable, Iterable, Sequence +from contextlib import nullcontext import numpy as np import torch @@ -19,6 +20,7 @@ from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( Qwen3OmniMoeCode2WavDecoderBlock, Qwen3OmniMoeCode2WavTransformerModel, Qwen3OmniMoeConvNeXtBlock, + Qwen3OmniMoeSnakeBeta, ) from vllm.config import VllmConfig # type: ignore from vllm.logger import init_logger # type: ignore @@ -31,9 +33,93 @@ from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta from vllm_omni.model_executor.models.qwen3_omni.quantization import ( Qwen3OmniNestedSupportsQuant, ) +from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) +# Fixed cost of one streaming decode call (graph replay or eager launch), in +# decoded-frame equivalents: a call on a single frame costs about as much GPU +# time as decoding this many extra frames inside a larger call. +_DECODE_CALL_COST_FRAMES = 32 + + +def use_fused_snake(module: nn.Module) -> int: + """Replace HF ``Qwen3OmniMoeSnakeBeta`` activations under ``module`` with the fused one. + + The HF decoder blocks recompute ``exp(alpha)`` / ``exp(beta)`` and launch + about seven broadcast kernels per call; the shared ``SnakeBeta`` computes the + same ``x + 1/b * sin^2(a*x)`` in one kernel from precomputed caches. + Parameter names and shapes match, so this runs before weight loading. + """ + count = 0 + for parent in list(module.modules()): + for name, child in list(parent.named_children()): + if isinstance(child, Qwen3OmniMoeSnakeBeta): + setattr(parent, name, SnakeBeta(child.in_features)) + count += 1 + return count + + +def plan_decode_groups( + lengths: Sequence[int], + bucket_of: Callable[[int], int | None], + batch_sizes_for: Callable[[int], Sequence[int]], + call_cost_frames: int = _DECODE_CALL_COST_FRAMES, +) -> list[tuple[list[int], int]]: + """Split a streaming batch into decode calls of rows with similar lengths. + + Decoding every row at the batch's longest window wastes most of the work + when short ramp chunks (1-15 frames) share a step with steady 50-frame + windows. Rows are sorted by length and cut into contiguous groups of + graph buckets (``bucket_of``) so that the modelled cost, per call + ``call_cost_frames`` plus padded rows (to the next captured row count of + that bucket, ``batch_sizes_for``) times the bucket length, is smallest. + Returns ``(row indices, decode length)`` per call; the decode length is + the group's longest row. + """ + count = len(lengths) + if count == 0: + return [] + order = sorted(range(count), key=lambda row: lengths[row]) + buckets = [bucket_of(int(lengths[row])) or int(lengths[row]) for row in order] + distinct = sorted(set(buckets)) + members = [[row for row, bucket in zip(order, buckets) if bucket == size] for size in distinct] + + def sizes_for(size: int) -> list[int]: + return sorted({int(b) for b in batch_sizes_for(size) if int(b) > 0} | {1}) + + def cost(rows: int, size: int) -> int: + batch_sizes = sizes_for(size) + calls, rest = divmod(rows, batch_sizes[-1]) + padded = calls * batch_sizes[-1] + if rest: + calls += 1 + padded += next(b for b in batch_sizes if b >= rest) + return calls * call_cost_frames + padded * size + + best = [0] + [None] * len(distinct) + cut = [0] * (len(distinct) + 1) + for end in range(1, len(distinct) + 1): + rows = 0 + for start in range(end - 1, -1, -1): + rows += len(members[start]) + total = best[start] + cost(rows, distinct[end - 1]) + if best[end] is None or total < best[end]: + best[end], cut[end] = total, start + groups: list[tuple[list[int], int]] = [] + end = len(distinct) + while end > 0: + start = cut[end] + groups.append(([row for bucket_rows in members[start:end] for row in bucket_rows], distinct[end - 1])) + end = start + calls: list[tuple[list[int], int]] = [] + for rows, size in reversed(groups): + max_batch = sizes_for(size)[-1] + for offset in range(0, len(rows), max_batch): + part = rows[offset : offset + max_batch] + calls.append((part, max(int(lengths[row]) for row in part))) + return calls + class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): """ @@ -74,6 +160,14 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): super().__init__() self.config: Qwen3OmniMoeCode2WavConfig = vllm_config.model_config.hf_config + connector = getattr(vllm_config.model_config, "stage_connector_config", None) + extra = connector.get("extra", {}) if isinstance(connector, dict) else getattr(connector, "extra", {}) + # Both the separate stage and the Talker's first-frame copy receive + # this connector config. Keep algorithm selection consistent between + # them, and leave other models' cuDNN settings unchanged. + self._cudnn_benchmark = (extra or {}).get("codec_cudnn_benchmark", False) + if not isinstance(self._cudnn_benchmark, bool): + raise ValueError("codec_cudnn_benchmark must be a boolean") # Calculate total upsampling factor self.total_upsample = np.prod(self.config.upsample_rates + self.config.upsampling_ratios) @@ -122,10 +216,19 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): Qwen3OmniMoeCausalConvNet(output_dim, 1, kernel_size=7), ] self.decoder = nn.ModuleList(decoder) + # Keep existing V1 and non-CUDA decoder blocks unchanged. + if ( + (extra or {}).get("codec_fused_snake", False) + and current_omni_platform.is_cuda() + and torch.device(vllm_config.device_config.device).type == "cuda" + ): + use_fused_snake(self.decoder) # CUDA Graph support — reuses CUDAGraphDecoderWrapper from Qwen3-TTS self._cudagraph_enabled = False self._cudagraph_wrapper = None + # Captured row counts per graph size; empty keeps one padded decode per batch. + self._streaming_batch_sizes: dict[int, list[int]] = {} def precompute_snake_caches(self): """Precompute exp(alpha) and 1/(exp(beta)+eps) for all SnakeBeta modules.""" @@ -142,8 +245,19 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): device: torch.device | None = None, codec_chunk_frames: int = 0, codec_left_context_frames: int = 0, + extra_capture_sizes: Iterable[int] = (), + streaming_batch_sizes: Iterable[int] = (), ): - """Enable CUDA graph acceleration (same pattern as Qwen3-TTS Code2Wav).""" + """Enable CUDA graph acceleration (same pattern as Qwen3-TTS Code2Wav). + + ``extra_capture_sizes`` adds exact-size graphs (e.g. a chunk ramp's + decode windows) to the default buckets, so those chunks replay without + zero padding. ``streaming_batch_sizes`` (e.g. ``[2, 4, 8, 16]``) also + captures multi-row graphs for every bucket up to one streaming window + (chunk + left context); a streaming batch is then decoded as a few + length groups (``plan_decode_groups``) instead of one eager call + padded to its longest row. + """ from vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper import ( CUDAGraphDecoderWrapper, ) @@ -154,8 +268,30 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): logger.warning("Cannot enable CUDA Graph: not on CUDA device (got %s)", device) return + extra = {int(size) for size in extra_capture_sizes if int(size) > 0} + batch_sizes = sorted({int(b) for b in streaming_batch_sizes if int(b) > 1}) + capture_sizes = None + extra_shapes: list[tuple[int, int]] = [] + if extra or batch_sizes: + capture_sizes = sorted( + set( + CUDAGraphDecoderWrapper.compute_capture_sizes( + codec_chunk_frames=codec_chunk_frames, + codec_left_context_frames=codec_left_context_frames, + ) + ) + | extra + ) + window = codec_chunk_frames + codec_left_context_frames + streaming_sizes = [size for size in capture_sizes if window <= 0 or size <= window] + # A multi-row graph never holds more frames than the largest + # single-row one, so the shared graph pool does not grow. + largest = max(capture_sizes) + extra_shapes = [(b, size) for b in batch_sizes for size in streaming_sizes if b * size <= largest] wrapper = CUDAGraphDecoderWrapper( decoder=self, + capture_sizes=capture_sizes, + extra_capture_shapes=extra_shapes, num_quantizers=self.config.num_quantizers, enabled=True, ) @@ -172,10 +308,18 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): raise self._cudagraph_wrapper = wrapper self._cudagraph_enabled = True + if extra_shapes and torch.cuda.is_available(): + # Return the eager warm-up activations: this stage shares the GPU. + torch.accelerator.empty_cache() + self._streaming_batch_sizes = { + size: [1, *sorted(b for b, shape_size in extra_shapes if shape_size == size)] + for size in {shape_size for _b, shape_size in extra_shapes} + } logger.info( - "CUDA Graph enabled for Code2Wav: num_quantizers=%d, sizes=%s", + "CUDA Graph enabled for Code2Wav: num_quantizers=%d, sizes=%s, multi-row shapes=%s", self.config.num_quantizers, self._cudagraph_wrapper.capture_sizes, + extra_shapes, ) def forward(self, codes: torch.Tensor) -> torch.Tensor: @@ -200,6 +344,30 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): hidden = self.pre_transformer(inputs_embeds=hidden).last_hidden_state # Shape: [batch, seq_len, hidden_size] + # Algorithm search happens during eager warmup before graph capture. + # Restore process settings afterward, including on a failed decode. + # Captured replays use the selected kernels without this Python path. + cudnn = torch.backends.cudnn + context = ( + cudnn.flags( + enabled=cudnn.enabled, + benchmark=True, + benchmark_limit=10, + deterministic=cudnn.deterministic, + allow_tf32=cudnn.allow_tf32, + ) + if ( + self._cudnn_benchmark + and hidden.is_cuda + and torch.version.hip is None + and not torch.cuda.is_current_stream_capturing() + ) + else nullcontext() + ) + with context: + return self._decode_waveform(hidden) + + def _decode_waveform(self, hidden: torch.Tensor) -> torch.Tensor: # Stage 3: Upsampling hidden = hidden.permute(0, 2, 1) # [batch, hidden_size, seq_len] for blocks in self.upsample: @@ -305,28 +473,38 @@ class Qwen3OmniMoeCode2Wav(nn.Module, Qwen3OmniNestedSupportsQuant): "defaulting to left_context_size=zeros(len(codes)). This is expected during cudagraph warmup." ) left_context_size = [0] * codes.shape[0] - # Decode chunk - wavs = [] - if self._cudagraph_enabled and self._cudagraph_wrapper is not None: - batch_wav = self._cudagraph_wrapper.decode(codes) - else: - batch_wav = self(codes) if seq_token_counts is not None: code_seq_lens = [n // self.config.num_quantizers for n in seq_token_counts] else: # Fallback: assume all batch elements share the same sequence length. code_seq_lens = [codes.shape[-1]] * codes.shape[0] - # The decoder trims the same right-edge tail from every batch row. - # Infer it from the decoded window, not a shorter request's length: - # padding that request with codec zeros does not provide valid context. - tail = max(0, int(codes.shape[-1] * self.total_upsample) - batch_wav.shape[-1]) - for idx, code_seq_len in enumerate(code_seq_lens): - # Refill the previous chunk's tail using the re-decoded left context, - # and exclude the current row's padded tail just as singleton decode does. - start = max(0, left_context_size[idx] * self.total_upsample - tail) - end = max(0, code_seq_len * self.total_upsample - tail) - wav_chunk = batch_wav[idx, :, start:end] - wavs.append(wav_chunk) + batch, width = int(codes.shape[0]), int(codes.shape[-1]) + graphs = self._cudagraph_wrapper if self._cudagraph_enabled else None + batch_sizes = getattr(self, "_streaming_batch_sizes", None) + if graphs is not None and batch_sizes and batch > 1: + groups = plan_decode_groups( + code_seq_lens, graphs._get_padded_size, lambda size: batch_sizes.get(size, (1,)) + ) + else: + groups = [(list(range(batch)), width)] + wavs: list[torch.Tensor | None] = [None] * batch + for rows, length in groups: + if len(rows) == batch and length == width: + group_codes = codes + else: + # Basic indexing only: a device index tensor would need a host copy. + group_codes = torch.stack([codes[row, :, :length] for row in rows]) + batch_wav = graphs.decode(group_codes) if graphs is not None else self(group_codes) + # The decoder trims the same right-edge tail from every batch row. + # Infer it from the decoded window, not a shorter request's length: + # padding that request with codec zeros does not provide valid context. + tail = max(0, int(length * self.total_upsample) - batch_wav.shape[-1]) + for position, row in enumerate(rows): + # Refill the previous chunk's tail using the re-decoded left context, + # and exclude the current row's padded tail just as singleton decode does. + start = max(0, left_context_size[row] * self.total_upsample - tail) + end = max(0, code_seq_lens[row] * self.total_upsample - tail) + wavs[row] = batch_wav[position, :, start:end] return wavs def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm_omni/model_executor/models/qwen3_tts/cached_code_predictor.py b/vllm_omni/model_executor/models/qwen3_tts/cached_code_predictor.py index 81eb8fdda..4481eb50d 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/cached_code_predictor.py +++ b/vllm_omni/model_executor/models/qwen3_tts/cached_code_predictor.py @@ -18,8 +18,7 @@ from vllm_omni.model_executor.models.common.qwen3_code_predictor import ( CodePredictorDecoderLayer, _rotate_half, ) - -from .short_kv_attention import short_kv_attention +from vllm_omni.model_executor.models.common.short_kv_attention import short_kv_attention def _attention_cached( diff --git a/vllm_omni/model_executor/models/qwen3_tts/first_audio.py b/vllm_omni/model_executor/models/qwen3_tts/first_audio.py index 1396c0dc4..4b10b7507 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/first_audio.py +++ b/vllm_omni/model_executor/models/qwen3_tts/first_audio.py @@ -4,9 +4,7 @@ from typing import Any -import torch - -from vllm_omni.platforms import current_omni_platform +from vllm_omni.model_executor.models.common.talker_first_audio import supports_talker_first_audio def talker_first_audio_enabled(vllm_config: Any) -> bool: @@ -18,16 +16,4 @@ def talker_first_audio_enabled(vllm_config: Any) -> bool: from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor extra = Predictor._stage_connector_extra_config(vllm_config) - model = vllm_config.model_config - parallel = vllm_config.parallel_config - return ( - Predictor._parse_bool_config(extra.get("talker_first_audio")) - and current_omni_platform.is_cuda() - and torch.device(vllm_config.device_config.device).type == "cuda" - and bool(getattr(model, "use_v2_model_runner", False)) - and bool(getattr(model, "async_chunk", False)) - and parallel.tensor_parallel_size == 1 - and parallel.pipeline_parallel_size == 1 - and parallel.distributed_executor_backend in (None, "uni") - and not vllm_config.cache_config.enable_prefix_caching - ) + return Predictor._parse_bool_config(extra.get("talker_first_audio")) and supports_talker_first_audio(vllm_config) diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py index 675c5eecd..9620d0d86 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py @@ -55,6 +55,9 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapp talker_hidden_size=int(talker_config.hidden_size), prefix=prefix, ) + # TTS owns its frame-local cache and overrides _predict_step_logits. + # Do not also activate the shared wrapper's Omni incremental loop. + self._kv_cache_enabled = False # Store talker_config for backward compat (accessed by some callers) self.talker_config = talker_config self._vllm_config = vllm_config diff --git a/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py b/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py index 9b9462ad9..774e5940e 100644 --- a/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.py +++ b/vllm_omni/model_executor/stage_input_processors/chunk_size_utils.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 logging from dataclasses import dataclass, field @@ -99,6 +99,21 @@ def ramp_cumulative(index: int, ramp: list[int], steady: int) -> int: return sum(ramp) + (index + 1 - ramp_len) * steady +def ramp_decode_windows(ramp: list[int], left_context: int) -> list[int]: + """Decode window (left context + new frames) of each ramp chunk for a stateless decoder. + + Chunk ``i`` re-decodes up to ``left_context`` earlier frames with its + ``ramp[i]`` new ones, e.g. ``[1, 2, 4, 8, 16, 25]`` with 25 frames of + context -> ``[1, 3, 7, 15, 31, 50]``. + """ + windows: list[int] = [] + emitted = 0 + for new_frames in ramp: + windows.append(new_frames + min(emitted, max(0, left_context))) + emitted += new_frames + return windows + + def compute_ramp_emit( length: int, chunk_index: int, diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index 3cef0eed8..3fee73ac3 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -1,12 +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 # Copyright 2025 The Qwen team. """Stage input processor for Qwen3 Omni MoE: Thinker → Talker transition.""" import logging from collections.abc import Mapping from dataclasses import dataclass, field -from typing import Any +from typing import Any, cast import torch from vllm.inputs import TextPrompt @@ -23,6 +23,7 @@ from vllm_omni.data_entry_keys import ( ) from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import compute_ramp_emit, parse_chunk_ramp from vllm_omni.model_executor.stage_input_processors.tts_utils import ( extract_language_from_prompt, extract_language_from_request, @@ -76,6 +77,43 @@ def _get_accept_hidden_layer_index(transfer_manager: Any) -> int: return int(mc.hf_config.talker_config.accept_hidden_layer) +# Set once chunk 0 carried the first generated token taken from the prefill +# sample: from then on each chunk carries the token its step sampled. +_SAMPLED_TEXT_STREAM = "_qwen3_omni_sampled_text_stream" + + +def _is_final_token(request: Any) -> bool | None: + """Whether the token sampled last ends the Thinker's turn (``None``: unknown).""" + output_token_ids = _ensure_list(getattr(request, "output_token_ids", None) or []) + count = getattr(request, "output_token_count", None) + count = len(output_token_ids) if count is None else count + token = getattr(request, "last_output_token_id", None) + if token is None and output_token_ids: + token = output_token_ids[-1] + if token is None or count == 0: + return None + params = getattr(request, "sampling_params", None) + stop_ids = getattr(params, "all_stop_token_ids", None) or getattr(params, "stop_token_ids", None) or () + max_tokens = getattr(params, "max_tokens", None) + return token in stop_ids or (max_tokens is not None and count >= max_tokens) + + +def _sampled_token_embed(thinker_embed: Mapping[str, Any], request: Any) -> torch.Tensor | None: + """Embedding of the token this step sampled, when published and not final. + + The Talker is fed the embeddings of the Thinker's generated tokens except + the final one (a stop token, or the last one ``max_tokens`` allows), which + no forward ever processes. For any other token the published embedding is + exactly what the next step's layer-0 capture holds, one step later. + """ + sampled = thinker_embed.get("sampled") + if not isinstance(sampled, torch.Tensor) or sampled.dim() != 2 or sampled.shape[0] != 1: + return None + if getattr(request, "resumable", False) or _is_final_token(request) is not False: + return None + return sampled.detach().cpu() + + def _compute_talker_prompt_ids_length(info: OmniPayload, device: torch.device | str = "cuda") -> int: im_start_token_id = 151644 system_token_id = 8948 @@ -298,6 +336,7 @@ def _get_qwen3_streaming_state( streaming_context: Any | None, ) -> _Qwen3OmniStreamingState: bridge_states = getattr(streaming_context, "bridge_states", None) + assert bridge_states is not None per_model_state = bridge_states.setdefault("qwen3_omni", {}) state = per_model_state.get(request_id) if state is None: @@ -387,7 +426,7 @@ def thinker2talker_async_chunk( return None thinker_hs = multimodal_output.get("hidden_states", {}) - thinker_layers = thinker_hs.get("layers", {}) if isinstance(thinker_hs, dict) else {} + thinker_layers = cast(Mapping[int | str, Any], thinker_hs.get("layers", {}) if isinstance(thinker_hs, dict) else {}) thinker_embed_raw = multimodal_output.get("embed", {}) thinker_embed = thinker_embed_raw if isinstance(thinker_embed_raw, dict) else {} @@ -441,27 +480,48 @@ def thinker2talker_async_chunk( speaker=speaker, language=language, ) - if transfer_manager.request_payload.get(request_id) is None: - if not is_finished: - transfer_manager.request_payload[request_id] = to_dict(payload) - return None - else: - save_payload = transfer_manager.request_payload.pop(request_id) + assert payload.embed is not None and payload.embed.prefill is not None + assert payload.hidden_states is not None and payload.hidden_states.output is not None + save_payload = transfer_manager.request_payload.pop(request_id, None) + if save_payload is not None: payload.embed.prefill = torch.cat( (save_payload.get("embed", {}).get("prefill"), payload.embed.prefill), dim=0 ) payload.hidden_states.output = torch.cat( (save_payload.get("hidden_states", {}).get("output"), payload.hidden_states.output), dim=0 ) - prefill_shape = payload.embed.prefill.shape[0] - if not is_finished and prefill_shape <= len(prompt_token_ids): - transfer_manager.request_payload[request_id] = to_dict(payload) - return None + if not is_finished and payload.embed.prefill.shape[0] == len(prompt_token_ids): + # The Talker prompt ends with the first generated token's embedding, + # which the next decode step would capture; take it from the + # prefill step's sample instead when the producer published it. + first_token_embed = _sampled_token_embed(thinker_embed, request) + if first_token_embed is not None: + payload.embed.prefill = torch.cat((payload.embed.prefill, first_token_embed), dim=0) + transfer_manager.request_payload[request_id] = {_SAMPLED_TEXT_STREAM: True} + if not is_finished and payload.embed.prefill.shape[0] <= len(prompt_token_ids): + transfer_manager.request_payload[request_id] = to_dict(payload) + return None else: if request.resumable: return _construct_thinker2talker_streaming_input_async_chunk( is_finished, request, thinker_emb, thinker_hid, transfer_manager ) + stream = transfer_manager.request_payload.get(request_id) + if stream is not None and stream.get(_SAMPLED_TEXT_STREAM): + # This step's capture was already sent as the previous step's sample. + sampled = _sampled_token_embed(thinker_embed, request) + if sampled is None: + # A partial (including resumed) prefill has no accepted sample. + # Its captures replay text already sent before preemption. + raw_sample = thinker_embed.get("sampled") + if isinstance(raw_sample, torch.Tensor) and raw_sample.numel() == 0: + return None + if _is_final_token(request) is not True and not is_finished: + raise RuntimeError(f"Thinker sample embedding missing for request {request_id}") + # The final token is never part of the text; a finish goes out + # as the bare finish marker. + return None + thinker_emb = sampled if thinker_emb.shape[0] > 1: logger.warning( "Unexpected multiple embeddings in thinker2talker_async_chunk for chunk_id %d: " @@ -486,7 +546,7 @@ def thinker2talker_full_payload( transfer_manager: Any, pooling_output: dict[str, Any], request: OmniEngineCoreRequest, -) -> dict[str, Any] | None: +) -> OmniPayload | None: """Pack complete thinker output for the non-async connector path.""" rid = getattr(request, "request_id", None) if not isinstance(pooling_output, Mapping): @@ -612,7 +672,7 @@ def thinker2talker_token_only( ) thinker_sequences = prompt_token_ids + output_ids thinker_input_ids = prompt_token_ids - info_for_len = {"ids": {"all": thinker_sequences, "prompt": thinker_input_ids}} + info_for_len: OmniPayload = {"ids": {"all": thinker_sequences, "prompt": thinker_input_ids}} prompt_len = _compute_talker_prompt_ids_length(info_for_len, device="cpu") # Keep this fallback until the connector reliably preserves voice metadata. additional_information = to_dict( @@ -637,6 +697,14 @@ def thinker2talker_token_only( # ========================= +def _last_codec_row(codes: Any) -> torch.Tensor | None: + """The last ``[1, Q]`` row of token-major codec codes, or None if there is none.""" + if not isinstance(codes, torch.Tensor) or codes.ndim != 2 or codes.shape[0] == 0: + return None + # Own the row: it outlives this step's (batch-wide) output tensor. + return codes[-1:].to(torch.long).clone() + + def talker2code2wav_async_chunk( transfer_manager: Any, multimodal_output: OmniPayload | dict[str, Any], @@ -648,27 +716,47 @@ def talker2code2wav_async_chunk( """ request_id = request.external_req_id code_predictor_codes = None + frame_valid = None + first_audio = False if isinstance(multimodal_output, Mapping): talker_codes = multimodal_output.get("codes", {}) if isinstance(talker_codes, dict): code_predictor_codes = talker_codes.get("audio") + meta = multimodal_output.get("meta", {}) + if isinstance(meta, Mapping): + frame_valid = meta.get("codec_frame_valid") + flag = meta.get("first_audio", False) + first_audio = ( + bool(flag.numel() and flag.reshape(-1)[-1].item()) if isinstance(flag, torch.Tensor) else bool(flag) + ) - sampling_params = getattr(request, "sampling_params", None) - stop_token_ids = set(getattr(sampling_params, "stop_token_ids", None) or []) - stop_token_id = getattr(sampling_params, "stop_token_id", None) - if stop_token_id is not None: - stop_token_ids.add(stop_token_id) + if isinstance(frame_valid, torch.Tensor) and frame_valid.numel() > 0: + # Token-major codes with explicit per-row validity (MRv2 eager + # frames): a step's frame, if any, is the last row of its span, and a + # multi-row prefill span carries its first frame there. + frame = _last_codec_row(code_predictor_codes) + append_codes = frame is not None and bool(frame_valid.reshape(-1)[-1].item()) + code_predictor_codes = frame + else: + # One row per step (V1 and MRv2 deferred frames): all-zero rows are + # prefill placeholders and a stop-token CB0 marks the post-EOS step. + sampling_params = getattr(request, "sampling_params", None) + stop_token_ids = set(getattr(sampling_params, "stop_token_ids", None) or []) + stop_token_id = getattr(sampling_params, "stop_token_id", None) + if stop_token_id is not None: + stop_token_ids.add(stop_token_id) - append_codes = ( - isinstance(code_predictor_codes, torch.Tensor) - and code_predictor_codes.numel() > 0 - and bool(code_predictor_codes.any()) - ) - if append_codes: - first_codebook = int(code_predictor_codes[0, 0].item()) - if first_codebook in stop_token_ids: - logger.debug("skip stop-token codec frame: first_codebook=%s", first_codebook) - append_codes = False + append_codes = ( + isinstance(code_predictor_codes, torch.Tensor) + and code_predictor_codes.numel() > 0 + and bool(code_predictor_codes.any()) + ) + if append_codes: + assert isinstance(code_predictor_codes, torch.Tensor) + first_codebook = int(code_predictor_codes[0, 0].item()) + if first_codebook in stop_token_ids: + logger.debug("skip stop-token codec frame: first_codebook=%s", first_codebook) + append_codes = False if append_codes: transfer_manager.code_prompt_token_ids[request_id].append(code_predictor_codes) elif not is_finished: @@ -680,12 +768,42 @@ def talker2code2wav_async_chunk( chunk_size_config = int(cfg.get("codec_chunk_frames", 25)) left_context_size_config = int(cfg.get("codec_left_context_frames", 25)) configured_initial_chunk_size = int(cfg.get("initial_codec_chunk_frames") or 0) + # The connector config is fixed for the adapter's lifetime: parse once. + if not hasattr(transfer_manager, "_qwen3_omni_chunk_ramp"): + transfer_manager._qwen3_omni_chunk_ramp = parse_chunk_ramp(cfg, steady=chunk_size_config) + ramp = transfer_manager._qwen3_omni_chunk_ramp chunk_id = transfer_manager.put_req_chunk[request_id] length = len(transfer_manager.code_prompt_token_ids[request_id]) if length <= 0: return None + if ramp is not None: + # Ramp (replaces initial_codec_chunk_frames): chunk i carries ramp[i] + # new frames, then codec_chunk_frames. Code2Wav decodes statelessly, so + # every chunk re-sends up to codec_left_context_frames earlier frames. + # The counter restarts with each segment, like code_prompt_token_ids. + chunk_index = transfer_manager.ramp_chunk_count.get(request_id, 0) + emit, context_length = compute_ramp_emit(length, chunk_index, ramp, chunk_size_config, bool(is_finished)) + if not emit or context_length <= 0: + # Finished on a chunk boundary: the transport sends the finish marker. + return None + left_context_size = min(length - context_length, left_context_size_config) + end_index = left_context_size + context_length + codes = ( + torch.cat(transfer_manager.code_prompt_token_ids[request_id][-end_index:], dim=0) + .transpose(0, 1) + .reshape(-1) + ) + return OmniPayloadStruct( + codes=CodesStruct(audio=codes), + meta=MetaStruct( + first_audio=torch.tensor(first_audio and chunk_id == 0), + left_context_size=left_context_size, + finished=torch.tensor(is_finished, dtype=torch.bool), + ), + ) + if configured_initial_chunk_size > 0: if chunk_id == 0: chunk_size_config = configured_initial_chunk_size @@ -715,6 +833,7 @@ def talker2code2wav_async_chunk( return OmniPayloadStruct( codes=CodesStruct(audio=codes), meta=MetaStruct( + first_audio=torch.tensor(first_audio and chunk_id == 0), left_context_size=left_context_size, finished=torch.tensor(is_finished, dtype=torch.bool), ), diff --git a/vllm_omni/worker_v2/model_states/omni_model_state.py b/vllm_omni/worker_v2/model_states/omni_model_state.py index 6700e26b5..f20a0f2db 100644 --- a/vllm_omni/worker_v2/model_states/omni_model_state.py +++ b/vllm_omni/worker_v2/model_states/omni_model_state.py @@ -890,15 +890,15 @@ class OmniModelState(DefaultModelState): finishing = computed + scheduled >= input_batch.prefill_len_np[:num_reqs] kept = ~input_batch.is_prefilling_np[:num_reqs] | finishing rows = np.flatnonzero(kept).tolist() - if not rows: - return None device = sampled_token_ids.device - first = sampled_token_ids.reshape(num_reqs, -1)[:, 0].index_select(0, index_to_device(rows, device)) - embeds = self.model.embed_input_ids(first.long()) - empty = torch.empty(0, dtype=embeds.dtype) - sampled: list[torch.Tensor] = [empty] * num_reqs - for k, row in enumerate(rows): - sampled[row] = embeds[k : k + 1] + # Preserve the empty per-row prefill marker even when every request + # is prefilling; absence would look like a broken decode handoff. + sampled: list[torch.Tensor] = [torch.empty(0)] * num_reqs + if rows: + first = sampled_token_ids.reshape(num_reqs, -1)[:, 0].index_select(0, index_to_device(rows, device)) + embeds = self.model.embed_input_ids(first.long()) + for k, row in enumerate(rows): + sampled[row] = embeds[k : k + 1] done = torch.cuda.Event() done.record() extra: dict[str, Any] = {"embed": {"sampled": sampled}} From 53c7340668a8d87490934329310bf483615ae881 Mon Sep 17 00:00:00 2001 From: andyluo7 <43718156+andyluo7@users.noreply.github.com> Date: Thu, 1 Oct 2026 23:07:20 -0700 Subject: [PATCH 07/22] test: measure offload memory inside diffusion workers (#8189) Signed-off-by: andyluo7 --- .../test_diffusion_layerwise_offload.py | 152 +++++++++++------- 1 file changed, 93 insertions(+), 59 deletions(-) diff --git a/tests/diffusion/offloader/test_diffusion_layerwise_offload.py b/tests/diffusion/offloader/test_diffusion_layerwise_offload.py index c579affe9..678393439 100644 --- a/tests/diffusion/offloader/test_diffusion_layerwise_offload.py +++ b/tests/diffusion/offloader/test_diffusion_layerwise_offload.py @@ -13,7 +13,6 @@ import torch from vllm.distributed.parallel_state import cleanup_dist_env_and_memory from tests.helpers.mark import hardware_marks -from tests.helpers.monitor import DeviceMemoryMonitor from tests.helpers.runtime import OmniRunner from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform @@ -55,6 +54,7 @@ IMAGE_VIDEO_MODELS_PARAMS: dict[str, dict[str, Any]] = { } _OFFLOAD_STATE_PROBE = "vllm_omni_layerwise_offload_test" +_OFFLOAD_MEMORY_PROBE = "vllm_omni_layerwise_offload_memory_test" # A second CI build can reverse the same exact-head measurement without a # source edit; the default keeps the existing baseline-first workload. _MEASUREMENT_ORDER_ENV = "VLLM_OMNI_OFFLOAD_TEST_ORDER" @@ -65,7 +65,7 @@ _MEASUREMENT_ORDERS: dict[str, tuple[bool, bool]] = { class OffloadStateProbe: - """Test-only worker extension exposing the configured offloader state.""" + """Test-only worker extension exposing offloader and allocator state.""" model_runner: Any @@ -86,6 +86,21 @@ class OffloadStateProbe: "block_count": sum(group_sizes), } + def reset_peak_memory_for_test(self) -> dict[str, str]: + """Start a process-local inference peak after model initialization.""" + current_omni_platform.synchronize() + torch.accelerator.reset_peak_memory_stats() + return {"probe": _OFFLOAD_MEMORY_PROBE, "operation": "reset"} + + def get_peak_memory_for_test(self) -> dict[str, str | float]: + """Return this worker's allocator peak without counting other jobs.""" + current_omni_platform.synchronize() + return { + "probe": _OFFLOAD_MEMORY_PROBE, + "operation": "snapshot", + "peak_allocated_mb": torch.accelerator.max_memory_allocated() / (1024**2), + } + def check_audio_determinism(audio1: np.ndarray, audio2: np.ndarray, atol: float = 1e-2) -> bool: if not np.allclose(audio1, audio2, atol=atol): @@ -96,13 +111,6 @@ def check_audio_determinism(audio1: np.ndarray, audio2: np.ndarray, atol: float return True -def _device_used_mb(device_index: int) -> float: - current_omni_platform.synchronize() - with current_omni_platform.device(device_index): - free_bytes, total_bytes = current_omni_platform.mem_get_info() - return (total_bytes - free_bytes) / (1024**2) - - def _extract_audio(output: Any) -> np.ndarray | None: if not output: return None @@ -117,17 +125,29 @@ def _extract_audio(output: Any) -> np.ndarray | None: return np.asarray(audio) -def _collect_offload_states(value: Any) -> list[dict[str, Any]]: +def _collect_probe_results(value: Any, probe: str) -> list[dict[str, Any]]: if isinstance(value, dict): - return [value] if value.get("probe") == _OFFLOAD_STATE_PROBE else [] + return [value] if value.get("probe") == probe else [] if isinstance(value, (list, tuple)): - states: list[dict[str, Any]] = [] + results: list[dict[str, Any]] = [] for item in value: - states.extend(_collect_offload_states(item)) - return states + results.extend(_collect_probe_results(item, probe)) + return results return [] +def _collect_offload_states(value: Any) -> list[dict[str, Any]]: + return _collect_probe_results(value, _OFFLOAD_STATE_PROBE) + + +def _collect_memory_results(value: Any, operation: str) -> list[dict[str, Any]]: + return [ + result + for result in _collect_probe_results(value, _OFFLOAD_MEMORY_PROBE) + if result.get("operation") == operation + ] + + def _measurement_order() -> tuple[bool, bool]: """Return the requested order without changing the default CI workload.""" configured = os.environ.get(_MEASUREMENT_ORDER_ENV, "baseline-first") @@ -144,8 +164,6 @@ def run_inference( num_inference_steps: int = 3, ) -> dict[str, Any]: current_omni_platform.empty_cache() - device_index = current_omni_platform.current_device() - initial_used_mb = _device_used_mb(device_index) if model_name in AUDIO_MODEL: params = AUDIO_MODEL_PARAMS @@ -166,52 +184,45 @@ def run_inference( if not offload_states: raise AssertionError("The offload-state worker probe returned no diffusion worker results") - # Measure steady-state inference memory, not model construction. Enabling - # layerwise offload first loads the model and then replaces each block's - # device storage with CPU-backed weights. Monitoring that transition - # captures both the original model and temporary staging allocations, - # which is not representative of layerwise-offloaded inference. - monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02) - current_omni_platform.reset_peak_memory_stats() - monitor.start() + # Reset and read allocator peaks in the diffusion worker process. A + # device-wide mem_get_info() sample includes unrelated sibling jobs on + # the same GPU and produced negative "savings" in shared ROCm CI. + reset_results = _collect_memory_results( + runner.omni.engine.collective_rpc(method="reset_peak_memory_for_test", timeout=60), + "reset", + ) + if not reset_results: + raise AssertionError("The worker peak-memory probe returned no reset acknowledgements") - try: - # Refer to tests/e2e/offline_inference/test_wan22.py - # Use minimal settings for testing - output = runner.omni.generate( - "A cat sitting on a table", - OmniDiffusionSamplingParams( - generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42), - guidance_scale=1.0, - num_inference_steps=num_inference_steps, - **params["sampler_params"], - ), - ) - finally: - monitor.stop() + # Refer to tests/e2e/offline_inference/test_wan22.py. + # Use minimal settings for testing. + output = runner.omni.generate( + "A cat sitting on a table", + OmniDiffusionSamplingParams( + generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42), + guidance_scale=1.0, + num_inference_steps=num_inference_steps, + **params["sampler_params"], + ), + ) + memory_results = _collect_memory_results( + runner.omni.engine.collective_rpc(method="get_peak_memory_for_test", timeout=60), + "snapshot", + ) + if not memory_results: + raise AssertionError("The worker peak-memory probe returned no snapshots") audio = _extract_audio(output) del output - # DeviceMemoryMonitor reports absolute device usage. Subtract this run's - # starting usage. Each mode runs in its own spawned process below, so model - # objects and process-local allocator/compiler/backend workspace state cannot - # carry into the other measurement. Shared filesystem caches can persist; - # _MEASUREMENT_ORDER_ENV enables a reversed-order validation run for that. - peak_used_mb = monitor.peak_used_mb - incremental_peak_mb = max(0.0, peak_used_mb - initial_used_mb) - del runner gc.collect() cleanup_dist_env_and_memory() current_omni_platform.empty_cache() - post_cleanup_used_mb = _device_used_mb(device_index) return { - "initial_used_mb": initial_used_mb, - "peak_used_mb": peak_used_mb, - "incremental_peak_mb": incremental_peak_mb, - "post_cleanup_used_mb": post_cleanup_used_mb, + "peak_allocated_mb": max(float(result["peak_allocated_mb"]) for result in memory_results), + "memory_results": memory_results, "audio": audio, "offload_states": offload_states, } @@ -259,12 +270,8 @@ def _assert_audio_outputs( def _print_measurement(label: str, measurement: dict[str, Any]) -> None: - print( - f"{label}: initial={measurement['initial_used_mb']:.1f} MB, " - f"peak={measurement['peak_used_mb']:.1f} MB, " - f"incremental_peak={measurement['incremental_peak_mb']:.1f} MB, " - f"post_cleanup={measurement['post_cleanup_used_mb']:.1f} MB" - ) + print(f"{label}: worker_peak_allocated={measurement['peak_allocated_mb']:.1f} MB") + print(f"{label} worker memory: {measurement['memory_results']}") print(f"{label} offload state: {measurement['offload_states']}") @@ -298,6 +305,33 @@ def test_measurement_order_rejects_unknown_value(monkeypatch) -> None: _measurement_order() +@pytest.mark.diffusion +@pytest.mark.core_model +@pytest.mark.cpu +def test_collect_memory_results_filters_nested_worker_responses() -> None: + value = [ + None, + [ + {"probe": _OFFLOAD_MEMORY_PROBE, "operation": "reset"}, + { + "probe": _OFFLOAD_MEMORY_PROBE, + "operation": "snapshot", + "peak_allocated_mb": 123.0, + }, + ], + {"probe": _OFFLOAD_STATE_PROBE, "operation": "snapshot"}, + ] + + assert _collect_memory_results(value, "reset") == [{"probe": _OFFLOAD_MEMORY_PROBE, "operation": "reset"}] + assert _collect_memory_results(value, "snapshot") == [ + { + "probe": _OFFLOAD_MEMORY_PROBE, + "operation": "snapshot", + "peak_allocated_mb": 123.0, + } + ] + + @pytest.mark.diffusion @pytest.mark.core_model @pytest.mark.cpu @@ -366,8 +400,8 @@ def test_layerwise_offload_diffusion_model(model_name: str): # Verify that layerwise offloading significantly reduces memory usage # Passes only if the actual savings meets the expected savings - no_offload_peak_memory = no_offload["incremental_peak_mb"] - layerwise_offload_peak_memory = layerwise_offload["incremental_peak_mb"] + no_offload_peak_memory = no_offload["peak_allocated_mb"] + layerwise_offload_peak_memory = layerwise_offload["peak_allocated_mb"] actual_saved_memory = no_offload_peak_memory - layerwise_offload_peak_memory assert layerwise_offload_peak_memory + expected_saved_memory <= no_offload_peak_memory, ( f"Layerwise offload peak memory {layerwise_offload_peak_memory} MB " From b742f86136d28423903a1213adde2a62a11067a2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B1=AA=E5=BF=97=E9=B9=8F?= Date: Fri, 2 Oct 2026 16:45:56 +0800 Subject: [PATCH 08/22] [CI/Build] Restore streaming executor fixture contract (#8418) Signed-off-by: princepride --- .../test_diffusion_streaming_output.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py b/tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py index 68d17d49f..e9514abc9 100644 --- a/tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py +++ b/tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py @@ -77,6 +77,7 @@ class _PipelineBackedEngine: def __init__(self, pipeline: _StepStreamingPipeline | _FailingStreamingPipeline) -> None: self.pipeline = pipeline self.executor = SimpleNamespace( + is_dead=False, register_failure_callback=MagicMock(), check_health=MagicMock(), ) @@ -127,9 +128,9 @@ def _streaming_diffusion_output( def _streaming_payload_and_metadata(output: DiffusionOutput) -> tuple[dict, dict]: envelope = output.output if isinstance(output.output, dict) else {} - payload = envelope.get("payload") if isinstance(envelope.get("payload"), dict) else {} - metadata = envelope.get("metadata") if isinstance(envelope.get("metadata"), dict) else {} - return payload, metadata + payload = envelope.get("payload") + metadata = envelope.get("metadata") + return payload if isinstance(payload, dict) else {}, metadata if isinstance(metadata, dict) else {} class TestPipelineStreamingOutputToStageDiffusionClient: @@ -230,7 +231,7 @@ class TestPipelineStreamingOutputToStageDiffusionClient: client._response_socket.bind(response_address) client._encoder = OmniMsgpackEncoder() client._decoder = OmniMsgpackDecoder() - client._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue() + client._output_queue = asyncio.Queue[OmniRequestOutput]() client._rpc_results = {} client._pending_rpcs = set() client._tasks = {} From 97d6c693a684299d3696460d8942417231bfdd7f Mon Sep 17 00:00:00 2001 From: MaciejBalaNV Date: Fri, 2 Oct 2026 11:35:09 +0200 Subject: [PATCH 09/22] [feature] Add HSDP pre-sharded option (#7948) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Maciej Bala Co-authored-by: 汪志鹏 --- .../configuration/hsdp_pre_sharded_loading.md | 22 + .../distributed/test_cosmos3_pre_sharded.py | 200 +++++++ tests/diffusion/distributed/test_hsdp.py | 92 ++++ .../model_loader/test_diffusers_loader.py | 204 +++++++ .../cosmos3/test_cosmos3_transformer.py | 87 +++ .../test_async_omni_diffusion_config.py | 31 ++ vllm_omni/config/omni_config.py | 1 + vllm_omni/config/stage_config.py | 1 + vllm_omni/diffusion/data.py | 7 + vllm_omni/diffusion/distributed/hsdp.py | 116 +++- .../model_loader/diffusers_loader.py | 512 +++++++++++++++++- .../model_loader/host_weight_plan.py | 41 +- .../models/cosmos3/transformer_cosmos3.py | 77 +-- .../cosmos3/transformer_cosmos3_edge.py | 34 +- vllm_omni/diffusion/models/utils.py | 22 + vllm_omni/engine/arg_utils.py | 1 + vllm_omni/entrypoints/cli/serve.py | 9 + 17 files changed, 1370 insertions(+), 87 deletions(-) create mode 100644 docs/configuration/hsdp_pre_sharded_loading.md create mode 100644 tests/diffusion/distributed/test_cosmos3_pre_sharded.py diff --git a/docs/configuration/hsdp_pre_sharded_loading.md b/docs/configuration/hsdp_pre_sharded_loading.md new file mode 100644 index 000000000..cb8507b3d --- /dev/null +++ b/docs/configuration/hsdp_pre_sharded_loading.md @@ -0,0 +1,22 @@ +# Pre-sharded HSDP loading + +Use `--use-hsdp --hsdp-shard-size 8 --hsdp-weight-load-strategy pre_sharded` +to shard transformer parameters before loading checkpoint values. Set the shard +size to the available GPU count (or configure replication explicitly). +The default strategy, `full`, retains the existing load-then-shard behavior. +The same `hsdp_weight_load_strategy` option is available through Python and stage +engine configuration. + +The pre-sharded strategy reads rank-local slices from ordinary Hugging Face +safetensors checkpoints into FSDP-owned storage. No converted checkpoint is +required. Cosmos 3 and Cosmos 3 Edge transformer blocks release their initialized +parameter storage to meta as they are constructed; nonpersistent buffers retain +their values. +Encoders and VAEs follow their existing loading paths. + +This strategy requires the default diffusion pipeline loader, complete dedicated +transformer checkpoint sources, and supported checkpoint key mappings. Tensor +layouts must match the runtime layout; tensor transforms are unsupported. +Checkpoint tensor dtypes must match runtime tensor dtypes at preflight; dtype conversion is unsupported. +Quantization and LoRA are unsupported. Every worker needs checkpoint access. +`num_weight_load_threads` controls the safetensors reader's threads per rank. diff --git a/tests/diffusion/distributed/test_cosmos3_pre_sharded.py b/tests/diffusion/distributed/test_cosmos3_pre_sharded.py new file mode 100644 index 000000000..57af9c72b --- /dev/null +++ b/tests/diffusion/distributed/test_cosmos3_pre_sharded.py @@ -0,0 +1,200 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Compare full and pre-sharded Cosmos loading with real two-rank FSDP.""" + +from contextlib import ExitStack +from datetime import timedelta +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch import nn + +from tests.helpers.mark import hardware_test + +pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.core_model, pytest.mark.gpu] + + +class _SDPAAttention(nn.Module): + """Use deterministic native attention while testing loading and FSDP.""" + + def __init__(self, *, causal, softmax_scale, **kwargs): + super().__init__() + self.causal = causal + self.scale = softmax_scale + + def forward(self, query, key, value): + return F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + is_causal=self.causal, + scale=self.scale, + ).transpose(1, 2) + + +def _load_cosmos_worker(rank, rendezvous, checkpoint_dir, edge): + from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules + + bootstrap_vllm_layer_custom_op_modules() + + from safetensors.torch import save_file + from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config + from vllm.config.load import LoadConfig + from vllm.model_executor import parameter + from vllm.model_executor.layers import linear + from vllm.utils.torch_utils import set_default_torch_dtype + + from vllm_omni.diffusion.distributed import hsdp + from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader + from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3, transformer_cosmos3_edge + from vllm_omni.platforms import current_omni_platform + + device = torch.device("cuda", rank) + current_omni_platform.set_device(device) + dist.init_process_group("nccl", init_method=rendezvous, rank=rank, world_size=2, timeout=timedelta(seconds=120)) + try: + with ExitStack() as stack: + stack.enter_context(set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cuda")))) + stack.enter_context(set_default_torch_dtype(torch.bfloat16)) + # HSDP spans both workers, while every vLLM linear uses TP=1. + for module in (linear, parameter): + stack.enter_context(patch.object(module, "get_tensor_model_parallel_rank", return_value=0)) + stack.enter_context(patch.object(module, "get_tensor_model_parallel_world_size", return_value=1)) + for module in (transformer_cosmos3, transformer_cosmos3_edge): + stack.enter_context(patch.object(module, "get_tensor_model_parallel_world_size", return_value=1)) + stack.enter_context(patch.object(module, "FrameworkAttention", _SDPAAttention)) + stack.enter_context(patch.object(transformer_cosmos3, "_get_ulysses_state", return_value=(1, 0, None))) + stack.enter_context( + patch.object(hsdp, "get_world_group", return_value=SimpleNamespace(world_size=2, rank_in_group=rank)) + ) + + model_cls = ( + transformer_cosmos3_edge.Cosmos3EdgeVFMTransformer + if edge + else transformer_cosmos3.Cosmos3VFMTransformer + ) + tf_config = dict( + hidden_size=16, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + head_dim=8, + intermediate_size=32, + vocab_size=32, + latent_patch_size=2, + latent_channel=48, + temporal_compression_factor=4, + rope_scaling={"mrope_section": [2, 1, 1]}, + ) + if edge: + tf_config.update( + backbone_type=transformer_cosmos3_edge.COSMOS3_EDGE_BACKBONE_TYPE, + qk_norm_for_text=False, + use_und_k_norm_for_gen=True, + ) + config = SimpleNamespace( + tf_model_config=tf_config, + dtype=torch.bfloat16, + quantization_config=None, + lora_path=None, + num_weight_load_threads=1, + hsdp_weight_load_strategy="full", + parallel_config=SimpleNamespace( + use_hsdp=True, hsdp_replicate_size=1, hsdp_shard_size=2, tensor_parallel_size=1 + ), + ) + + class Pipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = model_cls(config) + self.transformer.register_buffer("scratch", torch.tensor([7.0]), persistent=False) + self.vae = nn.Linear(2, 2) + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource(checkpoint_dir, None, None, "transformer.", False) + ] + + def load_weights(self, weights): + state = dict(weights) + self.transformer.load_state_dict( + {name.removeprefix("transformer."): t for name, t in state.items()} + ) + self.transformer.post_load_weights() + self.transformer.eval() + self.transformer.validate_loaded_weights(set(state)) + return set(state) + + full = Pipeline() + if rank == 0: + torch.manual_seed(42) + with torch.no_grad(): + for name, tensor in full.transformer.named_parameters(): + tensor.copy_(torch.randn_like(tensor) * 0.02) + if "norm" in name: + tensor.fill_(1) + save_file(full.transformer.state_dict(), str(Path(checkpoint_dir) / "model.safetensors")) + dist.barrier() + + def load(model): + loader = DiffusersPipelineLoader(LoadConfig(), config) + with patch.object(loader, "_init_from_load_format", return_value=model): + return loader._load_model_with_hsdp(device) + + full = load(full) + config.hsdp_weight_load_strategy = "pre_sharded" + pre_sharded = load(Pipeline()) + full_parameters = dict(full.transformer.named_parameters()) + from torch.distributed.tensor import DTensor + + for name, actual in pre_sharded.transformer.named_parameters(): + expected = full_parameters[name] + if name.startswith("time_embedder."): + assert not isinstance(actual, DTensor) + assert actual.dtype == torch.float32 + else: + assert isinstance(actual, DTensor) + actual = actual.full_tensor() + expected = expected.full_tensor() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert pre_sharded.vae.weight.device == device + torch.testing.assert_close(pre_sharded.transformer.scratch, full.transformer.scratch) + assert all(t.device.type != "meta" for t in pre_sharded.transformer.buffers()) + + inputs = dict( + hidden_states=torch.linspace(-1, 1, 48 * 4, device=device).reshape(1, 48, 1, 2, 2), + timestep=torch.tensor([1.0], device=device), + text_ids=torch.tensor([[1, 2]], device=device), + text_mask=torch.ones(1, 2, dtype=torch.long, device=device), + video_shape=(1, 2, 2), + fps=24.0, + ) + with torch.no_grad(): + # Also exercise the cached UND path on the second denoising step. + for step in (1.0, 0.5): + inputs["timestep"].fill_(step) + expected = full.transformer(**inputs) + actual = pre_sharded.transformer(**inputs) + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("edge", [False, True], ids=["cosmos3", "cosmos3_edge"]) +@hardware_test(res={"cuda": "L4"}, num_cards=2) +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.accelerator.device_count() < 2 or not dist.is_nccl_available(), + reason="requires two CUDA GPUs and NCCL", +) +def test_cosmos3_pre_sharded_matches_full_hsdp(tmp_path, edge): + torch.multiprocessing.spawn( + _load_cosmos_worker, + args=(f"file://{tmp_path / 'rendezvous'}", str(tmp_path), edge), + nprocs=2, + join=True, + ) diff --git a/tests/diffusion/distributed/test_hsdp.py b/tests/diffusion/distributed/test_hsdp.py index 3c3805a82..b1911e97c 100644 --- a/tests/diffusion/distributed/test_hsdp.py +++ b/tests/diffusion/distributed/test_hsdp.py @@ -3,11 +3,13 @@ """Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities.""" import gc +from types import SimpleNamespace import pytest import torch import torch.distributed as dist import torch.nn as nn +from torch.distributed.fsdp import fully_shard from torch.distributed.tensor import DeviceMesh, DTensor from vllm.utils.network_utils import get_file_store_init_method @@ -18,6 +20,7 @@ from vllm_omni.diffusion.distributed.hsdp import ( _unshardable_parameters, shard_model, ) +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.cpu, pytest.mark.core_model] @@ -338,3 +341,92 @@ class TestHSDPShardConditions: matched.append(name) assert "blocks.0" in matched assert "blocks.1" in matched + + +def test_pre_sharded_to_empty_materializes_fsdp_internal_storage(cpu_process_group): + with torch.device("meta"): + model = nn.Module() + model.transformer = nn.Sequential(nn.Linear(2, 2, bias=False)) + model.transformer.register_buffer("scratch", torch.tensor([3.0]), persistent=False) + shard_model( + model.transformer, + mesh=DeviceMesh("cpu", [0]), + hsdp_shard_conditions=[lambda name, _module: name == "0"], + ) + + load_plan = SimpleNamespace( + roots=[SimpleNamespace(name="transformer", module=model.transformer)], + groups=[SimpleNamespace(module=model.transformer[0])], + binding_names=frozenset({"transformer.0.weight"}), + ) + sharded_weight = model.transformer[0].weight + + DiffusersPipelineLoader._materialize_pre_sharded_hsdp_state(model, load_plan, torch.device("cpu")) + + fsdp_param = fully_shard.state(model.transformer[0])._fsdp_param_group.fsdp_params[0] + assert model.transformer[0].weight is sharded_weight + assert model.transformer[0].weight.to_local().device.type == "cpu" + assert fsdp_param._sharded_param_data.device.type == "cpu" + assert ( + fsdp_param._sharded_param_data.untyped_storage().data_ptr() + == model.transformer[0].weight.to_local().untyped_storage().data_ptr() + ) + assert torch.equal(model.transformer.scratch, torch.tensor([3.0])) + + +def test_pre_sharded_load_restores_checkpoint_values(cpu_process_group, tmp_path, monkeypatch): + from safetensors.torch import save_file + from vllm.config.load import LoadConfig + + from vllm_omni.diffusion.distributed import hsdp as hsdp_module + from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta + + class Transformer(nn.Module): + _hsdp_shard_conditions = [lambda name, module: name == "block"] + _hsdp_ignored_modules = ["time_embedder"] + + def __init__(self): + super().__init__() + self.block = nn.Linear(3, 4, bias=False) + self.time_embedder = nn.Linear(2, 2, bias=False, dtype=torch.bfloat16) + self.root_weight = nn.Parameter(torch.zeros(2)) + self.register_buffer("scratch", torch.tensor([7.0]), persistent=False) + release_module_parameters_to_meta(self.block) + + def post_load_weights(self): + self.time_embedder.to(torch.float32) + + model = nn.Module() + model.transformer = Transformer() + model.vae = nn.Linear(2, 2) + tensors = { + "block.weight": torch.arange(12, dtype=torch.float32).reshape(4, 3), + "root_weight": torch.tensor([2.0, 3.0]), + "time_embedder.weight": torch.arange(4, dtype=torch.bfloat16).reshape(2, 2), + } + save_file(tensors, str(tmp_path / "model.safetensors")) + model.weights_sources = [DiffusersPipelineLoader.ComponentSource(str(tmp_path), None, None, "transformer.", False)] + config = SimpleNamespace( + dtype=torch.float32, + quantization_config=None, + lora_path=None, + hsdp_weight_load_strategy="pre_sharded", + num_weight_load_threads=1, + parallel_config=SimpleNamespace(use_hsdp=True, hsdp_replicate_size=1, hsdp_shard_size=1), + ) + loader = DiffusersPipelineLoader(LoadConfig(), config) + monkeypatch.setattr(loader, "_init_from_load_format", lambda *args, **kwargs: model) + monkeypatch.setattr(hsdp_module, "get_world_group", lambda: SimpleNamespace(world_size=1, rank_in_group=0)) + monkeypatch.setattr(hsdp_module, "current_omni_platform", SimpleNamespace(device_type="cpu")) + loaded = loader._load_model_with_hsdp(torch.device("cpu")) + assert loaded is model + for name, expected in tensors.items(): + actual = dict(model.transformer.named_parameters())[name] + if name.startswith("time_embedder."): + assert not isinstance(actual, DTensor) + assert actual.dtype == torch.float32 + torch.testing.assert_close(actual, expected.float()) + else: + assert isinstance(actual, DTensor) + torch.testing.assert_close(actual.to_local(), expected) + torch.testing.assert_close(model.transformer.scratch, torch.tensor([7.0])) diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 8a2ee0900..60c0cd7d6 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -27,6 +27,7 @@ from vllm_omni.diffusion.model_loader.host_weight_plan import ( HostWeightPlan, HostWeightPlanResult, TensorBinding, + build_checkpoint_binding_plan, ) from vllm_omni.diffusion.model_loader.host_weights import source_identity as source_identity_module from vllm_omni.diffusion.models.helios import HeliosPipeline @@ -2089,3 +2090,206 @@ def test_hsdp_broadcast_weight_load_online_quant_multiprocess(): ) assert os.path.exists(os.path.join(temp_dir, "rank_0_success.flag")) assert os.path.exists(os.path.join(temp_dir, "rank_1_success.flag")) + + +def test_pre_sharded_hsdp_rejects_quantization_before_loading(): + od_config = SimpleNamespace( + dtype=torch.float32, + hsdp_weight_load_strategy="pre_sharded", + lora_path=None, + quantization_config=object(), + parallel_config=SimpleNamespace( + use_hsdp=True, + hsdp_replicate_size=1, + hsdp_shard_size=2, + ), + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + with pytest.raises(ValueError, match="does not support quantization"): + loader._load_model_with_hsdp(torch.device("cpu")) + + +def test_pre_sharded_hsdp_strategy_dispatches_without_using_full_loader(mocker): + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + from vllm_omni.diffusion.offloader.module_collector import PipelineModules + + od_config = SimpleNamespace( + dtype=torch.float32, + hsdp_weight_load_strategy="pre_sharded", + lora_path=None, + quantization_config=None, + parallel_config=SimpleNamespace( + use_hsdp=True, + hsdp_replicate_size=1, + hsdp_shard_size=2, + ), + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + discovered = PipelineModules( + dits=[model.transformer], + dit_names=["transformer"], + vaes=[], + encoders=[], + encoder_names=[], + resident_modules=[], + resident_names=[], + ) + loader._init_from_load_format = mocker.Mock(return_value=model) # type: ignore[method-assign] + loader.load_weights = mocker.Mock(side_effect=AssertionError("full loader must not run")) # type: ignore[method-assign] + mocker.patch.object(loader_mod.ModuleDiscovery, "discover", return_value=discovered) + pre_sharded_load = mocker.patch.object( + loader, + "_load_model_with_pre_sharded_hsdp", + return_value=model, + ) + + assert loader._load_model_with_hsdp(torch.device("cpu")) is model + pre_sharded_load.assert_called_once() + + +def test_pre_sharded_hsdp_rejects_tensor_transform(): + def transform(tensor): + return tensor.transpose(-2, -1) + + load_plan = SimpleNamespace( + bindings={ + "transformer.weight": TensorBinding( + checkpoint_key="weight", + file_path="model.safetensors", + transform=transform, + ) + } + ) + + with pytest.raises(ValueError, match="tensor transforms are unsupported"): + DiffusersPipelineLoader._validate_pre_sharded_hsdp_bindings(load_plan) + + +def test_pre_sharded_hsdp_accepts_direct_safetensors_binding(): + load_plan = SimpleNamespace( + bindings={ + "transformer.weight": TensorBinding( + checkpoint_key="weight", + file_path="model.safetensors", + ) + } + ) + + DiffusersPipelineLoader._validate_pre_sharded_hsdp_bindings(load_plan) + + +def test_pre_sharded_hsdp_uses_fsdp_destination_directly(): + target = torch.empty(8, 4, 3) + load_plan = SimpleNamespace( + bindings={ + "transformer.weight": TensorBinding( + checkpoint_key="source_weight", + file_path="transformer/model.safetensors", + ) + } + ) + + checkpoint_states, local_bytes = DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states( + load_plan, + {"transformer.weight": target}, + ) + checkpoint_target = checkpoint_states[Path("transformer")]["source_weight"] + + assert checkpoint_target is target + assert checkpoint_target.shape == (8, 4, 3) + assert checkpoint_target.device.type == "cpu" + assert checkpoint_target.untyped_storage().data_ptr() == target.untyped_storage().data_ptr() + assert local_bytes == target.numel() * target.element_size() + + +def test_pre_sharded_hsdp_checkpoint_writes_runtime_storage_directly(): + runtime_target = torch.empty(2, 3) + load_plan = SimpleNamespace( + bindings={ + "transformer.weight": TensorBinding( + checkpoint_key="source_weight", + file_path="transformer/model.safetensors", + ) + } + ) + checkpoint_tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3) + checkpoint_states, _ = DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states( + load_plan, + {"transformer.weight": runtime_target}, + ) + checkpoint_states[Path("transformer")]["source_weight"].copy_(checkpoint_tensor) + + assert torch.equal(runtime_target, checkpoint_tensor) + + +def test_pre_sharded_hsdp_rejects_duplicate_checkpoint_destinations(): + binding = TensorBinding(checkpoint_key="weight", file_path="transformer/model.safetensors") + load_plan = SimpleNamespace( + bindings={ + "transformer.first": binding, + "transformer.second": binding, + } + ) + + with pytest.raises(ValueError, match="multiple runtime tensors"): + DiffusersPipelineLoader._prepare_pre_sharded_checkpoint_states( + load_plan, + {"transformer.first": torch.empty(2, 3), "transformer.second": torch.empty(2, 3)}, + ) + + +def test_pre_sharded_hsdp_managed_module_traversal_is_child_first_and_skips_nested_roots(): + root = nn.Module() + root.left = nn.Sequential(nn.Linear(2, 2), nn.ReLU()) + root.right = nn.Sequential(nn.Linear(2, 2), nn.ReLU()) + + names = {module: name or "root" for name, module in root.named_modules()} + modules = DiffusersPipelineLoader._hsdp_managed_modules_post_order( + root, + nested_hsdp_roots={root.left}, + ) + + assert [names[module] for module in modules] == ["right.0", "right.1", "right", "root"] + + +@pytest.mark.parametrize("checkpoint_key", ["renamed", "unexpected"]) +def test_hsdp_checkpoint_plan_honors_remap_and_rejects_missing(tmp_path, checkpoint_key, monkeypatch): + import vllm_omni.diffusion.model_loader.host_weight_plan as plan_mod + + monkeypatch.setattr(plan_mod, "get_direct_mmap_adapter", lambda _model: None) + + class Pipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = nn.Linear(2, 2, bias=False) + + @staticmethod + def remap_checkpoint_key(name): + return "transformer.weight" if name == "transformer.renamed" else name + + checkpoint = tmp_path / "model.safetensors" + save_file({checkpoint_key: torch.ones(2, 2)}, str(checkpoint)) + source = SimpleNamespace( + model_or_path=str(tmp_path), + subfolder=None, + revision=None, + prefix="transformer.", + ) + model = Pipeline() + result = build_checkpoint_binding_plan( + model, + dit_modules=(("transformer", model.transformer),), + sources=(source,), + model_path=None, + tensor_parallel_size=1, + online_quantization=False, + ) + + if checkpoint_key == "renamed": + assert result.plan is not None + assert result.plan.bindings["transformer.weight"].checkpoint_key == "renamed" + else: + assert result.plan is None + assert "no checkpoint binding" in result.fallback_reason diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py index 58fe4b10b..56c33b9e0 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py @@ -868,6 +868,93 @@ def test_qwen_gen_mlp_remains_gated() -> None: assert hasattr(model.gen_layers[0].mlp, "gate_proj") +@pytest.mark.parametrize("edge", [False, True]) +@pytest.mark.parametrize( + "use_hsdp,strategy", [(False, "full"), (False, "pre_sharded"), (True, "full"), (True, "pre_sharded")] +) +def test_cosmos3_releases_blocks_only_for_pre_sharded_hsdp(edge, use_hsdp, strategy, monkeypatch): + from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3, transformer_cosmos3_edge + from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta + + released = [] + + def release_and_check(block): + parameters = dict(block.named_parameters()) + metadata = {name: dict(parameter.__dict__) for name, parameter in parameters.items()} + release_module_parameters_to_meta(block) + for name, parameter in block.named_parameters(): + assert parameter is parameters[name] + assert parameter.device.type == "meta" + assert parameter.__dict__ == metadata[name] + released.append(block) + + monkeypatch.setattr(transformer_cosmos3, "release_module_parameters_to_meta", release_and_check) + monkeypatch.setattr(transformer_cosmos3_edge, "release_module_parameters_to_meta", release_and_check) + model_cls = ( + transformer_cosmos3_edge.Cosmos3EdgeVFMTransformer if edge else transformer_cosmos3.Cosmos3VFMTransformer + ) + config = _tiny_cosmos3_edge_config if edge else _tiny_cosmos3_config + model = model_cls( + SimpleNamespace( + tf_model_config=config(num_hidden_layers=1), + dtype=torch.float32, + parallel_config=SimpleNamespace(use_hsdp=use_hsdp), + hsdp_weight_load_strategy=strategy, + ) + ) + blocks = [*model.language_model.layers, *model.gen_layers] + expected_meta = use_hsdp and strategy == "pre_sharded" + assert len(released) == (len(blocks) if expected_meta else 0) + assert all((p.device.type == "meta") == expected_meta for block in blocks for p in block.parameters()) + assert model.time_embedder.linear_1.weight.device.type != "meta" + assert all(buffer.device.type != "meta" for buffer in model.buffers()) + + +@pytest.mark.parametrize("edge", [False, True]) +def test_cosmos3_pre_sharded_meta_processing_accepts_real_unquantized_blocks(edge, monkeypatch): + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + + from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader + from vllm_omni.diffusion.models.cosmos3 import Cosmos3EdgeVFMTransformer, Cosmos3VFMTransformer + + model_cls = Cosmos3EdgeVFMTransformer if edge else Cosmos3VFMTransformer + config = _tiny_cosmos3_edge_config if edge else _tiny_cosmos3_config + model = model_cls( + SimpleNamespace( + tf_model_config=config(num_hidden_layers=1), + dtype=torch.float32, + parallel_config=SimpleNamespace(use_hsdp=True), + hsdp_weight_load_strategy="pre_sharded", + ) + ) + groups = [SimpleNamespace(module=block) for block in (*model.language_model.layers, *model.gen_layers)] + plan = SimpleNamespace(roots=[SimpleNamespace(module=model)], groups=groups) + processed = [] + # Exercise model selection with real vLLM linears, without invoking the + # CPU-specific weight packing hook on meta tensors in CPU-only CI. + monkeypatch.setattr( + UnquantizedLinearMethod, "process_weights_after_loading", lambda self, layer: processed.append(layer) + ) + DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan) + expected = [module for group in groups for module in group.module.modules() if hasattr(module, "quant_method")] + assert expected + assert processed == expected + + model._hsdp_pre_sharded_meta_post_load = False + with pytest.raises(ValueError, match="has not declared"): + DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan) + model._hsdp_pre_sharded_meta_post_load = True + + class UnsupportedMethod(UnquantizedLinearMethod): + pass + + expected[-1].quant_method = UnsupportedMethod() + processed.clear() + with pytest.raises(ValueError, match="UnsupportedMethod"): + DiffusersPipelineLoader._process_pre_sharded_hsdp_groups_on_meta(plan) + assert not processed + + def test_model_cpu_offload_swaps_back_to_generator(monkeypatch: pytest.MonkeyPatch) -> None: from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 1d106932d..496d19320 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -11,6 +11,7 @@ from pydantic import ValidationError from vllm_omni.config.config_factory import StageConfigFactory from vllm_omni.config.omni_config import VllmOmniDiffusionStageConfig, extract_diffusion_stage_config_kwargs from vllm_omni.config.resolver import OmniConfigResolution, resolve_omni_config +from vllm_omni.config.stage_config import build_stage_runtime_overrides 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 @@ -945,3 +946,33 @@ def test_generic_diffusion_structured_stage_reaches_standard_startup(mocker): assert launched["model"] == "generic-diffusion" assert launched["stage_id"] == 0 assert runtime.stage_pools[0].clients == [client] + + +def test_default_stage_config_preserves_pre_sharded_hsdp_weight_load() -> None: + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"hsdp_weight_load_strategy": "pre_sharded"})[0] + + assert stage_cfg["engine_args"]["hsdp_weight_load_strategy"] == "pre_sharded" + + +@pytest.mark.parametrize("prefix", ["", "stage_0_"]) +def test_pre_sharded_hsdp_weight_load_survives_stage_overrides(prefix: str) -> None: + overrides = build_stage_runtime_overrides(0, {f"{prefix}hsdp_weight_load_strategy": "pre_sharded"}) + + assert overrides["hsdp_weight_load_strategy"] == "pre_sharded" + + +@pytest.mark.parametrize("strategy", ["full", "pre_sharded"]) +def test_typed_default_stage_preserves_hsdp_weight_load_strategy(strategy: str) -> None: + config = StageConfigFactory.create_typed_default_diffusion( + "generic-diffusion", {"hsdp_weight_load_strategy": strategy} + ) + + assert config.stage_configs[0].diffusion_config.hsdp_weight_load_strategy == strategy + + +@pytest.mark.parametrize("strategy", ["memory_limited", "invalid"]) +def test_removed_hsdp_loading_strategies_are_rejected(strategy): + from vllm_omni.diffusion.data import OmniDiffusionConfig + + with pytest.raises(ValueError, match="hsdp_weight_load_strategy"): + OmniDiffusionConfig(hsdp_weight_load_strategy=strategy) diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 526f4c9e5..d77e44194 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -817,6 +817,7 @@ class _DiffusionConfigProjection: prompt_embed_cache_size: int = Field(default=32, ge=1) enable_session_state_manager: bool = False diffusion_load_format: str = "default" + hsdp_weight_load_strategy: str = "full" diffusers_load_kwargs: dict[str, Any] = field(default_factory=dict) diffusers_call_kwargs: dict[str, Any] = field(default_factory=dict) diffusers_pipeline_cls: Any = None diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 0c46e9c22..5f959050a 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -492,6 +492,7 @@ class StageDeployConfig: auxiliary_text_encoder: str | None = None # Runtime optimizations used by diffusion loading/execution. + hsdp_weight_load_strategy: str | None = None enable_multithread_weight_load: bool | None = None enable_broadcast_weight_load: bool | None = None num_weight_load_threads: int | None = None diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 5e54637c2..0811d2634 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -983,6 +983,9 @@ class OmniDiffusionConfig: enable_broadcast_weight_load: bool = False num_weight_load_threads: int = 4 + # Shard meta parameters before reading rank-local HF safetensors slices. + hsdp_weight_load_strategy: str = "full" + # Enable sleep mode enable_sleep_mode: bool = False @@ -1196,6 +1199,10 @@ class OmniDiffusionConfig: ) def __post_init__(self): + if self.hsdp_weight_load_strategy not in {"full", "pre_sharded"}: + raise ValueError( + f"hsdp_weight_load_strategy must be 'full' or 'pre_sharded', got {self.hsdp_weight_load_strategy!r}" + ) from vllm_omni.diffusion.offloader.config import ( OffloadStrategy, materialize_legacy_offload_flags, diff --git a/vllm_omni/diffusion/distributed/hsdp.py b/vllm_omni/diffusion/distributed/hsdp.py index 82c1d1986..7e20d8471 100644 --- a/vllm_omni/diffusion/distributed/hsdp.py +++ b/vllm_omni/diffusion/distributed/hsdp.py @@ -46,6 +46,14 @@ class HSDPInferenceConfig: reshard_after_forward: bool = True +@dataclass +class HSDPShardContext: + """Reusable state for incrementally sharding one HSDP model.""" + + hsdp_kwargs: dict[str, Any] + ignored_params: set[nn.Parameter] | None = None + + def _create_hsdp_mesh( device_type: str, replicate_size: int, @@ -111,6 +119,38 @@ def apply_hsdp_to_model( Returns: HSDP-wrapped model ready for inference """ + context = prepare_hsdp_shard_context( + model, + hsdp_config, + target_device=target_device, + ) + + hsdp_shard_conditions = getattr(model, "_hsdp_shard_conditions", None) + if not hsdp_shard_conditions: + raise ValueError(f"Model {type(model).__name__} has no _hsdp_shard_conditions defined") + + # Apply HSDP sharding, this will automatically handle weight distribution + shard_model( + model, + hsdp_shard_conditions=hsdp_shard_conditions, + context=context, + ) + + logger.info("HSDP applied to model: %s", type(model).__name__) + return model + + +def prepare_hsdp_shard_context( + model: nn.Module, + hsdp_config: HSDPInferenceConfig, + target_device: torch.device | None = None, +) -> HSDPShardContext: + """Prepare mesh, precision, and ignored-parameter state for HSDP. + + The returned context can be reused to shard completed child modules while + checkpoint weights are still streaming, then to shard the root once all + children have been loaded. + """ if not hsdp_config.enabled: raise ValueError("HSDP is not enabled in config") @@ -160,7 +200,7 @@ def apply_hsdp_to_model( ) hsdp_shard_conditions = getattr(model, "_hsdp_shard_conditions", None) - if not hsdp_shard_conditions or len(hsdp_shard_conditions) == 0: + if not hsdp_shard_conditions: raise ValueError(f"Model {type(model).__name__} has no _hsdp_shard_conditions defined") # Collect parameters of any modules the model wants excluded from FSDP sharding. @@ -188,7 +228,11 @@ def apply_hsdp_to_model( except AttributeError: logger.warning("_hsdp_ignored_modules entry %r not found on model", mod_name) continue - sub_mod.to(target_device) + has_meta_tensor = any(tensor.device.type == "meta" for tensor in (*sub_mod.parameters(), *sub_mod.buffers())) + if has_meta_tensor: + sub_mod.to_empty(device=target_device) + else: + sub_mod.to(target_device) ignored_params.update(sub_mod.parameters()) if ignored_params: logger.info( @@ -224,21 +268,36 @@ def apply_hsdp_to_model( target_device, ) - # Apply HSDP sharding, this will automatically handle weight distribution - shard_model( - model, - reshard_after_forward=hsdp_config.reshard_after_forward, - mp_policy=mp_policy, - mesh=device_mesh, - hsdp_shard_conditions=hsdp_shard_conditions, - ignored_params=ignored_params if ignored_params else None, + return HSDPShardContext( + hsdp_kwargs={ + "reshard_after_forward": hsdp_config.reshard_after_forward, + "mesh": device_mesh, + "mp_policy": mp_policy, + }, + ignored_params=ignored_params or None, ) - for param in model.parameters(): + +def shard_hsdp_module(module: nn.Module, context: HSDPShardContext) -> None: + """Shard one completed child module with a prepared HSDP context.""" + module_kwargs = dict(context.hsdp_kwargs) + if context.ignored_params: + module_ignored_params = context.ignored_params.intersection(module.parameters()) + if module_ignored_params: + module_kwargs["ignored_params"] = module_ignored_params + fully_shard(module, **module_kwargs) + for param in module.parameters(): param.requires_grad = False - logger.info("HSDP applied to model: %s", type(model).__name__) - return model + +def finalize_hsdp_root(model: nn.Module, context: HSDPShardContext) -> None: + """Shard the model root after all selected child modules are sharded.""" + root_kwargs = dict(context.hsdp_kwargs) + if context.ignored_params: + root_kwargs["ignored_params"] = context.ignored_params + fully_shard(model, **root_kwargs) + for param in model.parameters(): + param.requires_grad = False def shard_model( @@ -249,6 +308,7 @@ def shard_model( mesh: DeviceMesh | None = None, hsdp_shard_conditions: list[Callable[[str, nn.Module], bool]], ignored_params: set[nn.Parameter] | None = None, + context: HSDPShardContext | None = None, ) -> None: """Apply HSDP sharding to model modules based on shard conditions. @@ -259,32 +319,30 @@ def shard_model( This is required for packed integer parameters inside sharded transformer blocks; the root wrap receives the full set for all remaining parameters. """ - hsdp_kwargs: dict[str, Any] = { - "reshard_after_forward": reshard_after_forward, - "mesh": mesh, - "mp_policy": mp_policy, - } + if context is None: + context = HSDPShardContext( + hsdp_kwargs={ + "reshard_after_forward": reshard_after_forward, + "mesh": mesh, + "mp_policy": mp_policy, + }, + ignored_params=ignored_params, + ) + elif any(value is not None for value in (mp_policy, mesh, ignored_params)) or not reshard_after_forward: + raise ValueError("Pass either context or individual HSDP shard options, not both") num_sharded = 0 for name, module in reversed(list(model.named_modules())): if any(cond(name, module) for cond in hsdp_shard_conditions): - module_kwargs = dict(hsdp_kwargs) - if ignored_params: - module_ignored_params = ignored_params.intersection(module.parameters()) - if module_ignored_params: - module_kwargs["ignored_params"] = module_ignored_params - fully_shard(module, **module_kwargs) + shard_hsdp_module(module, context) num_sharded += 1 if num_sharded == 0: raise ValueError("No modules were sharded") - root_kwargs = dict(hsdp_kwargs) - if ignored_params: - root_kwargs["ignored_params"] = ignored_params - fully_shard(model, **root_kwargs) + finalize_hsdp_root(model, context) logger.info( "Sharded %d modules + root (ignored_params=%d)", num_sharded, - len(ignored_params) if ignored_params else 0, + len(context.ignored_params) if context.ignored_params else 0, ) diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 7cbff273c..c74808efb 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -16,6 +16,7 @@ import torch from torch import nn from vllm.config.load import LoadConfig from vllm.logger import init_logger +from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase from vllm.model_executor.model_loader.weight_utils import ( download_weights_from_hf, @@ -30,7 +31,14 @@ from vllm.utils.torch_utils import set_default_torch_dtype from vllm_omni.diffusion.config import set_current_diffusion_config from vllm_omni.diffusion.data import OmniDiffusionConfig -from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp_to_model +from vllm_omni.diffusion.distributed.hsdp import ( + HSDPInferenceConfig, + HSDPShardContext, + apply_hsdp_to_model, + finalize_hsdp_root, + prepare_hsdp_shard_context, + shard_hsdp_module, +) from vllm_omni.diffusion.lora.manager import LoRABackend from vllm_omni.diffusion.model_loader.checkpoint_adapters import ( get_checkpoint_adapter, @@ -38,6 +46,8 @@ from vllm_omni.diffusion.model_loader.checkpoint_adapters import ( from vllm_omni.diffusion.model_loader.host_weight_loader import HWRLoaderMixin, _HWRCommitError from vllm_omni.diffusion.model_loader.host_weight_plan import ( HostWeightPlan, + TensorBinding, + build_checkpoint_binding_plan, build_checkpoint_mmap_plan, has_online_quantization, ) @@ -49,7 +59,7 @@ from vllm_omni.diffusion.offloader.config import ( OffloadStrategy, resolve_offload, ) -from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery +from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery, PipelineModules from vllm_omni.diffusion.offloader.offload_plan import get_offload_plan from vllm_omni.diffusion.registry import initialize_model from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific @@ -138,6 +148,36 @@ def _resolve_custom_pipeline_cls(custom_pipeline_name: str | type | None) -> typ ) +def _is_distributed_rank_zero() -> bool: + return ( + not torch.distributed.is_available() + or not torch.distributed.is_initialized() + or torch.distributed.get_rank() == 0 + ) + + +@dataclasses.dataclass(frozen=True) +class _HSDPLoadRoot: + name: str + module: nn.Module + context: HSDPShardContext + + +@dataclasses.dataclass(frozen=True) +class _HSDPModuleLoadGroup: + name: str + module: nn.Module + context: HSDPShardContext + + +@dataclasses.dataclass(frozen=True) +class _PreShardedHSDPLoadPlan: + roots: list[_HSDPLoadRoot] + groups: list[_HSDPModuleLoadGroup] + bindings: dict[str, TensorBinding] + binding_names: frozenset[str] + + class DiffusersPipelineLoader(HWRLoaderMixin): """Model loader that can load diffusers pipeline components from disk.""" @@ -1222,6 +1262,451 @@ class DiffusersPipelineLoader(HWRLoaderMixin): raise ValueError(f"Unknown load_format: {load_format}") return model + def _load_model_with_pre_sharded_hsdp( + self, + model: nn.Module, + discovered_modules: PipelineModules, + hsdp_config: HSDPInferenceConfig, + target_device: torch.device, + ) -> nn.Module: + """Shard meta DiT parameters before reading their rank-local checkpoint slices.""" + from torch.distributed.checkpoint import HuggingFaceStorageReader + from torch.distributed.checkpoint import load as dcp_load + + if self.counter_before_loading_weights == 0.0: + self.counter_before_loading_weights = time.perf_counter() + + load_plan = self._prepare_pre_sharded_hsdp_load( + model, + discovered_modules, + hsdp_config, + target_device, + ) + self._check_no_unplanned_meta_tensors(model, load_plan.binding_names) + self._validate_pre_sharded_hsdp_bindings(load_plan) + + self._process_pre_sharded_hsdp_groups_on_meta(load_plan) + + for group in load_plan.groups: + logger.debug("Applying meta-first HSDP shard to %s", group.name) + shard_hsdp_module(group.module, group.context) + for root in load_plan.roots: + logger.debug("Finalizing meta-first HSDP root %s", root.name) + finalize_hsdp_root(root.module, root.context) + + self._materialize_pre_sharded_hsdp_state(model, load_plan, target_device) + + targets = self._get_pre_sharded_hsdp_targets(model, load_plan) + checkpoint_states, local_payload_bytes = self._prepare_pre_sharded_checkpoint_states( + load_plan, + targets, + ) + + thread_count = int(getattr(self.od_config, "num_weight_load_threads", 4)) + if _is_distributed_rank_zero(): + directory_label = "directory" if len(checkpoint_states) == 1 else "directories" + logger.info( + "Pre-sharded HSDP loading %.2f GiB of rank-local tensor slices directly into " + "FSDP-owned storage from %d HF safetensors %s with %d reader thread(s) per rank", + local_payload_bytes / 1024**3, + len(checkpoint_states), + directory_label, + thread_count, + ) + + for checkpoint_dir in sorted(checkpoint_states, key=str): + reader = HuggingFaceStorageReader(str(checkpoint_dir), thread_count=thread_count) + # Each worker has the full checkpoint available and each DTensor + # already describes its rank-local destination slice. Avoid DCP's + # coordinator collectives, which can leave rank 0 with less GPU + # memory headroom than the other HSDP workers. + dcp_load(checkpoint_states[checkpoint_dir], storage_reader=reader, no_dist=True) + + remaining_meta = self._pre_sharded_meta_target_names(model, load_plan) + if remaining_meta: + raise ValueError( + "Pre-sharded HSDP checkpoint loading left tensors on meta; " + f"first missing tensors: {remaining_meta[:16]}" + ) + + self._finalize_pre_sharded_hsdp_loaded_weights(load_plan) + + self.counter_after_loading_weights = time.perf_counter() + logger.info_once( + "Loading weights took %.2f seconds", + self.counter_after_loading_weights - self.counter_before_loading_weights, + ) + self._move_non_hsdp_modules(discovered_modules, target_device) + return model + + @staticmethod + def _validate_pre_sharded_hsdp_bindings(load_plan: _PreShardedHSDPLoadPlan) -> None: + unsupported = [name for name, binding in load_plan.bindings.items() if binding.transform is not None] + if unsupported: + raise ValueError( + "Pre-sharded HSDP requires checkpoint tensors in runtime layout; " + f"tensor transforms are unsupported: {unsupported[:16]}" + ) + non_safetensors = sorted( + { + binding.file_path + for binding in load_plan.bindings.values() + if not binding.file_path.endswith(".safetensors") + } + ) + if non_safetensors: + raise ValueError( + f"Pre-sharded HSDP requires existing HF safetensors files; unsupported files: {non_safetensors[:8]}" + ) + + @staticmethod + def _get_pre_sharded_hsdp_targets( + model: nn.Module, + load_plan: _PreShardedHSDPLoadPlan, + ) -> dict[str, torch.Tensor]: + targets: dict[str, torch.Tensor] = { + **dict(model.named_buffers()), + **dict(model.named_parameters()), + } + missing_targets = sorted(load_plan.binding_names - targets.keys()) + if missing_targets: + raise ValueError( + f"Pre-sharded HSDP changed checkpoint target names; first missing targets: {missing_targets[:16]}" + ) + return targets + + @staticmethod + def _materialize_pre_sharded_hsdp_state( + model: nn.Module, + load_plan: _PreShardedHSDPLoadPlan, + target_device: torch.device, + ) -> None: + """Materialize the installed FSDP rank-local state before checkpoint I/O. + + ``fully_shard`` supports meta initialization through ``Module.to_empty``: + its registered ``_apply`` hooks update the internal sharded storage to + point at the newly allocated rank-local tensors. Modules are materialized + in the same HSDP-group order and child-first traversal used by ordinary + ``fully_shard`` initialization instead of the forward traversal used by + a recursive ``to_empty`` call. + PyTorch's swap-on-conversion mode preserves Parameter identity so that + post-load kernel objects created before sharding cannot retain stale + references to the original meta Parameters. + + A complete checkpoint plan covers every parameter and persistent buffer + in each DiT. Non-persistent buffers are not checkpoint entries, so keep + their initialized values across ``to_empty``. + """ + unplanned_parameters: list[str] = [] + saved_buffers: dict[str, torch.Tensor] = {} + for root in load_plan.roots: + prefix = f"{root.name}." + for local_name, _parameter in root.module.named_parameters(): + full_name = prefix + local_name + if full_name not in load_plan.binding_names: + unplanned_parameters.append(full_name) + for local_name, buffer in root.module.named_buffers(): + full_name = prefix + local_name + if full_name in load_plan.binding_names: + continue + if buffer.device.type == "meta": + raise ValueError(f"Pre-sharded HSDP cannot preserve an unplanned meta buffer: {full_name}") + saved_buffers[full_name] = buffer.detach().cpu().clone() + + if unplanned_parameters: + raise ValueError( + "Pre-sharded HSDP requires checkpoint coverage for every DiT parameter; " + f"first unplanned parameters: {unplanned_parameters[:16]}" + ) + + previous_swap_mode = torch.__future__.get_swap_module_params_on_conversion() + torch.__future__.set_swap_module_params_on_conversion(True) + try: + sharded_group_roots = {group.module for group in load_plan.groups} + materialized_modules: set[nn.Module] = set() + for group in load_plan.groups: + for module in DiffusersPipelineLoader._hsdp_managed_modules_post_order( + group.module, + nested_hsdp_roots=sharded_group_roots - {group.module}, + ): + if module not in materialized_modules: + module.to_empty(device=target_device, recurse=False) + materialized_modules.add(module) + for root in load_plan.roots: + for module in DiffusersPipelineLoader._hsdp_managed_modules_post_order( + root.module, + nested_hsdp_roots=sharded_group_roots, + ): + if module not in materialized_modules: + module.to_empty(device=target_device, recurse=False) + materialized_modules.add(module) + finally: + torch.__future__.set_swap_module_params_on_conversion(previous_swap_mode) + + if saved_buffers: + materialized_buffers = dict(model.named_buffers()) + with torch.no_grad(): + for name, value in saved_buffers.items(): + target = materialized_buffers[name] + target.copy_(value.to(device=target.device, dtype=target.dtype)) + + @staticmethod + def _hsdp_managed_modules_post_order( + root: nn.Module, + *, + nested_hsdp_roots: set[nn.Module], + ) -> list[nn.Module]: + """Mirror FSDP's child-first managed-module traversal. + + Nested FSDP roots own their own parameters and are materialized in their + corresponding load group, so the enclosing traversal must not descend + into them. + """ + modules: list[nn.Module] = [] + visited: set[nn.Module] = set() + + def visit(module: nn.Module) -> None: + if module in visited: + return + visited.add(module) + for child in module.children(): + if child not in nested_hsdp_roots: + visit(child) + modules.append(module) + + visit(root) + return modules + + @staticmethod + def _prepare_pre_sharded_checkpoint_states( + load_plan: _PreShardedHSDPLoadPlan, + targets: dict[str, torch.Tensor], + ) -> tuple[dict[Path, dict[str, torch.Tensor]], int]: + from torch.distributed.tensor import DTensor + + checkpoint_states: dict[Path, dict[str, torch.Tensor]] = {} + checkpoint_owners: dict[tuple[Path, str], str] = {} + local_payload_bytes = 0 + for name, binding in load_plan.bindings.items(): + target = targets[name] + + checkpoint_dir = Path(binding.file_path).parent + owner_key = (checkpoint_dir, binding.checkpoint_key) + if previous_owner := checkpoint_owners.get(owner_key): + raise ValueError( + "Pre-sharded HSDP cannot bind one checkpoint tensor to multiple runtime tensors: " + f"{binding.checkpoint_key!r} maps to both {previous_owner!r} and {name!r}" + ) + checkpoint_owners[owner_key] = name + + if target.device.type == "meta": + raise ValueError(f"Pre-sharded HSDP target {name!r} was not materialized") + + local_target = target.to_local() if isinstance(target, DTensor) else target + local_payload_bytes += local_target.numel() * local_target.element_size() + checkpoint_states.setdefault(checkpoint_dir, {})[binding.checkpoint_key] = target + return checkpoint_states, local_payload_bytes + + @staticmethod + def _pre_sharded_meta_target_names( + model: nn.Module, + load_plan: _PreShardedHSDPLoadPlan, + ) -> list[str]: + from torch.distributed.tensor import DTensor + + targets = { + **dict(model.named_buffers()), + **dict(model.named_parameters()), + } + return [ + name + for name in load_plan.binding_names + if (target := targets[name]).device.type == "meta" + or (isinstance(target, DTensor) and target.to_local().device.type == "meta") + ] + + @staticmethod + def _process_pre_sharded_hsdp_groups_on_meta(load_plan: _PreShardedHSDPLoadPlan) -> None: + modules_to_process: list[tuple[nn.Module, QuantizeMethodBase]] = [] + for group in load_plan.groups: + for _, module in group.module.named_modules(): + quant_method = getattr(module, "quant_method", None) + if quant_method is None or not isinstance(quant_method, QuantizeMethodBase): + continue + modules_to_process.append((module, quant_method)) + + if not modules_to_process: + return + for root in load_plan.roots: + if not getattr(root.module, "_hsdp_pre_sharded_meta_post_load", False): + raise ValueError( + f"Model {type(root.module).__name__} has not declared its post-load processing safe " + "for pre-sharded meta initialization" + ) + for module, quant_method in modules_to_process: + if type(quant_method) is not UnquantizedLinearMethod: + raise ValueError( + "Pre-sharded HSDP does not support the selected post-load processing; " + f"{type(module).__name__} uses {type(quant_method).__name__}" + ) + for module, quant_method in modules_to_process: + quant_method.process_weights_after_loading(module) + + def _prepare_pre_sharded_hsdp_load( + self, + model: nn.Module, + discovered_modules: PipelineModules, + hsdp_config: HSDPInferenceConfig, + target_device: torch.device, + ) -> _PreShardedHSDPLoadPlan: + outer_dit_names, outer_dits = discovered_modules.outermost_dits() + if not outer_dits: + raise ValueError("No DiT modules discovered for HSDP sharding") + + sources = self._get_weight_sources(model) + plan_result = build_checkpoint_binding_plan( + model, + dit_modules=tuple(zip(outer_dit_names, outer_dits)), + sources=sources, + model_path=str(getattr(self.od_config, "model", "")) or None, + tensor_parallel_size=int(getattr(self.parallel_config, "tensor_parallel_size", 1)), + online_quantization=False, + ) + if plan_result.plan is None: + raise ValueError(f"Pre-sharded HSDP checkpoint loading is incompatible: {plan_result.fallback_reason}") + plan = plan_result.plan + unsupported_sources = [source.prefix for source in sources if source.prefix not in plan.planned_source_prefixes] + if unsupported_sources: + raise ValueError( + "Pre-sharded HSDP checkpoint loading requires dedicated DiT sources; " + f"unsupported source prefixes: {unsupported_sources}" + ) + + roots: list[_HSDPLoadRoot] = [] + groups: list[_HSDPModuleLoadGroup] = [] + for outer_name, outer_dit in zip(outer_dit_names, outer_dits): + conditions = getattr(outer_dit, "_hsdp_shard_conditions", None) + if not conditions: + raise ValueError(f"Model {type(outer_dit).__name__} has no _hsdp_shard_conditions defined") + + context = prepare_hsdp_shard_context( + outer_dit, + hsdp_config, + target_device=target_device, + ) + roots.append(_HSDPLoadRoot(outer_name, outer_dit, context)) + for local_name, module in reversed(list(outer_dit.named_modules())): + if not any(condition(local_name, module) for condition in conditions): + continue + if not local_name: + raise ValueError("_hsdp_shard_conditions must not select the DiT root") + full_name = f"{outer_name}.{local_name}" if local_name else outer_name + groups.append(_HSDPModuleLoadGroup(full_name, module, context)) + + parameters = dict(model.named_parameters()) + targets: dict[str, torch.Tensor] = {**dict(model.named_buffers()), **parameters} + preserved_params = {param for root in roots for param in root.context.ignored_params or ()} + binding_names = frozenset(plan.bindings) + self._release_hsdp_checkpoint_targets_to_meta( + model, + targets, + parameters, + binding_names, + preserved_params, + ) + + return _PreShardedHSDPLoadPlan( + roots=roots, + groups=groups, + bindings=plan.bindings, + binding_names=binding_names, + ) + + @staticmethod + def _check_no_unplanned_meta_tensors(model: nn.Module, planned_names: frozenset[str]) -> None: + unplanned_meta = [ + name + for name, tensor in (*model.named_parameters(), *model.named_buffers()) + if tensor.device.type == "meta" and name not in planned_names + ] + if unplanned_meta: + raise ValueError( + "Pre-sharded HSDP initialized tensors on meta that are not covered by the " + f"checkpoint plan. First unplanned tensors: {unplanned_meta[:16]}" + ) + + @staticmethod + def _release_hsdp_checkpoint_targets_to_meta( + model: nn.Module, + targets: dict[str, torch.Tensor], + parameters: dict[str, nn.Parameter], + binding_names: frozenset[str], + preserved_params: set[nn.Parameter], + ) -> None: + released: dict[str, tuple[int, float]] = {} + for name in binding_names: + target = targets[name] + if target in preserved_params or target.device.type == "meta": + continue + logical_gib = target.numel() * target.element_size() / 1024**3 + device_type = target.device.type + count, gib = released.get(device_type, (0, 0.0)) + released[device_type] = (count + 1, gib + logical_gib) + if name in parameters: + replacement = DiffusersPipelineLoader._make_parameter_like( + target, + torch.empty_like(target, device="meta"), + ) + torch.utils.swap_tensors(target, replacement) + else: + replacement = torch.empty_like(target, device="meta") + torch.utils.swap_tensors(target, replacement) + if released: + logger.info("Released HSDP checkpoint-covered tensors to meta before pre-sharded load: %s", released) + + @staticmethod + def _make_parameter_like(template: nn.Parameter, tensor: torch.Tensor) -> nn.Parameter: + """Create a Parameter on new storage while preserving loader metadata.""" + if type(template) is nn.Parameter: + replacement = nn.Parameter(tensor, requires_grad=template.requires_grad) + replacement.__dict__.update(getattr(template, "__dict__", {})) + return replacement + + replacement = template.__class__.__new__(template.__class__, tensor) + replacement.__dict__.update(getattr(template, "__dict__", {})) + replacement.requires_grad_(template.requires_grad) + return replacement + + @staticmethod + def _finalize_pre_sharded_hsdp_loaded_weights( + load_plan: _PreShardedHSDPLoadPlan, + ) -> None: + """Run model post-load hooks for tensors materialized by pre-sharded loading.""" + for root in load_plan.roots: + post_load = getattr(root.module, "post_load_weights", None) + if callable(post_load): + post_load() + root.module.eval() + validate = getattr(root.module, "validate_loaded_weights", None) + if callable(validate): + prefix = f"{root.name}." + validate({name.removeprefix(prefix) for name in load_plan.binding_names if name.startswith(prefix)}) + + @staticmethod + def _move_non_hsdp_modules( + discovered_modules: PipelineModules, + target_device: torch.device, + ) -> None: + modules_to_move: list[nn.Module] = [] + if discovered_modules.vaes is not None: + modules_to_move.extend(discovered_modules.vaes) + if discovered_modules.encoders is not None: + modules_to_move.extend(discovered_modules.encoders) + if discovered_modules.resident_modules is not None: + modules_to_move.extend(discovered_modules.resident_modules) + for module in modules_to_move: + module.to(target_device) + def _load_model_with_hsdp( self, target_device: torch.device, @@ -1251,6 +1736,29 @@ class DiffusersPipelineLoader(HWRLoaderMixin): # mapping (QKV fusion, etc.). if load_format == "diffusers": raise ValueError("HSDP is not supported with the diffusers adapter load format") + strategy = getattr(self.od_config, "hsdp_weight_load_strategy", "full") + if strategy == "pre_sharded": + if load_format != "default": + raise ValueError( + f"hsdp_weight_load_strategy={strategy!r} currently supports only diffusion_load_format='default'" + ) + if self.quant_config is not None: + raise ValueError(f"hsdp_weight_load_strategy={strategy!r} does not support quantization yet") + if getattr(self.od_config, "lora_path", None): + raise ValueError(f"hsdp_weight_load_strategy={strategy!r} does not support LoRA yet") + elif strategy != "full": + raise ValueError(f"Unknown hsdp_weight_load_strategy: {strategy!r}") + + if strategy == "pre_sharded": + model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True) + discovered_modules = ModuleDiscovery.discover(model) + return self._load_model_with_pre_sharded_hsdp( + model, + discovered_modules, + hsdp_config, + target_device, + ) + # Same host-fallback bound as the ordinary path: with a quant config the # HSDP path initializes on the accelerator (hsdp_defer_to_cpu=False), so # over-wide NPU-unquantizable fallback weights would otherwise round-trip diff --git a/vllm_omni/diffusion/model_loader/host_weight_plan.py b/vllm_omni/diffusion/model_loader/host_weight_plan.py index f0ea8024d..3bb9f0c5a 100644 --- a/vllm_omni/diffusion/model_loader/host_weight_plan.py +++ b/vllm_omni/diffusion/model_loader/host_weight_plan.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Loader-owned host-weight plans shared with diffusion offload backends.""" +"""Loader-owned host-weight plans shared by direct-checkpoint consumers.""" from __future__ import annotations @@ -205,7 +205,7 @@ def _build_source_map( if not model_to_ckpt: raise _PlanIncompatibleError("no compatible safetensors entries were found in the loader's model sources") logger.info( - "Indexed %d runtime tensor names from %d checkpoint keys for DLO preflight", + "Indexed %d runtime tensor names from %d checkpoint keys for checkpoint preflight", len(model_to_ckpt), indexed_keys, ) @@ -313,28 +313,28 @@ def _validate_source_metadata( ) -def build_checkpoint_mmap_plan( +def build_checkpoint_binding_plan( pipeline: nn.Module, *, dit_modules: Sequence[tuple[str, nn.Module]], sources: Sequence[object], model_path: str | None, tensor_parallel_size: int, - use_hsdp: bool, online_quantization: bool, + use_pipeline_key_remap: bool = True, has_distilled_lora: bool = False, ) -> HostWeightPlanResult: - """Build a complete direct-checkpoint plan or return a fallback reason.""" + """Build a complete binding plan for the pre-sharded HSDP consumer.""" if tensor_parallel_size != 1: return HostWeightPlanResult(None, f"TP={tensor_parallel_size} requires the ordinary loader") - if use_hsdp: - return HostWeightPlanResult(None, "HSDP requires the ordinary loader") if online_quantization: return HostWeightPlanResult(None, "online quantization requires the ordinary loader") if has_distilled_lora: return HostWeightPlanResult(None, "distilled LoRA requires the ordinary loader") remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None) + if not callable(remap_fn) and use_pipeline_key_remap: + remap_fn = getattr(pipeline, "remap_checkpoint_key", None) if not callable(remap_fn): remap_fn = None adapter = get_direct_mmap_adapter(pipeline) @@ -387,10 +387,37 @@ def build_checkpoint_mmap_plan( ) +def build_checkpoint_mmap_plan( + pipeline: nn.Module, + *, + dit_modules: Sequence[tuple[str, nn.Module]], + sources: Sequence[object], + model_path: str | None, + tensor_parallel_size: int, + use_hsdp: bool, + online_quantization: bool, + has_distilled_lora: bool = False, +) -> HostWeightPlanResult: + """Build the DLO mmap plan without changing its HSDP eligibility.""" + if use_hsdp: + return HostWeightPlanResult(None, "HSDP requires the ordinary loader") + return build_checkpoint_binding_plan( + pipeline, + dit_modules=dit_modules, + sources=sources, + model_path=model_path, + tensor_parallel_size=tensor_parallel_size, + online_quantization=online_quantization, + use_pipeline_key_remap=False, + has_distilled_lora=has_distilled_lora, + ) + + __all__ = [ "HostWeightPlan", "HostWeightPlanResult", "TensorBinding", + "build_checkpoint_binding_plan", "build_checkpoint_mmap_plan", "has_online_quantization", ] diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index 7241082e8..b060ea420 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -38,6 +38,7 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available from vllm_omni.diffusion.layers.norm import RMSNorm as _VllmRMSNorm +from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta from vllm_omni.diffusion.offloader.config import ( OffloadStrategy, resolve_offload_strategy, @@ -1084,6 +1085,7 @@ class Cosmos3LanguageModel(nn.Module): rope_theta: float, mrope_section: list[int], quant_config: QuantizationConfig | None = None, + release_completed_blocks_to_meta: bool = False, prefix: str = "", ) -> None: super().__init__() @@ -1093,21 +1095,21 @@ class Cosmos3LanguageModel(nn.Module): rope_theta=rope_theta, mrope_section=mrope_section, ) - self.layers = nn.ModuleList( - [ - Cosmos3UndDecoderLayer( - hidden_size=hidden_size, - intermediate_size=intermediate_size, - num_attention_heads=num_attention_heads, - num_key_value_heads=num_key_value_heads, - head_dim=head_dim, - rms_norm_eps=rms_norm_eps, - quant_config=quant_config, - prefix=f"{prefix}.layers.{i}", - ) - for i in range(num_hidden_layers) - ] - ) + self.layers = nn.ModuleList() + for i in range(num_hidden_layers): + layer = Cosmos3UndDecoderLayer( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + head_dim=head_dim, + rms_norm_eps=rms_norm_eps, + quant_config=quant_config, + prefix=f"{prefix}.layers.{i}", + ) + if release_completed_blocks_to_meta: + release_module_parameters_to_meta(layer) + self.layers.append(layer) # TODO: Not used right now, will be used in the future for prompt upsampler. self.norm = RMSNorm(hidden_size, eps=rms_norm_eps) @@ -1228,6 +1230,10 @@ class Cosmos3VFMTransformer(nn.Module): _hsdp_shard_conditions = [_is_transformer_block] + # Standard unquantized vLLM linears need no value-dependent post-load + # processing on accelerators. Edge inherits the same loading contract. + _hsdp_pre_sharded_meta_post_load = True + # Modules whose parameters must NOT be FSDP-sharded at the root level. # time_embedder is cast to fp32 by post_load_weights for precision; if it # were swept into the root flat-parameter under MixedPrecisionPolicy(param_dtype=bf16), @@ -1346,6 +1352,10 @@ class Cosmos3VFMTransformer(nn.Module): language_model_quant_config, gen_layers_quant_config = _resolve_cosmos3_quant_configs( getattr(od_config, "quantization_config", None) ) + release_completed_blocks_to_meta = bool( + getattr(getattr(od_config, "parallel_config", None), "use_hsdp", False) + and getattr(od_config, "hsdp_weight_load_strategy", "full") == "pre_sharded" + ) mixed_precision_config, mixed_precision_source = resolve_mixed_precision_config(od_config) if mixed_precision_config is None: if mixed_precision_source == "additional_config_disabled": @@ -1372,6 +1382,7 @@ class Cosmos3VFMTransformer(nn.Module): rope_theta=self.rope_theta, mrope_section=self.mrope_section, quant_config=language_model_quant_config, + release_completed_blocks_to_meta=release_completed_blocks_to_meta, prefix="language_model", **self._language_model_kwargs(), ) @@ -1399,24 +1410,24 @@ class Cosmos3VFMTransformer(nn.Module): self.audio_proj_out = nn.Linear(self.hidden_size, self.sound_dim) self.audio_modality_embed = nn.Parameter(torch.zeros(self.hidden_size)) - self.gen_layers = nn.ModuleList( - [ - Cosmos3GenDecoderLayer( - layer_idx=i, - hidden_size=self.hidden_size, - intermediate_size=self.intermediate_size, - num_attention_heads=self.num_attention_heads, - num_key_value_heads=self.num_key_value_heads, - head_dim=self.head_dim, - rms_norm_eps=self.rms_norm_eps, - quant_config=gen_layers_quant_config, - mlp_cls=self._gen_mlp_cls, - qk_norm=self.qk_norm_for_diffusion, - prefix=f"gen_layers.{i}", - ) - for i in range(self.num_hidden_layers) - ] - ) + self.gen_layers = nn.ModuleList() + for i in range(self.num_hidden_layers): + layer = Cosmos3GenDecoderLayer( + layer_idx=i, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + num_attention_heads=self.num_attention_heads, + num_key_value_heads=self.num_key_value_heads, + head_dim=self.head_dim, + rms_norm_eps=self.rms_norm_eps, + quant_config=gen_layers_quant_config, + mlp_cls=self._gen_mlp_cls, + qk_norm=self.qk_norm_for_diffusion, + prefix=f"gen_layers.{i}", + ) + if release_completed_blocks_to_meta: + release_module_parameters_to_meta(layer) + self.gen_layers.append(layer) self.mixed_precision_runtime: Cosmos3MixedPrecisionRuntime | None = None if mixed_precision_config is not None: diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py index c9a7d856e..fcc19123b 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py @@ -14,6 +14,7 @@ from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelL from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm_omni.diffusion.attention.layer import Attention as FrameworkAttention +from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta from .transformer_cosmos3 import ( Cosmos3VFMTransformer, @@ -258,6 +259,7 @@ class Cosmos3EdgeLanguageModel(nn.Module): mrope_section: list[int], use_und_k_norm_for_gen: bool, quant_config: QuantizationConfig | None = None, + release_completed_blocks_to_meta: bool = False, prefix: str = "", ) -> None: super().__init__() @@ -267,22 +269,22 @@ class Cosmos3EdgeLanguageModel(nn.Module): rope_theta=rope_theta, mrope_section=mrope_section, ) - self.layers = nn.ModuleList( - [ - Cosmos3EdgeUndDecoderLayer( - hidden_size=hidden_size, - intermediate_size=intermediate_size, - num_attention_heads=num_attention_heads, - num_key_value_heads=num_key_value_heads, - head_dim=head_dim, - rms_norm_eps=rms_norm_eps, - use_und_k_norm_for_gen=use_und_k_norm_for_gen, - quant_config=quant_config, - prefix=f"{prefix}.layers.{i}", - ) - for i in range(num_hidden_layers) - ] - ) + self.layers = nn.ModuleList() + for i in range(num_hidden_layers): + layer = Cosmos3EdgeUndDecoderLayer( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + head_dim=head_dim, + rms_norm_eps=rms_norm_eps, + use_und_k_norm_for_gen=use_und_k_norm_for_gen, + quant_config=quant_config, + prefix=f"{prefix}.layers.{i}", + ) + if release_completed_blocks_to_meta: + release_module_parameters_to_meta(layer) + self.layers.append(layer) self.norm = RMSNorm(hidden_size, eps=rms_norm_eps) def forward( diff --git a/vllm_omni/diffusion/models/utils.py b/vllm_omni/diffusion/models/utils.py index f55996a76..05a6a0a31 100644 --- a/vllm_omni/diffusion/models/utils.py +++ b/vllm_omni/diffusion/models/utils.py @@ -168,6 +168,28 @@ def init_parameters( init_parameters(child, dtype, device) +def release_module_parameters_to_meta(module: nn.Module) -> None: + """Release a fully constructed module's parameter storage to meta. + + The module must finish construction first so initialization code and vLLM + weight-loader metadata follow the ordinary CPU path. ``swap_tensors`` keeps + each registered Parameter object's identity while dropping only its storage. + Persistent buffers intentionally remain materialized. + """ + for param in module.parameters(): + if param.device.type == "meta": + continue + + meta_tensor = torch.empty_like(param, device="meta") + if type(param) is nn.Parameter: + meta_param = nn.Parameter(meta_tensor, requires_grad=param.requires_grad) + else: + meta_param = param.__class__.__new__(param.__class__, meta_tensor) + meta_param.requires_grad_(param.requires_grad) + meta_param.__dict__.update(getattr(param, "__dict__", {})) + torch.utils.swap_tensors(param, meta_param) + + def init_parameters_preserving_vllm_linear( module: nn.Module, dtype: torch.dtype | None, diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index b9ed35de0..458a6237d 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -558,6 +558,7 @@ class OrchestratorArgs: # === Diffusion model config === num_gpus: int | None = None model_class_name: str | None = None + hsdp_weight_load_strategy: str | None = None diffusion_load_format: str | None = None lora_path: list[str] | None = None lora_backend: str | None = None diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index ac1c3cda1..9c3aeb5b5 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -560,6 +560,15 @@ class OmniServeCommand(CLISubcommand): "custom_pipeline, dummy, or diffusers for the HF diffusers adapter." ), ) + omni_config_group.add_argument( + "--hsdp-weight-load-strategy", + choices=["full", "pre_sharded"], + default=None, + help=( + "HSDP checkpoint loading strategy: 'full' loads then shards; " + "'pre_sharded' shards meta parameters first and reads rank-local HF safetensors slices." + ), + ) omni_config_group.add_argument( "--lora-path", type=str, From 596487fe62bad27df5a9257e7b5dbdeb6e22d51e Mon Sep 17 00:00:00 2001 From: Rahul Steiger Date: Fri, 2 Oct 2026 11:58:47 +0200 Subject: [PATCH 10/22] [Core] Prototype diffusion attention execution contracts with FA4 (#7379) --- .../feature/attention_backend_selection.md | 3 + .../attention_execution_contract_poc.md | 133 +++++++ docs/design/index.md | 1 + .../attention/test_attention_capabilities.py | 359 ++++++++++++++++++ .../test_attention_state_lifetime.py | 118 ++++++ tests/diffusion/attention/test_flash_attn.py | 223 +++++++++++ .../attention/test_flash_attn_compile.py | 308 +++++++++++++++ tests/diffusion/test_attention_compile.py | 10 +- .../diffusion/attention/backends/abstract.py | 31 ++ .../attention/backends/flash_attn.py | 300 ++++++++++++++- .../diffusion/attention/backends/utils/fa.py | 25 ++ vllm_omni/diffusion/attention/capabilities.py | 126 ++++++ vllm_omni/diffusion/attention/layer.py | 93 ++++- 13 files changed, 1703 insertions(+), 27 deletions(-) create mode 100644 docs/design/feature/attention_execution_contract_poc.md create mode 100644 tests/diffusion/attention/test_attention_capabilities.py create mode 100644 tests/diffusion/attention/test_attention_state_lifetime.py create mode 100644 tests/diffusion/attention/test_flash_attn_compile.py create mode 100644 vllm_omni/diffusion/attention/capabilities.py diff --git a/docs/design/feature/attention_backend_selection.md b/docs/design/feature/attention_backend_selection.md index 3ae1115b1..3e62e2cd9 100644 --- a/docs/design/feature/attention_backend_selection.md +++ b/docs/design/feature/attention_backend_selection.md @@ -4,6 +4,9 @@ This document defines the selection and extension contract for diffusion attention backends. User-facing backend choices, installation, and tuning are in the [attention backend guides](../../user_guide/diffusion/attention_backends.md). +The [execution contract PoC](attention_execution_contract_poc.md) describes the +experimental post-selection capability and lifecycle contract from RFC #7226. + ## Scope The contract applies to DiT and other diffusion attention layers. It is diff --git a/docs/design/feature/attention_execution_contract_poc.md b/docs/design/feature/attention_execution_contract_poc.md new file mode 100644 index 000000000..8c7486239 --- /dev/null +++ b/docs/design/feature/attention_execution_contract_poc.md @@ -0,0 +1,133 @@ +# Diffusion attention execution contract PoC + +This prototype accompanies [RFC #7226](https://github.com/vllm-project/vllm-omni/issues/7226) +and extends the [attention selection design](attention_backend_selection.md). +It demonstrates path-specific capabilities with dense BF16 FA4, NPU/ROCm +routing examples, and tensor-state lifetime with a test-only attention module. + +## Execution contract + +`ExecutionContext` describes the requested execution path. `ExecutionPathResult` +reports its identity, support status and reason, and compilation mode. +`AttentionBackend.resolve_capabilities()` provides conservative pre-construction +results. After initialization, `Attention.resolve_execution_path()` supplies the +active parallel, paged-KV, and HSDP context; the backend combines it with the +selected kernel and normalized metadata. + +Resolution runs outside compiled execution. FA4 reports `SUPPORTED` and +`CUSTOM_OP` for dense, noncausal BF16 without parallel or HSDP boundaries when +its kernel accepts the head dimensions. Dimension validation delegates to FA4's +architecture-specific rules through `backends/utils/fa.py`. Kernel rejections +report `UNSUPPORTED` with an actionable reason. Missing private validators and +other unmigrated paths report `UNMIGRATED` with advisory `EAGER_ONLY` defaults. +`requested_support()` checks a fullgraph request; it does not enforce selection +or change existing execution. + +FA4 execution uses an opaque custom op. Its fake output preserves Q's batch, +sequence, and head count and uses V's head dimension. Output is contiguous even +when Q is noncontiguous. Metadata normalization is +shared with dispatch. Producers must update published mask semantics when masks +change; unpublished masks remain runtime-dependent to avoid synchronization. +Callers resolve again when execution metadata changes. + +## NPU and ROCm worked examples + +The same `FLASH_ATTN` backend resolves MindIE paths on NPU and AITER paths on +ROCm. Both examples report `UNMIGRATED`: their route is inspected, while hardware +correctness and compilation remain unvalidated. `EAGER_ONLY` is an advisory +default, not a measured compiler limitation. Existing execution is unchanged. + +| Input | NPU / MindIE | ROCm / AITER | +| --- | --- | --- | +| Dense | `npu_dense` | `rocm_dense` | +| Padding mask | `npu_masked`: expand the key mask to `[B, 1, Q, K]` | `rocm_masked_varlen`: unpad, call varlen attention, then restore rows | +| Packed inputs | Opt-in `[real, pad]` contract; choose varlen or prefix-KV slicing from the environment | Forward complete packed-document boundaries to varlen attention | +| Incomplete packing | Rebuild a mask when possible; otherwise reject before the kernel | Shared CUDA-like normalization rejects incomplete metadata | +| Unpublished mask semantics | Route to the masked wrapper | Keep resolution runtime-dependent without reading mask values | + +NPU resolution reuses `_resolve_packed_seq_npu`; ROCm reuses the metadata +normalization used by dispatch. AITER identity comes from initialization rather +than caller-provided context. Quantized, paged, piecewise, and parallel paths +remain outside these worked examples. + +The new L1 tests use CPU tensors and replace only vendor kernel calls. They +check path identity, forwarded masks/boundaries, fallback rejection, and +conservative fullgraph-request handling. They do not test vendor numerics or +compiler support. Existing L1 CI already collects this test file. + +```bash +python -m pytest tests/diffusion/attention/test_flash_attn.py \ + -k 'npu_contract or rocm_contract' -m 'core_model and cpu' -q +``` + +### Hardware validation required before migration + +On an Ascend/MindIE or ROCm/AITER runner, record device, driver, PyTorch, vendor +library, and compiler versions, then validate the corresponding rows above: + +1. Compare eager outputs with FP32 SDPA for dense and padding-mask inputs, + including unequal Q/K lengths. Compare valid query rows; do not assume both + platforms define padded-query outputs identically. +2. Compare packed outputs with separate per-document references. On NPU, cover + both varlen and laser prefix-slicing modes, plus mask reconstruction and the + missing-mask error. A padding fallback does not establish multi-document + isolation. On ROCm, cover multiple real documents. +3. Check output shape, dtype, device, finite values, and unchanged inputs. + Start BF16 comparisons at `atol=rtol=1e-2` and have the platform maintainer + confirm the tolerance against its reference tests. +4. Run `torch.compile(fullgraph=True, dynamic=True)` with the platform's supported + compiler over multiple sequence lengths. Record graph breaks and recompiles; + choose `TRACEABLE`, `CUSTOM_OP`, or verified `EAGER_ONLY` from that evidence. + A custom-op path additionally needs schema/fake checks and repeated replay. + +Promote only the paths validated on that platform. These hardware checks have +not been run locally; ROCm/NPU fullgraph support is not claimed. + +## State lifetime example + +`test_attention_state_lifetime.py` prepares an owned query-scale tensor eagerly +and passes it explicitly to an opaque attention op. The compiled module retains +its state after the caller drops its reference and releases ownership when the +compiled callable is deleted. Equivalent instances reuse a graph with their own +state values. + +This example tests ownership and release, dynamic replay across instances, +custom-op schema/fake behavior, and Inductor execution. It adds no production +backend or state registry. Planning keys, caching, fallback-policy declarations, +opaque handles, CUDA graph capture, and TRTLLM integration are deferred. + +## Validation + +```bash +python -m pytest tests/diffusion/attention/test_flash_attn_compile.py \ + tests/diffusion/attention/test_flash_attn.py \ + tests/diffusion/attention/test_attention_capabilities.py \ + tests/diffusion/attention/test_attention_state_lifetime.py -q -rs +``` + +The focused suite covers capability decisions and metadata changes, existing +padding/mask regressions, tensor-state lifetime, and real FA4 fullgraph execution +against FP32 SDPA and eager FA4. Representative Q/K and V dimensions are +(32, 32), (64, 64), (80, 48), (192, 128), and (256, 256), with batches 1 and 2 +and sequence lengths up to 1,024. Invalid dimensions are checked against actual +kernel errors. These samples exercise the contract rather than define support. +BF16 comparisons use `atol=rtol=1e-2`; `torch.library.opcheck` checks schema and +fake-output correctness. + +Single-graph reuse is asserted in the tensor-state test. The FA4 numerical tests +retain compiler caches across shapes within each case but do not require one +graph across different batch sizes, Q/K length equality, or head dimensions. + +Tested environment: GB300, PyTorch `2.13.0+cu130`, FA4 `4.0.0b18`, CUTLASS DSL +`4.6.2` with CUDA 13 libraries, Quack `0.6.4`, and TVM FFI `0.1.11`. +Real-kernel tests require CUDA and CuTe FA4; their availability is checked at runtime. + +Before the rebase, the focused suite passed all 97 tests without skips, including +11 CPU NPU/ROCm routing tests. After rebasing onto upstream `e3be42e05`, the +expanded suite reports 99 passed and 11 failed in the local vLLM 0.28 environment. +All 11 failures are import errors: upstream now requires `compute_layout_strides` +from vLLM, which this environment lacks. Both noncontiguous FA4 regression cases +and real-kernel schema/fake checks pass. The full suite must be rerun with the +vLLM version required by upstream. NPU/ROCm device numerics and compilation remain +unvalidated. Pre-commit passes with the CI hook skips. Upstream PyTorch/CUTLASS +warnings remain. diff --git a/docs/design/index.md b/docs/design/index.md index af8dbf438..64144a5df 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -68,6 +68,7 @@ guides list selectable backends, platform defaults, installation, and tuning. The design contracts separate selection mechanics from backend algorithms: - [Attention Backend Selection](feature/attention_backend_selection.md) +- [Attention Execution Contract PoC](feature/attention_execution_contract_poc.md) - [Skip-Softmax](feature/skip_softmax.md) #### CPU offloading diff --git a/tests/diffusion/attention/test_attention_capabilities.py b/tests/diffusion/attention/test_attention_capabilities.py new file mode 100644 index 000000000..a3d649aaa --- /dev/null +++ b/tests/diffusion/attention/test_attention_capabilities.py @@ -0,0 +1,359 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import sys +from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from vllm_omni.diffusion.attention.backends.abstract import ( + AttentionBackend, + AttentionMetadata, +) +from vllm_omni.diffusion.attention.backends.flash_attn import ( + FlashAttentionBackend, + FlashAttentionImpl, +) +from vllm_omni.diffusion.attention.backends.utils import fa +from vllm_omni.diffusion.attention.capabilities import ( + CapabilityResult, + CompilationMode, + ExecutionContext, + OuterBoundary, + ParallelStrategy, + SupportStatus, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + +_validate_fa4_head_dims = fa.validate_fa4_head_dims + + +@pytest.fixture(autouse=True) +def kernel_validator(monkeypatch): + # Capability unit tests do not require FA4 or a CUDA device. Real-kernel + # validation is exercised by test_flash_attn_compile.py. + monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: True) + + +def _impl(*, kernel_variant: str | None = "fa4", causal: bool = False): + impl = FlashAttentionImpl.__new__(FlashAttentionImpl) + impl._kernel_variant = kernel_variant + impl.causal = causal + return impl + + +def _resolve( + impl, + *, + context: ExecutionContext | None = None, + attn_metadata: AttentionMetadata | None = None, + dtype=torch.bfloat16, +): + tensor = torch.empty((1, 16, 8, 64), dtype=dtype) + return impl.resolve_execution_path( + context or ExecutionContext(platform="cuda"), + tensor, + tensor, + tensor, + attn_metadata, + ) + + +class _UnmigratedBackend(AttentionBackend): + @staticmethod + def get_name() -> str: + return "UNMIGRATED_TEST" + + @staticmethod + def get_impl_cls(): + return None + + @staticmethod + def get_metadata_cls(): + return None + + @staticmethod + def get_builder_cls(): + return None + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [] + + +def test_capability_result_is_tri_state(): + assert CapabilityResult.supported().status is SupportStatus.SUPPORTED + assert CapabilityResult.unsupported("not available").status is SupportStatus.UNSUPPORTED + assert CapabilityResult.unmigrated().status is SupportStatus.UNMIGRATED + with pytest.raises(ValueError, match="actionable reason"): + CapabilityResult(SupportStatus.UNSUPPORTED) + + +def test_unmigrated_backend_default_is_conservative(): + context = ExecutionContext(platform="cuda", require_fullgraph=True) + result = _UnmigratedBackend.resolve_capabilities(context) + + assert result.support.status is SupportStatus.UNMIGRATED + assert result.compilation_mode is CompilationMode.EAGER_ONLY + requested = result.requested_support(context) + assert requested.status is SupportStatus.UNSUPPORTED + assert "no verified fullgraph declaration" in requested.reason + assert result.support.status is SupportStatus.UNMIGRATED + assert result.requested_support(replace(context, require_fullgraph=False)).status is SupportStatus.UNMIGRATED + + +@pytest.mark.parametrize("mode", list(CompilationMode)) +def test_fullgraph_request_checks_compilation_mode(mode): + result = replace(_resolve(_impl()), compilation_mode=mode) + context = ExecutionContext(platform="cuda", require_fullgraph=True) + expected = SupportStatus.UNSUPPORTED if mode is CompilationMode.EAGER_ONLY else SupportStatus.SUPPORTED + assert result.requested_support(context).status is expected + assert result.requested_support(replace(context, require_fullgraph=False)).status is SupportStatus.SUPPORTED + + +def test_request_preserves_verified_rejection_reason(): + result = replace(_resolve(_impl()), support=CapabilityResult.unsupported("unsupported head size")) + assert result.requested_support(ExecutionContext(platform="cuda", require_fullgraph=True)) == result.support + + +def test_exact_cuda_fa4_dense_path_is_supported(): + context = ExecutionContext( + platform="cuda", + require_fullgraph=True, + ) + result = _resolve(_impl(), context=context) + + assert result.path == "fa4_dense" + assert result.support.status is SupportStatus.SUPPORTED + assert result.compilation_mode is CompilationMode.CUSTOM_OP + assert result.requested_support(context).status is SupportStatus.SUPPORTED + + +def test_causal_fa4_path_remains_unmigrated(): + result = _resolve(_impl(causal=True)) + + assert result.support.status is SupportStatus.UNMIGRATED + + +def test_fa4_delegates_dimensions_to_kernel(monkeypatch): + head_dims = (80, 48) + validator = Mock(return_value=True) + monkeypatch.setattr(fa, "validate_fa4_head_dims", validator) + query = torch.empty((1, 16, 8, head_dims[0]), dtype=torch.bfloat16) + value = torch.empty((1, 16, 8, head_dims[1]), dtype=torch.bfloat16) + result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), query, query, value, None) + assert result.support.status is SupportStatus.SUPPORTED + validator.assert_called_once_with(*head_dims, 8) + + +def test_fa4_reports_kernel_rejection(monkeypatch): + monkeypatch.setattr(fa, "validate_fa4_head_dims", Mock(side_effect=AssertionError("kernel dimension constraint"))) + result = _resolve(_impl()) + assert result.support.status is SupportStatus.UNSUPPORTED + assert "kernel dimension constraint" in result.support.reason + assert "another backend" in result.support.reason + + +def test_fa4_missing_validator_remains_unmigrated(monkeypatch): + monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: False) + result = _resolve(_impl()) + assert result.support.status is SupportStatus.UNMIGRATED + assert result.compilation_mode is CompilationMode.EAGER_ONLY + + +def test_fa4_rejects_mismatched_query_key_dimensions(): + query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16) + key = torch.empty((1, 16, 8, 32), dtype=torch.bfloat16) + result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), query, key, key, None) + assert result.support.status is SupportStatus.UNSUPPORTED + assert "Q and K head dimensions must match" in result.support.reason + + +@pytest.mark.parametrize("arch", [90, 100, 120]) +def test_fa4_validator_adapter_uses_selected_architecture(monkeypatch, arch): + validator = Mock() + monkeypatch.setitem( + sys.modules, + "flash_attn.cute.interface", + SimpleNamespace(_get_device_arch=lambda: arch, _validate_head_dims=validator), + ) + # Call the real adapter, not the fixture's stand-in. + verified = _validate_fa4_head_dims(80, 48, 8) + if arch == 120: + assert not verified + validator.assert_not_called() + else: + assert verified + validator.assert_called_once_with(80, 48, arch // 10, 8) + + +def test_fa4_validator_adapter_handles_missing_private_api(monkeypatch): + monkeypatch.setitem(sys.modules, "flash_attn.cute.interface", SimpleNamespace()) + assert not _validate_fa4_head_dims(64, 64, 8) + + +def test_initialized_kernel_identity_overrides_caller_claim(): + result = _resolve( + _impl(kernel_variant=None), + context=ExecutionContext( + platform="cuda", + kernel_variant="fa4", + ), + ) + + assert result.support.status is SupportStatus.UNMIGRATED + + +def test_preconstruction_without_resolved_cuda_kernel_is_unmigrated(): + result = FlashAttentionBackend.resolve_capabilities(ExecutionContext(platform="cuda")) + + assert result.support.status is SupportStatus.UNMIGRATED + assert result.path == "unverified" + + +@pytest.mark.parametrize( + "changes", + [ + {"platform": "rocm"}, + {"paged_kv": True}, + {"parallel_strategy": ParallelStrategy.ULYSSES}, + {"parallel_strategy": ParallelStrategy.RING}, + {"parallel_strategy": ParallelStrategy.HYBRID_ULYSSES_RING}, + {"outer_boundaries": frozenset({OuterBoundary.HSDP})}, + ], +) +def test_unverified_flash_attention_variants_remain_unmigrated(changes): + result = _resolve( + _impl(), + context=replace(ExecutionContext(platform="cuda"), **changes), + ) + + assert result.support.status is SupportStatus.UNMIGRATED + assert result.compilation_mode is CompilationMode.EAGER_ONLY + + +def test_unverified_dtype_remains_unmigrated(): + result = _resolve(_impl(), dtype=torch.float32) + + assert result.support.status is SupportStatus.UNMIGRATED + + +@pytest.mark.parametrize( + "attn_metadata", + [ + AttentionMetadata(full_attn_spans=[[(0, 8)]]), + AttentionMetadata( + extra={ + "cu_seqlens_q": torch.tensor([0, 16], dtype=torch.int32), + "cu_seqlens_k": torch.tensor([0, 16], dtype=torch.int32), + "max_seqlen_q": 16, + "max_seqlen_k": 16, + } + ), + AttentionMetadata(extra={"kv_cache_dtype": "fp8"}), + ], +) +def test_metadata_normalization_prevents_dense_support(attn_metadata): + result = _resolve(_impl(), attn_metadata=attn_metadata) + + assert result.support.status is SupportStatus.UNMIGRATED + + +def test_unpublished_all_true_mask_remains_runtime_dependent(): + metadata = AttentionMetadata( + attn_mask=torch.ones((1, 16), dtype=torch.bool), + ) + result = _resolve(_impl(), attn_metadata=metadata) + + assert result.path == "runtime_mask_dependent" + assert result.support.status is SupportStatus.UNMIGRATED + + +def test_producer_published_noop_mask_resolves_dense_fa4(): + metadata = AttentionMetadata( + attn_mask=torch.ones((1, 16), dtype=torch.bool), + extra={"attention_mask_mode": "none"}, + ) + result = _resolve(_impl(), attn_metadata=metadata) + + assert result.path == "fa4_dense" + assert result.support.status is SupportStatus.SUPPORTED + + +def test_fixed_shape_mask_semantic_change_invalidates_dense_resolution(): + metadata = AttentionMetadata( + attn_mask=torch.ones((1, 16), dtype=torch.bool), + extra={"attention_mask_mode": "none"}, + ) + impl = _impl() + dense = _resolve(impl, attn_metadata=metadata) + assert dense.support.status is SupportStatus.SUPPORTED + + metadata.attn_mask[:, 8:] = False + metadata.extra["attention_mask_mode"] = "padding" + masked = _resolve(impl, attn_metadata=metadata) + assert masked.support.status is SupportStatus.UNMIGRATED + request = ExecutionContext(platform="cuda", require_fullgraph=True) + assert masked.requested_support(request).status is SupportStatus.UNSUPPORTED + + +def test_piecewise_dispatch_ignores_unused_incomplete_packed_metadata(monkeypatch): + from vllm_omni.diffusion.attention.backends import flash_attn + from vllm_omni.diffusion.attention.backends.utils import fa + + metadata = AttentionMetadata(full_attn_spans=[[(0, 16)]], extra={"max_seqlen_q": 16}) + impl = _impl() + impl.fa_deterministic = False + impl.softmax_scale = 0.125 + assert _resolve(impl, attn_metadata=metadata).support.status is SupportStatus.UNMIGRATED + + query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16) + monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True) + monkeypatch.setattr(fa, "flash_attn_func", lambda *args, **kwargs: query) + monkeypatch.setattr(flash_attn, "piecewise_attn", lambda *args, **kwargs: query) + assert impl.forward_cuda(query, query, query, metadata) is query + + +@pytest.mark.parametrize("input_index", [1, 2]) +@pytest.mark.parametrize("mismatch", ["dtype", "device"]) +def test_fa4_rejects_incompatible_key_value(input_index, mismatch): + tensors = [torch.empty((1, 16, 8, 64), dtype=torch.bfloat16) for _ in range(3)] + tensors[input_index] = ( + tensors[input_index].to(dtype=torch.float16) if mismatch == "dtype" else tensors[input_index].to("meta") + ) + result = _impl().resolve_execution_path(ExecutionContext(platform="cuda"), *tensors, None) + assert result.support.status is SupportStatus.UNSUPPORTED + assert f"{mismatch}s must match" in result.support.reason + + +@pytest.mark.parametrize("skip_quant", [False, True]) +def test_layer_resolution_uses_effective_kv_quantization(skip_quant): + from vllm_omni.diffusion.attention.layer import Attention + from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention + + layer = Attention.__new__(Attention) + torch.nn.Module.__init__(layer) + layer.attention = _impl(kernel_variant=None) + layer._hsdp_compile_boundary_enabled = False + layer.skip_sequence_parallel = False + layer._no_parallel_strategy = NoParallelAttention() + layer.parallel_strategy = layer._no_parallel_strategy + layer.use_ring = False + layer.paged_kv_cache_role = None + layer._kv_cache_dtype = "fp8" + layer._disable_kv_quant = False + layer._kv_cache_skip_steps = None + layer._kv_cache_skip_layers = {0} if skip_quant else None + layer.layer_idx = 0 + query = torch.empty((1, 16, 8, 64), dtype=torch.bfloat16) + metadata = AttentionMetadata(extra={"kv_cache_dtype": "fp8"}) + result = layer.resolve_execution_path(ExecutionContext(platform="npu"), query, query, query, metadata) + assert result.path == ("npu_dense" if skip_quant else "npu_unverified") + assert result.support.status is SupportStatus.UNMIGRATED + assert metadata.extra == {"kv_cache_dtype": "fp8"} diff --git a/tests/diffusion/attention/test_attention_state_lifetime.py b/tests/diffusion/attention/test_attention_state_lifetime.py new file mode 100644 index 000000000..a85fa5c93 --- /dev/null +++ b/tests/diffusion/attention/test_attention_state_lifetime.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Stateful contract example, independent of an optional attention kernel. + +The prepared module owns tensor state and passes it explicitly to an opaque op. +This deliberately does not model opaque Python/C++ planning handles. +""" + +import gc +import weakref + +import pytest +import torch +import torch.nn.functional as F + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +@torch.library.custom_op("vllm_omni_test::planned_attention", mutates_args=()) +def _planned_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + query_scale: torch.Tensor, +) -> torch.Tensor: + return F.scaled_dot_product_attention(query * query_scale, key, value).contiguous() + + +@_planned_attention.register_fake +def _planned_attention_fake(query, key, value, query_scale): + return query.new_empty((*query.shape[:-1], value.shape[-1])) + + +class _PreparedAttention(torch.nn.Module): + """Initialization/planning runs eagerly; forward only executes tensors.""" + + def __init__(self, query_scale: torch.Tensor): + super().__init__() + self.register_buffer("query_scale", query_scale.detach().clone()) + + def forward(self, query, key, value): + return _planned_attention(query, key, value, self.query_scale) + + +def _inputs(length): + generator = torch.Generator().manual_seed(length) + return tuple(torch.randn(2, 3, length, 8, generator=generator) for _ in range(3)) + + +def _reference(plan, inputs): + query, key, value = inputs + return F.scaled_dot_product_attention(query * plan.query_scale, key, value) + + +def test_dynamic_replay_uses_each_instances_state_without_recompilation(): + compile_count = 0 + + def counting_backend(graph_module, _example_inputs): + nonlocal compile_count + compile_count += 1 + return graph_module.forward + + # Same tensor metadata, different values: state must be a graph input, + # never a Python instance identifier or a captured first-instance constant. + plans = [_PreparedAttention(torch.full((1, 3, 1, 1), scale)) for scale in (0.5, 2.0)] + compiled = torch.compile( + lambda plan, q, k, v: plan(q, k, v), + backend=counting_backend, + fullgraph=True, + dynamic=True, + ) + for length in (5, 9, 13): + inputs = _inputs(length) + for plan in plans: + torch.testing.assert_close(compiled(plan, *inputs), _reference(plan, inputs)) + assert not torch.allclose(_reference(plans[0], inputs), _reference(plans[1], inputs)) + assert compile_count == 1 + + +def test_compiled_callable_retains_state_until_released(): + plan = _PreparedAttention(torch.full((1, 3, 1, 1), 0.5)) + plan_ref = weakref.ref(plan) + state_ref = weakref.ref(plan.query_scale) + inputs = _inputs(5) + expected = _reference(plan, inputs) + compiled = torch.compile(plan, backend="eager", fullgraph=True, dynamic=True) + torch.testing.assert_close(compiled(*inputs), expected) + + # Deleting the caller's reference must not invalidate subsequent replay. + del plan + gc.collect() + assert plan_ref() is not None + assert state_ref() is not None + torch.testing.assert_close(compiled(*inputs), expected) + + del compiled + gc.collect() + assert plan_ref() is None + assert state_ref() is None + + +def test_stateful_custom_op_schema_and_fake_implementation(): + inputs = _inputs(5) + # Unequal Q/K lengths and V head dimensions exercise the fake output shape. + query, _, _ = inputs + key = torch.randn(2, 3, 7, 8) + value = torch.randn(2, 3, 7, 4) + scale = torch.full((1, 3, 1, 1), 0.5) + torch.library.opcheck(_planned_attention, (query, key, value, scale)) + + +def test_stateful_attention_compiles_with_inductor(): + plan = _PreparedAttention(torch.full((1, 3, 1, 1), 0.5)) + compiled = torch.compile(plan, fullgraph=True, dynamic=True) + for length in (5, 9): + inputs = _inputs(length) + torch.testing.assert_close(compiled(*inputs), _reference(plan, inputs)) diff --git a/tests/diffusion/attention/test_flash_attn.py b/tests/diffusion/attention/test_flash_attn.py index 5dbd8d794..9a5a942c8 100644 --- a/tests/diffusion/attention/test_flash_attn.py +++ b/tests/diffusion/attention/test_flash_attn.py @@ -21,6 +21,7 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl from vllm_omni.diffusion.attention.backends.utils import fa # noqa: E402 +from vllm_omni.diffusion.attention.capabilities import CompilationMode, ExecutionContext, SupportStatus from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] @@ -432,6 +433,35 @@ def test_varlen_masked_routing_by_role(monkeypatch): assert cu_seqlens_k.tolist() == [0, 2, 5] +@pytest.mark.cpu +def test_published_noop_mask_uses_dense_path(monkeypatch): + calls = [] + + def fake_dense_func(q, k, v, **_kwargs): + calls.append("dense") + return q + + def fail_varlen(*_args, **_kwargs): + raise AssertionError("published no-op mask must not use varlen attention") + + monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True) + monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", False) + monkeypatch.setattr(fa, "flash_attn_func", fake_dense_func) + monkeypatch.setattr(fa, "flash_attn_varlen_func", fail_varlen) + + impl = FlashAttentionImpl(num_heads=2, head_size=8, softmax_scale=0.5, causal=False) + query = torch.randn(1, 4, 2, 8) + metadata = AttentionMetadata( + attn_mask=torch.ones((1, 4), dtype=torch.bool), + extra={"attention_mask_mode": "none"}, + ) + + output = impl.forward_cuda(query, query, query, metadata) + + assert torch.equal(output, query) + assert calls == ["dense"] + + @pytest.mark.cpu def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch): calls = [] @@ -550,6 +580,199 @@ def _fake_torch_npu(monkeypatch, *, npu_fusion_attention=None): return fake +@pytest.fixture +def npu_contract(monkeypatch, mocker): + monkeypatch.delenv("MINDIE_SD_FA_TYPE", raising=False) + monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: pytest.fail("NPU must not validate FA4")) + query = torch.zeros(1, 8, 2, 4) + dense = mocker.Mock(return_value=torch.zeros_like(query)) + packed = mocker.Mock(return_value=torch.zeros_like(query.squeeze(0))) + _fake_mindiesd(monkeypatch, attention_forward=dense, attention_forward_varlen=packed) + return _npu_impl(), query, dense, packed + + +def _assert_platform_path(impl, platform, query, metadata, path): + context = ExecutionContext(platform=platform, require_fullgraph=True) + result = impl.resolve_execution_path(context, query, query, query, metadata) + assert result.path == path + assert result.kernel_variant == {"npu": "mindiesd", "rocm": "aiter"}[platform] + assert result.support.status is SupportStatus.UNMIGRATED + assert result.compilation_mode is CompilationMode.EAGER_ONLY + assert result.requested_support(context).status is SupportStatus.UNSUPPORTED + + +def _real_pad_metadata(): + cu = torch.tensor([0, 5, 8], dtype=torch.int32) + return AttentionMetadata( + extra={ + "npu_attn_varlen": True, + "cu_seqlens_q": cu, + "cu_seqlens_k": cu, + "max_seqlen_q": 5, + "max_seqlen_k": 5, + } + ) + + +@pytest.mark.cpu +@pytest.mark.parametrize("kv_cache_dtype", [None, "float", "auto"]) +def test_npu_contract_dense_dispatch(npu_contract, kv_cache_dtype): + impl, query, dense, packed = npu_contract + metadata = AttentionMetadata(extra={"kv_cache_dtype": kv_cache_dtype}) + _assert_platform_path(impl, "npu", query, metadata, "npu_dense") + impl.forward_npu(query, query, query, metadata) + dense.assert_called_once() + packed.assert_not_called() + assert dense.call_args.kwargs["attn_mask"] is None + + +@pytest.mark.cpu +def test_npu_contract_mask_layout(npu_contract): + impl, query, dense, packed = npu_contract + mask = torch.arange(8)[None] < 5 + metadata = AttentionMetadata(attn_mask=mask) + _assert_platform_path(impl, "npu", query, metadata, "npu_masked") + impl.forward_npu(query, query, query, metadata) + dense.assert_called_once() + packed.assert_not_called() + torch.testing.assert_close(dense.call_args.kwargs["attn_mask"], mask[:, None, None, :].expand(1, 1, 8, 8)) + + +@pytest.mark.cpu +def test_npu_contract_packed_dispatch(npu_contract): + impl, query, dense, packed = npu_contract + metadata = _real_pad_metadata() + _assert_platform_path(impl, "npu", query, metadata, "npu_packed_varlen") + impl.forward_npu(query, query, query, metadata) + dense.assert_not_called() + packed.assert_called_once() + assert packed.call_args.args[3:] == ([0, 5, 8], [0, 5, 8]) + + +@pytest.mark.cpu +def test_npu_contract_prefix_dispatch(npu_contract, monkeypatch): + impl, query, dense, packed = npu_contract + metadata = _real_pad_metadata() + _assert_platform_path(impl, "npu", query, metadata, "npu_packed_varlen") + # The same metadata resolves differently after the kernel choice changes. + monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention") + _assert_platform_path(impl, "npu", query, metadata, "npu_prefix_kv_slice") + impl.forward_npu(query, query, query, metadata) + dense.assert_called_once() + packed.assert_not_called() + torch.testing.assert_close(dense.call_args.args[1], query[:, :5]) + torch.testing.assert_close(dense.call_args.args[2], query[:, :5]) + assert dense.call_args.kwargs["attn_mask"] is None + + +@pytest.mark.cpu +def test_npu_contract_incomplete_packing_rebuilds_mask(npu_contract): + impl, query, dense, packed = npu_contract + # CUDA-like normalization rejects this metadata, but NPU can rebuild a mask. + metadata = AttentionMetadata(extra={"npu_attn_varlen": True, "max_seqlen_q": 5, "valid_kv_length": 5}) + _assert_platform_path(impl, "npu", query, metadata, "npu_masked_fallback") + impl.forward_npu(query, query, query, metadata) + dense.assert_called_once() + packed.assert_not_called() + mask = (torch.arange(8) < 5)[None, None, None, :].expand(1, 1, 8, 8) + torch.testing.assert_close(dense.call_args.kwargs["attn_mask"], mask) + + +@pytest.mark.cpu +def test_npu_contract_unusable_packing_does_not_run_unmasked(npu_contract): + impl, query, dense, packed = npu_contract + metadata = AttentionMetadata(extra={"npu_attn_varlen": True}) + _assert_platform_path(impl, "npu", query, metadata, "npu_masked_fallback") + with pytest.raises(ValueError, match="refusing to run unmasked attention"): + impl.forward_npu(query, query, query, metadata) + dense.assert_not_called() + packed.assert_not_called() + + +@pytest.fixture +def rocm_contract(monkeypatch, mocker): + query = torch.zeros(1, 8, 2, 4) + dense = mocker.Mock(return_value=torch.zeros_like(query)) + packed = mocker.Mock(side_effect=lambda *args, **kwargs: torch.zeros_like(args[0] if args else kwargs["q"])) + monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", False) + monkeypatch.setattr(fa, "IS_AITER", True) + monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True) + monkeypatch.setattr(fa, "flash_attn_func", dense) + monkeypatch.setattr(fa, "flash_attn_varlen_func", packed) + monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: pytest.fail("ROCm must not validate FA4")) + return _npu_impl(), query, dense, packed + + +@pytest.mark.cpu +def test_rocm_contract_dense_dispatch(rocm_contract): + impl, query, dense, packed = rocm_contract + _assert_platform_path(impl, "rocm", query, None, "rocm_dense") + impl.forward_hip(query, query, query) + dense.assert_called_once() + packed.assert_not_called() + + +@pytest.mark.cpu +def test_rocm_contract_masked_unpadding(rocm_contract): + impl, query, dense, packed = rocm_contract + metadata = AttentionMetadata(attn_mask=torch.arange(8)[None] < 5, extra={"attention_mask_mode": "padding"}) + _assert_platform_path(impl, "rocm", query, metadata, "rocm_masked_varlen") + impl.forward_hip(query, query, query, metadata) + dense.assert_not_called() + packed.assert_called_once() + assert packed.call_args.args[0].shape == (5, 2, 4) + torch.testing.assert_close(packed.call_args.kwargs["cu_seqlens_k"], torch.tensor([0, 5], dtype=torch.int32)) + + +@pytest.mark.cpu +def test_rocm_contract_multiple_packed_documents(rocm_contract): + impl, query, dense, packed = rocm_contract + cu = torch.tensor([0, 2, 4, 8], dtype=torch.int32) + metadata = AttentionMetadata( + extra={ + "cu_seqlens_q": cu, + "cu_seqlens_k": cu, + "max_seqlen_q": 4, + "max_seqlen_k": 4, + } + ) + _assert_platform_path(impl, "rocm", query, metadata, "rocm_packed_varlen") + # NPU's [real, pad] path cannot consume this packing contract. + assert impl._resolve_packed_seq_npu(query, query, metadata.extra) is None + impl.forward_hip(query, query, query, metadata) + dense.assert_not_called() + packed.assert_called_once() + assert packed.call_args.kwargs["cu_seqlens_q"] is cu + assert packed.call_args.kwargs["cu_seqlens_k"] is cu + + +@pytest.mark.cpu +def test_rocm_contract_unknown_mask_stays_runtime_dependent(rocm_contract, monkeypatch): + impl, query, dense, packed = rocm_contract + metadata = AttentionMetadata(attn_mask=torch.ones(1, 8, dtype=torch.bool)) + # Resolution may inspect shape/metadata, but must not read mask values. + with monkeypatch.context() as patch: + patch.setattr(torch, "any", lambda *_args, **_kwargs: pytest.fail("resolution read mask values")) + _assert_platform_path(impl, "rocm", query, metadata, "rocm_runtime_mask_dependent") + impl.forward_hip(query, query, query, metadata) + dense.assert_called_once() + packed.assert_not_called() + + +@pytest.mark.cpu +def test_rocm_contract_requires_initialized_aiter(monkeypatch): + monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", False) + monkeypatch.setattr(fa, "IS_AITER", False) + impl = _npu_impl() + query = torch.zeros(1, 8, 2, 4) + result = impl.resolve_execution_path( + ExecutionContext(platform="rocm", kernel_variant="aiter"), query, query, query, None + ) + assert result.path == "rocm_unverified" + assert result.kernel_variant is None + assert result.support.status is SupportStatus.UNMIGRATED + + # --- Test group A: boundary resolution (_resolve_packed_seq_npu) ------------- diff --git a/tests/diffusion/attention/test_flash_attn_compile.py b/tests/diffusion/attention/test_flash_attn_compile.py new file mode 100644 index 000000000..5e485bd5b --- /dev/null +++ b/tests/diffusion/attention/test_flash_attn_compile.py @@ -0,0 +1,308 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import functools +from dataclasses import dataclass + +import pytest +import torch + +from tests.helpers.mark import hardware_test +from vllm_omni.diffusion.attention.backends.flash_attn import ( + FlashAttentionBackend, + FlashAttentionImpl, +) +from vllm_omni.diffusion.attention.backends.utils import fa +from vllm_omni.diffusion.attention.capabilities import ( + CompilationMode, + ExecutionContext, + OuterBoundary, + ParallelStrategy, + SupportStatus, +) +from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention +from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig +from vllm_omni.diffusion.forward_context import ForwardContext +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda] + + +@pytest.fixture(autouse=True) +def isolated_compiler_cache(): + # Each case changes kernel/configuration on the same Attention.forward code + # object. Keep unrelated cases from exhausting its Dynamo recompile limit; + # retain the cache across all shapes and replays within an individual test. + torch.compiler.reset() + yield + torch.compiler.reset() + + +@dataclass +class _ParallelStrategyStub: + name: str + enabled: bool = True + + +def _make_attention_layer(monkeypatch, *, head_size=64): + from vllm_omni.diffusion.attention import layer as attention_layer + + monkeypatch.setattr( + attention_layer, + "get_attn_backend_for_role", + lambda **_kwargs: (FlashAttentionBackend, None), + ) + monkeypatch.setattr( + attention_layer, + "build_parallel_attention_strategy", + lambda **_kwargs: NoParallelAttention(), + ) + return attention_layer.Attention(num_heads=8, head_size=head_size, softmax_scale=head_size**-0.5, causal=False) + + +@pytest.fixture +def fake_fa4(monkeypatch, tmp_path): + marker = tmp_path / "kernel" + marker.write_text("loaded", encoding="utf-8") + + @functools.cache + def cached_kernel_loader(): + with open(marker, encoding="utf-8") as handle: + handle.read() + + def fake_attention(query, key, value, **kwargs): + cached_kernel_loader() + return ( + torch.nn.functional.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + scale=kwargs["softmax_scale"], + is_causal=kwargs["causal"], + ) + .transpose(1, 2) + .contiguous() + ) + + monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True) + monkeypatch.setattr(fa, "IS_FLASH_ATTN_4", True) + monkeypatch.setattr(fa, "flash_attn_func", fake_attention) + monkeypatch.setattr(fa, "validate_fa4_head_dims", lambda *_args: True) + + return FlashAttentionImpl( + num_heads=8, + head_size=64, + softmax_scale=0.125, + causal=False, + ) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("noncontiguous", [False, True]) +def test_fa4_dense_dispatch_is_opaque_to_dynamic_torch_compile(fake_fa4, noncontiguous): + """Fast Dynamo/custom-op regression; Inductor coverage is separate.""" + impl = fake_fa4 + q = torch.randn(1, 16, 8, 64, device="cuda", dtype=torch.bfloat16) + if noncontiguous: + q = q.transpose(1, 2).contiguous().transpose(1, 2) + assert not q.is_contiguous() + path = impl.resolve_execution_path( + ExecutionContext( + platform="cuda", + kernel_variant="fa4", + dtype="bfloat16", + require_fullgraph=True, + ), + q, + q, + q, + None, + ) + assert path.support.status is SupportStatus.SUPPORTED + assert path.compilation_mode is CompilationMode.CUSTOM_OP + + compile_count = 0 + + def counting_backend(graph_module, _example_inputs): + nonlocal compile_count + compile_count += 1 + return graph_module.forward + + compiled = torch.compile( + lambda query, key, value: impl.forward_cuda(query, key, value), + backend=counting_backend, + fullgraph=True, + dynamic=True, + ) + out = compiled(q, q, q) + q2 = torch.randn(1, 24, 8, 64, device="cuda", dtype=torch.bfloat16) + if noncontiguous: + q2 = q2.transpose(1, 2).contiguous().transpose(1, 2) + out2 = compiled(q2, q2, q2) + assert out.is_contiguous() + assert out2.is_contiguous() + + assert out.shape == q.shape + assert out2.shape == q2.shape + torch.testing.assert_close(out, impl.forward_cuda(q, q, q)) + torch.testing.assert_close(out2, impl.forward_cuda(q2, q2, q2)) + assert compile_count == 1 + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_fa4_production_attention_entry_compiles_with_inductor( + monkeypatch, + fake_fa4, +): + """Compile through Attention.forward with the default Inductor backend.""" + from vllm_omni.diffusion.attention import layer as attention_layer + + layer = _make_attention_layer(monkeypatch) + layer.attention = fake_fa4 + q = torch.randn(1, 16, 8, 64, device="cuda", dtype=torch.bfloat16) + + resolved_contexts = [] + resolve_execution_path = fake_fa4.resolve_execution_path + + def record_execution_path(context, query, key, value, attn_metadata): + resolved_contexts.append(context) + return resolve_execution_path( + context, + query, + key, + value, + attn_metadata, + ) + + monkeypatch.setattr(fake_fa4, "resolve_execution_path", record_execution_path) + context = ExecutionContext( + platform="cuda", + kernel_variant="fa4", + dtype="bfloat16", + paged_kv=True, + parallel_strategy=ParallelStrategy.ULYSSES, + require_fullgraph=True, + ) + path = layer.resolve_execution_path(context, q, q, q, None) + assert path.support.status is SupportStatus.SUPPORTED + assert path.compilation_mode is CompilationMode.CUSTOM_OP + assert resolved_contexts[-1].paged_kv is False + assert resolved_contexts[-1].parallel_strategy is ParallelStrategy.NONE + + compiled = torch.compile(layer, fullgraph=True, dynamic=True) + out = compiled(q, q, q) + + assert out.shape == q.shape + torch.testing.assert_close(out, layer(q, q, q)) + q2 = torch.randn(1, 24, 8, 64, device="cuda", dtype=torch.bfloat16) + torch.testing.assert_close(compiled(q2, q2, q2), layer(q2, q2, q2)) + + layer.paged_kv_cache_role = "self" + monkeypatch.setattr(layer, "_active_paged_kv_adapter", lambda: object()) + paged_path = layer.resolve_execution_path(context, q, q, q, None) + assert resolved_contexts[-1].paged_kv is True + assert paged_path.support.status is SupportStatus.UNMIGRATED + layer.paged_kv_cache_role = None + + layer.use_ring = True + layer.parallel_strategy = _ParallelStrategyStub(name="ulysses") + hybrid_path = layer.resolve_execution_path(context, q, q, q, None) + assert resolved_contexts[-1].parallel_strategy is ParallelStrategy.HYBRID_ULYSSES_RING + assert hybrid_path.support.status is SupportStatus.UNMIGRATED + layer.use_ring = False + layer.parallel_strategy = NoParallelAttention() + + layer._hsdp_compile_boundary_enabled = True + hsdp_path = layer.resolve_execution_path(context, q, q, q, None) + assert OuterBoundary.HSDP in resolved_contexts[-1].outer_boundaries + assert hsdp_path.support.status is SupportStatus.UNMIGRATED + assert hsdp_path.compilation_mode is CompilationMode.EAGER_ONLY + + layer._hsdp_compile_boundary_enabled = False + runtime_context = ForwardContext( + omni_diffusion_config=OmniDiffusionConfig( + parallel_config=DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=1), + ), + ) + monkeypatch.setattr(attention_layer, "is_forward_context_available", lambda: True) + monkeypatch.setattr(attention_layer, "get_forward_context", lambda: runtime_context) + layer.use_ring = True + layer.parallel_strategy = _ParallelStrategyStub(name="ring") + runtime_hsdp_path = layer.resolve_execution_path(context, q, q, q, None) + assert OuterBoundary.HSDP in resolved_contexts[-1].outer_boundaries + assert resolved_contexts[-1].parallel_strategy is ParallelStrategy.NONE + assert runtime_hsdp_path.support.status is SupportStatus.UNMIGRATED + + +@hardware_test(res={"cuda": "B200"}, num_cards=1) +@pytest.mark.parametrize("head_dims", [(32, 32), (64, 64), (80, 48), (192, 128), (256, 256)]) +def test_real_fa4_fullgraph_matches_sdpa_across_shapes(monkeypatch, head_dims): + """Validate real FA4 numerics through the complete compiled attention layer.""" + capability = current_omni_platform.get_device_capability() if current_omni_platform.is_cuda() else None + if capability is None or capability.major < 10: + pytest.skip("Requires Blackwell") + if not fa.IS_FLASH_ATTN_4 or fa.flash_attn_func is None: + pytest.skip("Requires CuTe FlashAttention-4") + + from torch.nn.attention import SDPBackend, sdpa_kernel + + head_size, value_head_size = head_dims + layer = _make_attention_layer(monkeypatch, head_size=head_size) + compiled = torch.compile(layer, fullgraph=True, dynamic=True) + generator = torch.Generator(device="cuda").manual_seed(42) + context = ExecutionContext(platform="cuda", require_fullgraph=True) + # Tile boundaries, cross-attention lengths, batching, and a longer sequence. + for batch, q_length, kv_length in ((1, 16, 16), (1, 129, 257), (2, 257, 129), (1, 1024, 1024)): + query, key, value = ( + torch.randn(batch, length, 8, dim, device="cuda", dtype=torch.bfloat16, generator=generator) + for length, dim in ((q_length, head_size), (kv_length, head_size), (kv_length, value_head_size)) + ) + path = layer.resolve_execution_path(context, query, key, value, None) + assert path.requested_support(context).status is SupportStatus.SUPPORTED + with sdpa_kernel(SDPBackend.MATH): + reference = torch.nn.functional.scaled_dot_product_attention( + query.transpose(1, 2).float(), + key.transpose(1, 2).float(), + value.transpose(1, 2).float(), + scale=head_size**-0.5, + ).transpose(1, 2) + actual = compiled(query, key, value) + assert actual.dtype == query.dtype + assert actual.device == query.device + torch.testing.assert_close(actual.float(), reference, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(actual, layer(query, key, value)) + + +@hardware_test(res={"cuda": "B200"}, num_cards=1) +@pytest.mark.parametrize("head_dims", [(192, 192), (65, 64), (64, 65)]) +def test_real_fa4_dimension_rejection_matches_kernel(monkeypatch, head_dims): + capability = current_omni_platform.get_device_capability() if current_omni_platform.is_cuda() else None + if capability is None or capability.major not in (10, 11) or not fa.IS_FLASH_ATTN_4: + pytest.skip("Requires SM100/SM110 CuTe FlashAttention-4") + head_size, value_head_size = head_dims + layer = _make_attention_layer(monkeypatch, head_size=head_size) + query = torch.randn(1, 16, 8, head_size, device="cuda", dtype=torch.bfloat16) + value = torch.randn(1, 16, 8, value_head_size, device="cuda", dtype=torch.bfloat16) + context = ExecutionContext(platform="cuda", require_fullgraph=True) + path = layer.resolve_execution_path(context, query, query, value, None) + assert path.requested_support(context).status is SupportStatus.UNSUPPORTED + with pytest.raises(AssertionError) as error: + layer(query, query, value) + assert str(error.value) in path.support.reason + + +@hardware_test(res={"cuda": "B200"}, num_cards=1) +@pytest.mark.parametrize("noncontiguous", [False, True]) +def test_real_fa4_custom_op_unequal_value_dimension_schema(noncontiguous): + if not current_omni_platform.is_cuda() or not fa.IS_FLASH_ATTN_4: + pytest.skip("Requires CuTe FlashAttention-4") + query = torch.randn(1, 17, 8, 80, device="cuda", dtype=torch.bfloat16) + key = torch.randn(1, 25, 8, 80, device="cuda", dtype=torch.bfloat16) + value = torch.randn(1, 25, 8, 48, device="cuda", dtype=torch.bfloat16) + if noncontiguous: + query, key, value = (t.transpose(1, 2).contiguous().transpose(1, 2) for t in (query, key, value)) + assert not query.is_contiguous() + torch.library.opcheck( + torch.ops.vllm_omni.fa4_dense_attention.default, + (query, key, value, 80**-0.5, False, False), + ) diff --git a/tests/diffusion/test_attention_compile.py b/tests/diffusion/test_attention_compile.py index d39744641..93ec9e4cb 100644 --- a/tests/diffusion/test_attention_compile.py +++ b/tests/diffusion/test_attention_compile.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 types import SimpleNamespace @@ -14,6 +14,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] def test_attention_uses_compile_boundary_for_hsdp(monkeypatch): attention = object.__new__(Attention) + attention._hsdp_compile_boundary_enabled = False calls = [] def _boundary(query, key, value, attn_metadata=None): @@ -38,6 +39,7 @@ def test_attention_uses_compile_boundary_for_hsdp(monkeypatch): def test_attention_keeps_compiled_impl_without_hsdp(monkeypatch): attention = object.__new__(Attention) + attention._hsdp_compile_boundary_enabled = False calls = [] def _boundary(query, key, value, attn_metadata=None): @@ -60,8 +62,10 @@ def test_attention_keeps_compiled_impl_without_hsdp(monkeypatch): assert calls == ["impl"] -def test_attention_keeps_compiled_impl_without_diffusion_config(monkeypatch): +@pytest.mark.parametrize("initialized_with_hsdp", [False, True]) +def test_attention_compile_boundary_without_diffusion_config(monkeypatch, initialized_with_hsdp): attention = object.__new__(Attention) + attention._hsdp_compile_boundary_enabled = initialized_with_hsdp calls = [] def _boundary(query, key, value, attn_metadata=None): @@ -80,4 +84,4 @@ def test_attention_keeps_compiled_impl_without_diffusion_config(monkeypatch): with set_forward_context(vllm_config=None): assert Attention.forward(attention, query, query, query) is query - assert calls == ["impl"] + assert calls == (["boundary"] if initialized_with_hsdp else ["impl"]) diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index 592f5ba9b..ed6a9a452 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -7,6 +7,10 @@ from typing import Any, Generic, Literal, TypeVar import torch +from vllm_omni.diffusion.attention.capabilities import ( + ExecutionContext, + ExecutionPathResult, +) from vllm_omni.platforms import current_omni_platform @@ -132,6 +136,15 @@ class AttentionBackend(ABC): return False + @classmethod + def resolve_capabilities(cls, context: ExecutionContext) -> ExecutionPathResult: + """Return conservative pre-construction capabilities. + + Unmigrated backends remain advisory so introducing the contract cannot + reject an execution path that works today. + """ + return ExecutionPathResult.unmigrated(cls.get_name(), context) + @dataclass(frozen=True, slots=True) class QueryRange: @@ -218,6 +231,9 @@ class AttentionMetadata: # "valid_kv_length": int — contiguous valid K/V prefix length. Local # backends may slice it directly; ring backends trim each circulated # K/V block from the same global prefix. + # "attention_mask_mode": "none" | "padding" | "arbitrary" — producer + # assertion used by capability planning. "none" means a present mask + # is semantically a no-op; absent means its contents are runtime-dependent. # "npu_attn_varlen": bool — model opt-in for the NPU packed varlen path # (TND npu_fusion_attention driven by cu_seqlens, mask never read). # Requires the [real, pad] two-document packing contract; see @@ -277,6 +293,21 @@ class AttentionImpl(ABC, Generic[T]): return True return kv_cache_dtype in cls._supported_kv_cache_dtypes.get(platform_key, set()) + def resolve_execution_path( + self, + context: ExecutionContext, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: T | None, + ) -> ExecutionPathResult: + """Resolve the complete path after backend initialization. + + Subclasses report verified results; this inspection API does not enforce them. + """ + del query, key, value, attn_metadata + return ExecutionPathResult.unmigrated(type(self).__name__, context) + def forward( self, query: torch.Tensor, diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index 679bf4955..2ba6b415a 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -3,7 +3,9 @@ import os from collections.abc import Callable +from dataclasses import replace from functools import cache, partial +from typing import NamedTuple import torch from vllm.logger import init_logger @@ -14,6 +16,16 @@ from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import ( piecewise_attn, run_paged_piecewise_plan, ) +from vllm_omni.diffusion.attention.capabilities import ( + CapabilityResult, + CompilationMode, + ExecutionContext, + ExecutionPathResult, + MaskMode, + PackingMode, + ParallelStrategy, + SupportStatus, +) from vllm_omni.diffusion.config import get_current_diffusion_config_or_none from vllm_omni.platforms import current_omni_platform @@ -29,6 +41,133 @@ def _get_npu_compressed_causal_mask(device: torch.device) -> torch.Tensor: ).contiguous() +if not hasattr(torch.ops.vllm_omni, "fa4_dense_attention"): + + @torch.library.custom_op("vllm_omni::fa4_dense_attention", mutates_args=()) + def _fa4_dense_attention_op( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + softmax_scale: float, + causal: bool, + deterministic: bool, + ) -> torch.Tensor: + from vllm_omni.diffusion.attention.backends.utils.fa import flash_attn_func + + if flash_attn_func is None: + raise RuntimeError("CuTe FlashAttention-4 is unavailable") + kwargs = { + "causal": causal, + "softmax_scale": softmax_scale, + } + if deterministic: + kwargs["deterministic"] = True + out = flash_attn_func(query, key, value, **kwargs) + return out[0] if isinstance(out, tuple) else out + + @_fa4_dense_attention_op.register_fake + def _fa4_dense_attention_fake( + query, + key, + value, + softmax_scale, + causal, + deterministic, + ): + # FA4 returns contiguous output even when Q is noncontiguous; V sets its head dimension. + return query.new_empty((*query.shape[:-1], value.shape[-1])) + + +_fa4_dense_attention_op = torch.ops.vllm_omni.fa4_dense_attention + + +_PACKED_KEYS = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k") + + +class _FlashAttentionMetadataPlan(NamedTuple): + attention_mask: torch.Tensor | None + mask_mode: MaskMode + full_attn_spans: list[list[tuple[int, int]]] | None + extra: dict + packing_mode: PackingMode + + +def _normalize_flash_attention_metadata( + attn_metadata: AttentionMetadata | None, +) -> _FlashAttentionMetadataPlan: + attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None + full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None + extra = attn_metadata.extra if attn_metadata is not None else {} + published_mask_mode = extra.get("attention_mask_mode") + if attention_mask is None: + mask_mode = MaskMode.NONE + elif published_mask_mode is None: + mask_mode = MaskMode.UNKNOWN + else: + try: + mask_mode = MaskMode(published_mask_mode) + except ValueError as error: + raise ValueError(f"Unknown attention_mask_mode {published_mask_mode!r}") from error + if mask_mode is MaskMode.UNKNOWN: + raise ValueError("attention_mask_mode='unknown' is reserved for unpublished semantics") + # Piecewise dispatch takes precedence and does not consume packed metadata. + present_packed_keys = [key for key in _PACKED_KEYS if key in extra] if full_attn_spans is None else [] + if present_packed_keys and len(present_packed_keys) != len(_PACKED_KEYS): + missing = sorted(set(_PACKED_KEYS) - set(present_packed_keys)) + raise ValueError(f"Incomplete packed FlashAttention metadata; missing {missing}") + + packing_mode = PackingMode.NONE + if present_packed_keys: + cu_seqlens_q = extra["cu_seqlens_q"] + packing_mode = PackingMode.MULTI_DOCUMENT if cu_seqlens_q.shape[0] > 3 else PackingMode.PACKED_PADDING + return _FlashAttentionMetadataPlan( + attention_mask=attention_mask, + mask_mode=mask_mode, + full_attn_spans=full_attn_spans, + extra=extra, + packing_mode=packing_mode, + ) + + +def _flash_attention_execution_path( + context: ExecutionContext, +) -> ExecutionPathResult: + if context.mask_mode is MaskMode.UNKNOWN: + return ExecutionPathResult.unmigrated( + "FLASH_ATTN", + context, + path="runtime_mask_dependent", + ) + if not ( + context.platform == "cuda" + and context.kernel_variant == "fa4" + and context.dtype in {"bfloat16", "torch.bfloat16"} + and context.causal is False + and context.mask_mode is MaskMode.NONE + and context.packing_mode is PackingMode.NONE + and not context.piecewise + and not context.paged_kv + and context.kv_cache_dtype is None + and context.parallel_strategy is ParallelStrategy.NONE + and not context.outer_boundaries + ): + return ExecutionPathResult.unmigrated( + "FLASH_ATTN", + context, + path="unverified", + ) + + return ExecutionPathResult( + backend="FLASH_ATTN", + path="fa4_dense", + support=CapabilityResult.supported(), + compilation_mode=CompilationMode.CUSTOM_OP, + platform=context.platform, + kernel_variant=context.kernel_variant, + parallel_strategy=context.parallel_strategy, + ) + + class FlashAttentionBackend(AttentionBackend): accept_output_buffer: bool = True supports_piecewise_spans: bool = True @@ -74,6 +213,10 @@ class FlashAttentionBackend(AttentionBackend): def get_impl_cls() -> type["FlashAttentionImpl"]: return FlashAttentionImpl + @classmethod + def resolve_capabilities(cls, context: ExecutionContext) -> ExecutionPathResult: + return _flash_attention_execution_path(replace(context, kernel_variant=None)) + class FlashAttentionImpl(AttentionImpl[AttentionMetadata]): # Per-platform attention quantization support. @@ -108,11 +251,135 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]): self.softmax_scale = softmax_scale self.qkv_layout = qkv_layout self.is_cross_attn = role == "cross" + from vllm_omni.diffusion.attention.backends.utils.fa import IS_AITER, IS_FLASH_ATTN_4 + + self._kernel_variant = "fa4" if IS_FLASH_ATTN_4 else "aiter" if IS_AITER else None cfg = get_current_diffusion_config_or_none() self.fa_deterministic = bool(getattr(cfg, "fa_deterministic", False)) if cfg is not None else False if backend_kwargs: logger.warning("FlashAttentionImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys())) + def resolve_execution_path( + self, + context: ExecutionContext, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None, + ) -> ExecutionPathResult: + if context.platform == "npu": + return self._resolve_npu_execution_path(context, query, key, attn_metadata) + metadata = _normalize_flash_attention_metadata(attn_metadata) + context = replace( + context, + kernel_variant=self._kernel_variant, + dtype=str(query.dtype).removeprefix("torch."), + causal=self.causal, + mask_mode=metadata.mask_mode, + packing_mode=metadata.packing_mode, + piecewise=metadata.full_attn_spans is not None, + kv_cache_dtype=metadata.extra.get("kv_cache_dtype"), + ) + if context.platform == "rocm": + return self._resolve_rocm_execution_path(context) + result = _flash_attention_execution_path(context) + if result.support.status is not SupportStatus.SUPPORTED: + return result + + from vllm_omni.diffusion.attention.backends.utils.fa import validate_fa4_head_dims + + try: + if query.dtype != key.dtype or query.dtype != value.dtype: + raise ValueError("Q, K, and V dtypes must match") + if query.device != key.device or query.device != value.device: + raise ValueError("Q, K, and V devices must match") + if query.shape[-1] != key.shape[-1]: + raise ValueError("Q and K head dimensions must match") + verified = validate_fa4_head_dims(query.shape[-1], value.shape[-1], 16 // value.element_size()) + except (AssertionError, ValueError) as error: + return replace( + result, + support=CapabilityResult.unsupported(f"FA4: {error} Select compatible inputs or another backend."), + ) + if not verified: + return replace( + result, + support=CapabilityResult.unmigrated("Selected FA4 kernel has no available head-dimension validator"), + compilation_mode=CompilationMode.EAGER_ONLY, + ) + return result + + @staticmethod + def _resolve_rocm_execution_path(context: ExecutionContext) -> ExecutionPathResult: + """Describe AITER dispatch; device and compiler validation is pending.""" + result = ExecutionPathResult.unmigrated("FLASH_ATTN", context, path="rocm_unverified") + if ( + context.kernel_variant != "aiter" + or context.parallel_strategy is not ParallelStrategy.NONE + or context.outer_boundaries + or context.paged_kv + or context.piecewise + or context.kv_cache_dtype is not None + ): + return result + if context.packing_mode is not PackingMode.NONE: + path = "rocm_packed_varlen" + elif context.mask_mode is MaskMode.UNKNOWN: + path = "rocm_runtime_mask_dependent" + elif context.mask_mode is not MaskMode.NONE: + path = "rocm_masked_varlen" + else: + path = "rocm_dense" + return replace( + result, + path=path, + support=CapabilityResult.unmigrated("AITER routing only; kernel and compilation validation require ROCm"), + ) + + def _resolve_npu_execution_path( + self, + context: ExecutionContext, + query: torch.Tensor, + key: torch.Tensor, + attn_metadata: AttentionMetadata | None, + ) -> ExecutionPathResult: + """Describe existing MindIE routing without claiming hardware validation. + + NPU packing has its own opt-in and [real, pad] contract. In particular, + incomplete packed metadata can fall back to a mask, so CUDA metadata + normalization must not run here. No tensor values are inspected. + """ + context = replace(context, kernel_variant="mindiesd") + result = ExecutionPathResult.unmigrated("FLASH_ATTN", context, path="npu_unverified") + extra = attn_metadata.extra if attn_metadata else {} + if ( + context.parallel_strategy is not ParallelStrategy.NONE + or context.outer_boundaries + or context.paged_kv + or self.causal + or extra.get("kv_cache_dtype") not in (None, "float", "auto") + or (attn_metadata is not None and attn_metadata.full_attn_spans is not None) + ): + return result + + path = "npu_masked" if attn_metadata is not None and attn_metadata.attn_mask is not None else "npu_dense" + if extra.get("npu_attn_varlen", False): + if self._resolve_packed_seq_npu(query, key, extra) is not None: + path = ( + "npu_prefix_kv_slice" + if os.environ.get("MINDIE_SD_FA_TYPE") == "ascend_laser_attention" + else "npu_packed_varlen" + ) + else: + # This route may reject the request if neither an explicit mask + # nor usable valid_kv_length is supplied; it is not support. + path = "npu_masked_fallback" + return replace( + result, + path=path, + support=CapabilityResult.unmigrated("NPU routing only; kernel and compilation validation require Ascend"), + ) + def _warn_fa_deterministic_non_dense(self, path: str) -> None: if not self.fa_deterministic: return @@ -401,6 +668,7 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]): """CUDA/ROCm/MUSA flash attention implementation.""" from vllm_omni.diffusion.attention.backends.utils.fa import ( HAS_FLASH_ATTN, + IS_FLASH_ATTN_4, flash_attn_func, flash_attn_varlen_func, ) @@ -412,9 +680,10 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]): "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA" ) - attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None - full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None - extra = attn_metadata.extra if attn_metadata is not None else {} + metadata_plan = _normalize_flash_attention_metadata(attn_metadata) + attention_mask = metadata_plan.attention_mask + full_attn_spans = metadata_plan.full_attn_spans + extra = metadata_plan.extra # Try piecewise attention if full_attn_spans is not None: @@ -443,12 +712,7 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]): query_ranges=None if attn_metadata is None else attn_metadata.query_ranges, ) - packed_keys = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k") - present_packed_keys = [key for key in packed_keys if key in extra] - if present_packed_keys: - if len(present_packed_keys) != len(packed_keys): - missing = sorted(set(packed_keys) - set(present_packed_keys)) - raise ValueError(f"Incomplete packed FlashAttention metadata; missing {missing}") + if metadata_plan.packing_mode is not PackingMode.NONE: self._warn_fa_deterministic_non_dense("packed-varlen") return self._forward_varlen_packed( query, @@ -460,7 +724,14 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]): max_seqlen_k=extra["max_seqlen_k"], ) - if attention_mask is not None and torch.any(~attention_mask): + use_masked_path = False + if attention_mask is not None: + use_masked_path = ( + torch.any(~attention_mask) + if metadata_plan.mask_mode is MaskMode.UNKNOWN + else metadata_plan.mask_mode is not MaskMode.NONE + ) + if use_masked_path: self._warn_fa_deterministic_non_dense("masked-varlen") return self._forward_varlen_masked( query, @@ -470,6 +741,15 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]): ) if flash_attn_func is not None: + if IS_FLASH_ATTN_4: + return _fa4_dense_attention_op( + query, + key, + value, + self.softmax_scale, + self.causal, + self.fa_deterministic, + ) fa_kwargs = { "causal": self.causal, "softmax_scale": self.softmax_scale, diff --git a/vllm_omni/diffusion/attention/backends/utils/fa.py b/vllm_omni/diffusion/attention/backends/utils/fa.py index e2f96abab..25eed8b24 100644 --- a/vllm_omni/diffusion/attention/backends/utils/fa.py +++ b/vllm_omni/diffusion/attention/backends/utils/fa.py @@ -59,6 +59,8 @@ def is_flash_attn_4_available() -> bool: # Bind via aliases so mypy does not treat each candidate import as a redefinition. flash_attn_func: FlashAttnFn | None = None flash_attn_varlen_func: FlashAttnFn | None = None +IS_FLASH_ATTN_4 = False +IS_AITER = False if current_omni_platform.is_rocm(): # ROCm: try Aiter first @@ -71,6 +73,7 @@ if current_omni_platform.is_rocm(): flash_attn_func = _fa_func flash_attn_varlen_func = _fa_varlen + IS_AITER = True except (ImportError, ModuleNotFoundError): pass elif current_omni_platform.is_xpu(): @@ -102,6 +105,7 @@ else: flash_attn_func = _fa_func flash_attn_varlen_func = _fa_varlen + IS_FLASH_ATTN_4 = True logger.info("Using CuTe FlashAttention-4 on Blackwell") except Exception as exc: # Optional FA4 dependencies may be present but ABI-incompatible @@ -166,6 +170,27 @@ else: HAS_FLASH_ATTN = flash_attn_func is not None or flash_attn_varlen_func is not None +def validate_fa4_head_dims(head_dim: int, head_dim_v: int, alignment: int) -> bool: + """Use the selected FA4 implementation's constraints during path resolution. + + Return False when its validator is unavailable; propagate kernel validation + errors for the backend to report. Keep the private FA4 API dependency here, + rather than duplicating architecture/version-specific dimension rules. + This is a planning-time call, never part of compiled tensor execution. + """ + try: + from flash_attn.cute.interface import _get_device_arch, _validate_head_dims + except ImportError: + return False + + arch = _get_device_arch() // 10 + # FA4 routes SM80/SM120 through separate kernels and bypasses this validator. + if arch not in (9, 10, 11): + return False + _validate_head_dims(head_dim, head_dim_v, arch, alignment) + return True + + def _choose_vllm_flash_attn_version( device_major: int, requested: str | int | None, diff --git a/vllm_omni/diffusion/attention/capabilities.py b/vllm_omni/diffusion/attention/capabilities.py new file mode 100644 index 000000000..00a328b8c --- /dev/null +++ b/vllm_omni/diffusion/attention/capabilities.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum + + +class SupportStatus(str, Enum): + SUPPORTED = "supported" + UNSUPPORTED = "unsupported" + UNMIGRATED = "unmigrated" + + +class CompilationMode(str, Enum): + TRACEABLE = "traceable" + CUSTOM_OP = "custom_op" + EAGER_ONLY = "eager_only" + + +class PackingMode(str, Enum): + NONE = "none" + PACKED_PADDING = "packed_padding" + MULTI_DOCUMENT = "multi_document" + + +class MaskMode(str, Enum): + NONE = "none" + PADDING = "padding" + ARBITRARY = "arbitrary" + UNKNOWN = "unknown" + + +class ParallelStrategy(str, Enum): + NONE = "none" + ULYSSES = "ulysses" + RING = "ring" + HYBRID_ULYSSES_RING = "hybrid_ulysses_ring" + ALLGATHER_KV = "allgather_kv" + + +class OuterBoundary(str, Enum): + HSDP = "hsdp" + + +@dataclass(frozen=True, slots=True) +class CapabilityResult: + status: SupportStatus + reason: str | None = None + + def __post_init__(self) -> None: + if self.status is SupportStatus.UNSUPPORTED and not self.reason: + raise ValueError("UNSUPPORTED capability results require an actionable reason") + + @classmethod + def supported(cls) -> CapabilityResult: + return cls(SupportStatus.SUPPORTED) + + @classmethod + def unsupported(cls, reason: str) -> CapabilityResult: + return cls(SupportStatus.UNSUPPORTED, reason) + + @classmethod + def unmigrated(cls, reason: str | None = None) -> CapabilityResult: + return cls(SupportStatus.UNMIGRATED, reason) + + +@dataclass(frozen=True, slots=True) +class ExecutionContext: + platform: str + kernel_variant: str | None = None + dtype: str | None = None + causal: bool | None = None + mask_mode: MaskMode = MaskMode.NONE + packing_mode: PackingMode = PackingMode.NONE + piecewise: bool = False + paged_kv: bool = False + kv_cache_dtype: str | None = None + parallel_strategy: ParallelStrategy = ParallelStrategy.NONE + require_fullgraph: bool = False + outer_boundaries: frozenset[OuterBoundary] = frozenset() + + +@dataclass(frozen=True, slots=True) +class ExecutionPathResult: + backend: str + path: str + support: CapabilityResult + compilation_mode: CompilationMode + platform: str + kernel_variant: str | None + parallel_strategy: ParallelStrategy + + @classmethod + def unmigrated( + cls, + backend: str, + context: ExecutionContext, + *, + path: str = "unmigrated", + ) -> ExecutionPathResult: + return cls( + backend=backend, + path=path, + support=CapabilityResult.unmigrated(), + compilation_mode=CompilationMode.EAGER_ONLY, + platform=context.platform, + kernel_variant=context.kernel_variant, + parallel_strategy=context.parallel_strategy, + ) + + def requested_support(self, context: ExecutionContext) -> CapabilityResult: + """Validate request-level guarantees without changing path migration state.""" + if context.require_fullgraph and self.support.status is SupportStatus.UNMIGRATED: + return CapabilityResult.unsupported( + f"{self.backend} path {self.path!r} has no verified fullgraph declaration; " + "use a verified path or disable the guaranteed-fullgraph request" + ) + if self.support.status is not SupportStatus.SUPPORTED: + return self.support + if context.require_fullgraph and self.compilation_mode is CompilationMode.EAGER_ONLY: + return CapabilityResult.unsupported( + f"{self.backend} path {self.path!r} does not support guaranteed fullgraph compilation" + ) + return CapabilityResult.supported() diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 0ad41e80f..d239525c2 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -19,7 +19,13 @@ from vllm.model_executor.models.utils import extract_layer_index from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheSpec from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend, AttentionImpl, AttentionMetadata -from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend +from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend, SDPAImpl +from vllm_omni.diffusion.attention.capabilities import ( + ExecutionContext, + ExecutionPathResult, + OuterBoundary, + ParallelStrategy, +) from vllm_omni.diffusion.attention.parallel import build_parallel_attention_strategy from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention @@ -119,6 +125,8 @@ class Attention(nn.Module): config = get_current_diffusion_config_or_none() attention_config = config.diffusion_attention_config if config is not None else None + parallel_config = getattr(config, "parallel_config", None) + self._hsdp_compile_boundary_enabled = bool(getattr(parallel_config, "use_hsdp", False)) from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata @@ -131,6 +139,9 @@ class Attention(nn.Module): is DiffusionKVCacheMode.PAGED_SCHEDULER ) self._scheduler_paged_kv = scheduler_paged_kv + self.attn_backend: type[AttentionBackend] | None + self.attention: AttentionImpl | nn.Module + self.sdpa_fallback: SDPAImpl | None if custom_attention is None: attn_backend_cls, spec = get_attn_backend_for_role( role=role, @@ -166,7 +177,6 @@ class Attention(nn.Module): attn_backend_cls.get_name(), dense_backend_name, ) - parallel_config = getattr(config, "parallel_config", None) allgather_degree = getattr(parallel_config, "allgather_degree", 1) # TODO: Move AllGather-KV compatibility into an AttentionBackend capability # so validation does not depend on backend names. @@ -279,6 +289,7 @@ class Attention(nn.Module): if self.paged_kv_cache_role is None: return None + assert self.attn_backend is not None # Custom attention cannot opt into paged KV. dtype = self.paged_kv_cache_dtype or vllm_config.model_config.dtype # Keep backend layout discovery under the same config context used by # upstream vLLM's attention-spec collector. vLLM 0.29 moved the @@ -378,17 +389,68 @@ class Attention(nn.Module): value: torch.Tensor, attn_metadata: AttentionMetadata | None = None, ) -> torch.Tensor: - if torch.compiler.is_compiling() and is_forward_context_available(): - od_config = get_forward_context().omni_diffusion_config - parallel_config = getattr(od_config, "parallel_config", None) - if getattr(parallel_config, "use_hsdp", False): - # Keep HSDP/FSDP2 parameter all-gather outside Inductor's - # attention graph; otherwise scheduler dependency analysis can - # fail on the fused attention region. - return self._forward_hsdp_compile_boundary(query, key, value, attn_metadata) + if torch.compiler.is_compiling() and self._uses_hsdp_compile_boundary(): + # Keep HSDP/FSDP2 parameter all-gather outside Inductor's + # attention graph; otherwise scheduler dependency analysis can + # fail on the fused attention region. + return self._forward_hsdp_compile_boundary(query, key, value, attn_metadata) return self._forward_impl(query, key, value, attn_metadata) + def _uses_hsdp_compile_boundary(self) -> bool: + if self._hsdp_compile_boundary_enabled: + return True + if not is_forward_context_available(): + return False + od_config = get_forward_context().omni_diffusion_config + parallel_config = getattr(od_config, "parallel_config", None) + return bool(getattr(parallel_config, "use_hsdp", False)) + + def resolve_execution_path( + self, + context: ExecutionContext, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None, + ) -> ExecutionPathResult: + """Compose backend capabilities with outer Attention boundaries.""" + boundaries = set(context.outer_boundaries) + if self._uses_hsdp_compile_boundary(): + boundaries.add(OuterBoundary.HSDP) + active_strategy = self._get_active_parallel_strategy() + strategy_name = active_strategy.name + if self.use_ring and active_strategy.enabled and strategy_name == "ulysses": + parallel_strategy = ParallelStrategy.HYBRID_ULYSSES_RING + elif self.use_ring and active_strategy.enabled: + parallel_strategy = ParallelStrategy.RING + else: + parallel_strategy = { + "allgather_kv": ParallelStrategy.ALLGATHER_KV, + "ulysses": ParallelStrategy.ULYSSES, + }.get(strategy_name, ParallelStrategy.NONE) + resolved_context = replace( + context, + outer_boundaries=frozenset(boundaries), + paged_kv=self.is_paged_kv_active(), + parallel_strategy=parallel_strategy, + ) + resolver = getattr(self.attention, "resolve_execution_path", None) + if not callable(resolver): + return ExecutionPathResult.unmigrated( + type(self.attention).__name__, + resolved_context, + ) + if not resolved_context.paged_kv: + attn_metadata = self._with_kv_cache_dtype(attn_metadata) + return resolver( + resolved_context, + query, + key, + value, + attn_metadata, + ) + @torch.compiler.disable def _forward_hsdp_compile_boundary( self, @@ -422,10 +484,11 @@ class Attention(nn.Module): ) use_paged_attention = paged_adapter is not None and self.paged_kv_cache_role is not None if use_paged_attention and not getattr(self.attn_backend, "supports_paged_kv", False): - assert self.attn_backend is not None + backend_name = ( + self.attn_backend.get_name() if self.attn_backend is not None else type(self.attention).__name__ + ) raise NotImplementedError( - f"Diffusion paged KV requires an Omni backend with paged support; " - f"selected {self.attn_backend.get_name()}" + f"Diffusion paged KV requires an Omni backend with paged support; selected {backend_name}" ) if use_paged_attention and strategy is not self._no_parallel_strategy: strategy_name = strategy.name @@ -554,9 +617,11 @@ class Attention(nn.Module): self._scheduler_paged_kv and self.paged_kv_cache_role is not None and in_kv_memory_profile - and cast(type[AttentionBackend], self.attn_backend).get_name() == "FLASH_ATTN" + and self.attn_backend is not None + and self.attn_backend.get_name() == "FLASH_ATTN" and not current_omni_platform.supports_diffusion_dense_flash_attention() ): + assert self.sdpa_fallback is not None logger.warning_once( "The startup KV memory profile is using SDPA because dense FLASH_ATTN is unavailable. " "Formal paged requests still use the platform-native paged attention backend." From e84014dffb972c62ca35d83437d6248425fbad18 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Fri, 2 Oct 2026 04:34:30 -0700 Subject: [PATCH 11/22] [Model] Add single-stage Qwen3-TTS with in-Talker streaming codec decode (#8259) Signed-off-by: Sy03 <1370724210@qq.com> --- docs/configuration/stage_configs.md | 65 +- tests/config/test_config_merge.py | 158 ++++ tests/config/test_omni_config.py | 1 + .../test_omni_ar_scheduler_stale_drain.py | 102 ++- tests/engine/test_orchestrator_first_audio.py | 59 +- tests/engine/test_stage_engine_core_proc.py | 30 +- .../openai_api/test_audio_format.py | 16 + .../qwen3_tts/test_code_predictor_dtype.py | 23 + .../qwen3_tts/test_code_predictor_sampling.py | 119 +++ .../qwen3_tts/test_first_audio_graphs.py | 164 ++++ .../qwen3_tts/test_streaming_col2im_cuda.py | 57 ++ .../test_qwen3_tts_async_chunk.py | 34 + tests/worker_v2/test_index_staging.py | 21 + tests/worker_v2/test_omni_ar_model_runner.py | 359 ++++++++ tests/worker_v2/test_omni_model_state.py | 6 + vllm_omni/config/model.py | 6 +- vllm_omni/config/omni_config.py | 4 + vllm_omni/config/pipeline_registry.py | 3 +- vllm_omni/config/stage_config.py | 14 +- vllm_omni/core/sched/omni_ar_scheduler.py | 40 +- vllm_omni/core/sched/omni_scheduler_mixin.py | 2 + .../deploy/qwen3_tts_fused_single_gpu.yaml | 54 ++ vllm_omni/engine/arg_utils.py | 2 + vllm_omni/engine/orchestrator.py | 4 +- vllm_omni/engine/stage_engine_core_proc.py | 5 +- vllm_omni/entrypoints/openai/api_server.py | 30 +- .../entrypoints/openai/audio_utils_mixin.py | 21 +- .../openai/tts_adapters/qwen3_tts.py | 48 +- .../models/common/qwen3_code_predictor.py | 32 +- .../models/qwen3_tts/fused_code_predictor.py | 312 +++++++ .../models/qwen3_tts/pipeline.py | 30 + .../models/qwen3_tts/prompt_embeds_builder.py | 57 +- .../qwen3_tts_code_predictor_vllm.py | 64 ++ .../models/qwen3_tts/qwen3_tts_talker.py | 94 ++- .../models/qwen3_tts/stream_decode.py | 60 ++ .../tokenizer_12hz/streaming_decoder.py | 799 ++++++++++++++++++ .../stage_input_processors/qwen3_tts.py | 13 +- vllm_omni/worker_v2/model_states/eager_mtp.py | 337 ++++++++ .../model_states/eager_mtp_kernels.py | 126 +++ .../model_states/omni_model_state.py | 64 +- vllm_omni/worker_v2/omni_ar_model_runner.py | 53 +- vllm_omni/worker_v2/omni_model_runner.py | 13 + vllm_omni/worker_v2/streaming_audio.py | 154 ++++ 43 files changed, 3565 insertions(+), 90 deletions(-) create mode 100644 tests/model_executor/models/qwen3_tts/test_streaming_col2im_cuda.py create mode 100644 tests/worker_v2/test_index_staging.py create mode 100644 vllm_omni/deploy/qwen3_tts_fused_single_gpu.yaml create mode 100644 vllm_omni/model_executor/models/qwen3_tts/fused_code_predictor.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/stream_decode.py create mode 100644 vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/streaming_decoder.py create mode 100644 vllm_omni/worker_v2/model_states/eager_mtp_kernels.py create mode 100644 vllm_omni/worker_v2/streaming_audio.py diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 1085b4e64..2cf7327dd 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -319,7 +319,9 @@ ROCm and MUSA. Select one of these deployment profiles: | `qwen3_tts_mrv2.yaml` | V2 | B1 | Explicit MRV2 profile (same runner selection as the default) | | `qwen3_tts_high_concurrency_mrv2.yaml` | V2 | B1, B2 | Opt-in throughput tuning | | `qwen3_tts_high_concurrency_mrv2_b4.yaml` | V2 | B1, B2, B3, B4 | Experimental throughput / buffered playback | +| `qwen3_tts_high_concurrency_mrv2_single_gpu.yaml` | V2 | B1–B8 | Experimental two-stage deployment on one GPU, with MPS | | `qwen3_tts_high_concurrency.yaml` | V1 | Existing defaults | V1 high-concurrency control | +| `qwen3_tts_fused_single_gpu.yaml` | V2 | N/A (in-Talker decoder) | Opt-in single-stage CUDA pipeline on one GPU | ```bash # MRV2 is the default; pass a copy with `model_runner: v1` to force V1. @@ -350,12 +352,73 @@ supply `mtp_sampling_params` and `get_mtp_seed(sampling_params)` for model-local (with `mtp_sample_steps` and `mtp_sample_vocab_size`) as needed. Qwen3-TTS retains its existing `talker_mtp` entry point for V1. +### Selecting one or two Qwen3-TTS stages + +Choose the deployment profile explicitly with `--deploy-config`. The default +`qwen3_tts` pipeline retains two stages: the Talker publishes codec chunks to +Code2Wav, which produces audio. Its stage devices can be placed on the same +GPU or on separate GPUs. `qwen3_tts_fused` is a separate, opt-in pipeline for +one CUDA GPU: the Talker owns the stateful codec decoder and emits PCM directly, +without a Code2Wav engine or stage connector. + +```bash +# Two stages on one GPU; the two-stage pipeline remains independently selectable. +vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \ + --deploy-config vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml + +# One stage on one GPU; no inter-stage transport or Code2Wav scheduling. +vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \ + --deploy-config vllm_omni/deploy/qwen3_tts_fused_single_gpu.yaml +``` + +The single-stage profile requires CUDA, V2, asynchronous chunks, TP=PP=1, +an in-process worker and disabled prefix caching. Select the full profile; +adding `talker_stream_decode: true` to a two-stage Talker is rejected because +its latent output belongs to the Code2Wav input contract. Single-stage options +live in `stages[].additional_config`; two-stage transport and first-frame +options live in `connectors.*.extra`. + +The single-stage profile defaults `talker_stream_first_audio` to `false` for +throughput. To opt into a separate first-frame delivery queue, save this +overlay beside the bundled deployment files and pass its path to +`--deploy-config`: + +```yaml +base_config: qwen3_tts_fused_single_gpu.yaml +stages: + - stage_id: 0 + additional_config: + talker_stream_decode: true + talker_stream_first_audio: true + ref_code_context_frames: 72 + code_predictor_kv_cache: true + code_predictor_fused_sampling: true + code_predictor_fused: true +``` + +This queue uses PCM from the single-stage decoder. The two-stage +`talker_first_audio` option below uses its own first-frame decoder and does +not require single-stage execution. The two options cannot be enabled +together. Single-stage prefix-cache reset with +`reset_running_requests=True` returns `False` while requests are active; +wait for completion or abort them before retrying. + +Reference-cloning context is selected independently by each profile. The +single-stage profile explicitly uses 72 reference frames, matching the +`qwen3_tts.yaml` context boundary. A shorter +`additional_config.ref_code_context_frames` (for example 25) can change +timbre continuity even when WER is similar; +the throughput measurements for CustomVoice do not qualify Base voice +cloning quality. Reference layout, grouping and priming are model hooks; +the shared worker does not select Qwen's context length or codebook layout. + ### Qwen3-TTS first-frame delivery and rollback The standard `qwen3_tts.yaml`, `qwen3_tts_high_concurrency.yaml`, `qwen3_tts_mrv2.yaml` and `qwen3_tts_high_concurrency_mrv2.yaml` profiles enable the `talker_first_audio` connector option, as does the experimental -single-GPU profile. This changes the default CUDA MRV2 streaming path: +two-stage `qwen3_tts_high_concurrency_mrv2_single_gpu.yaml` profile. +This changes the default CUDA MRV2 streaming path: residual prediction runs eagerly after the Talker sample, the Talker loads an additional first-frame decoder with its weights and CUDA graphs, and the orchestrator orders its audio before subsequent Code2Wav chunks. diff --git a/tests/config/test_config_merge.py b/tests/config/test_config_merge.py index 0bcc15b0b..54b9f4816 100644 --- a/tests/config/test_config_merge.py +++ b/tests/config/test_config_merge.py @@ -35,3 +35,161 @@ def test_recursive_deploy_inheritance(tmp_path): merged = resolve_deploy_yaml(leaf) assert merged["speech_cache"] == {"nested": {"a": 1, "b": 0}} assert merged["async_chunk"] is False + + +def test_single_stage_streaming_capability_uses_central_async_resolver(): + from vllm_omni.config.stage_config import DeployConfig, _resolve_pipeline_async_chunk_enabled + from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_FUSED_PIPELINE + + assert _resolve_pipeline_async_chunk_enabled(QWEN3_TTS_FUSED_PIPELINE, DeployConfig()) + with pytest.raises(ValueError, match="requires async_chunk=True"): + _resolve_pipeline_async_chunk_enabled(QWEN3_TTS_FUSED_PIPELINE, DeployConfig(async_chunk=False)) + + +@pytest.mark.parametrize( + "pipeline_key,profile,stage_count", + [ + ("qwen3_tts", "qwen3_tts_high_concurrency_mrv2_single_gpu.yaml", 2), + ("qwen3_tts_fused", "qwen3_tts_fused_single_gpu.yaml", 1), + ], +) +def test_qwen3_tts_deployment_selection_preserves_output_and_transport_contract(pipeline_key, profile, stage_count): + from tests.helpers.stage_config import get_deploy_config_path + from vllm_omni.config.pipeline_registry import resolve_pipeline_config + from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy + + pipeline = resolve_pipeline_config(pipeline_key) + assert pipeline is not None + deploy = load_deploy_config(get_deploy_config_path(profile)) + stages = merge_pipeline_deploy(pipeline, deploy) + assert len(stages) == stage_count + talker = stages[0] + if stage_count == 1: + assert talker.final_output and talker.final_output_type == "audio" + assert talker.yaml_engine_args["engine_output_type"] == "audio" + assert "custom_process_next_stage_input_func" not in talker.yaml_engine_args + assert not talker.yaml_extras.get("output_connectors") + assert talker.yaml_engine_args["supports_running_prefix_cache_reset"] is False + assert talker.yaml_engine_args["additional_config"]["ref_code_context_frames"] == 72 + else: + assert not talker.final_output + assert talker.yaml_engine_args["engine_output_type"] == "latent" + assert talker.yaml_engine_args["custom_process_next_stage_input_func"].endswith("talker2code2wav_async_chunk") + assert talker.yaml_extras["output_connectors"] + assert stages[1].input_sources == [0] and stages[1].final_output_type == "audio" + assert talker.yaml_engine_args.get("supports_running_prefix_cache_reset", True) + + +def test_single_stage_running_reset_capability_cannot_be_overridden(): + from vllm_omni.config.omni_config import VllmOmniConfig + from vllm_omni.config.stage_config import DeployConfig, StageDeployConfig, merge_pipeline_deploy + from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_FUSED_PIPELINE + + deploy = DeployConfig( + stages=[StageDeployConfig(stage_id=0, engine_extras={"supports_running_prefix_cache_reset": True})] + ) + legacy = merge_pipeline_deploy(QWEN3_TTS_FUSED_PIPELINE, deploy)[0] + legacy.runtime_overrides["supports_running_prefix_cache_reset"] = True + assert legacy.to_omegaconf().engine_args.supports_running_prefix_cache_reset is False + with pytest.raises(ValueError, match="no structured config owner: supports_running_prefix_cache_reset"): + VllmOmniConfig.from_pipeline_config(QWEN3_TTS_FUSED_PIPELINE, user_deploy_config=deploy) + structured = VllmOmniConfig.from_pipeline_config(QWEN3_TTS_FUSED_PIPELINE) + assert structured.stage_by_id(0).model_config.supports_running_prefix_cache_reset is False + + +@pytest.fixture +def single_stage_tts_config(mocker): + from vllm.config import VllmConfig + + config = mocker.Mock(spec=VllmConfig) + config.additional_config = {"talker_stream_decode": True} + config.model_config = mocker.Mock( + stage_connector_config={}, + engine_output_type="audio", + final_output=True, + supports_running_prefix_cache_reset=False, + use_v2_model_runner=True, + async_chunk=True, + ) + config.parallel_config = mocker.Mock( + tensor_parallel_size=1, pipeline_parallel_size=1, distributed_executor_backend="uni" + ) + config.cache_config = mocker.Mock(enable_prefix_caching=False) + config.device_config = mocker.Mock(device="cuda") + mocker.patch( + "vllm_omni.model_executor.models.qwen3_tts.stream_decode.current_omni_platform.is_cuda", return_value=True + ) + return config + + +def test_single_stage_tts_rejects_unsupported_decoder_configuration(single_stage_tts_config): + from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled + + config = single_stage_tts_config + config.cache_config.enable_prefix_caching = True + with pytest.raises(ValueError, match="enable_prefix_caching=False"): + talker_stream_decode_enabled(config) + config.cache_config.enable_prefix_caching = False + assert talker_stream_decode_enabled(config) + + +@pytest.mark.parametrize("option_location", ["additional_config", "connector"]) +@pytest.mark.parametrize("output_type", ["latent", "audio"]) +def test_two_stage_talker_rejects_single_stage_pcm_option(single_stage_tts_config, option_location, output_type): + from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled + + config = single_stage_tts_config + config.model_config.engine_output_type = output_type + config.model_config.final_output = False + if option_location == "connector": + config.model_config.stage_connector_config = {"extra": config.additional_config} + config.additional_config = {} + with pytest.raises(ValueError, match="final audio-output Talker"): + talker_stream_decode_enabled(config) + + +def test_two_stage_first_audio_remains_independent(single_stage_tts_config): + from vllm_omni.model_executor.models.qwen3_tts.first_audio import talker_first_audio_enabled + from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled + + config = single_stage_tts_config + config.model_config.engine_output_type = "latent" + config.additional_config = {} + config.model_config.stage_connector_config = {"extra": {"talker_first_audio": True}} + assert not talker_stream_decode_enabled(config) + assert talker_first_audio_enabled(config) + config.model_config.stage_connector_config["extra"]["talker_stream_first_audio"] = True + with pytest.raises(ValueError, match="requires the single-stage"): + talker_stream_decode_enabled(config) + + +def test_single_stage_pcm_requires_cuda_device_and_own_first_audio_option(single_stage_tts_config): + from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled + + config = single_stage_tts_config + config.device_config.device = "cpu" + with pytest.raises(ValueError, match="requires CUDA"): + talker_stream_decode_enabled(config) + config.device_config.device = "cuda" + config.additional_config["talker_first_audio"] = True + with pytest.raises(ValueError, match="mutually exclusive"): + talker_stream_decode_enabled(config) + config.additional_config.pop("talker_first_audio") + config.additional_config["talker_stream_first_audio"] = True + assert talker_stream_decode_enabled(config) + + +def test_single_stage_audio_talker_requires_pcm_option(single_stage_tts_config): + from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled + + single_stage_tts_config.additional_config = {} + with pytest.raises(ValueError, match="requires talker_stream_decode=True"): + talker_stream_decode_enabled(single_stage_tts_config) + + +def test_single_stage_pcm_requires_reset_capability_declaration(single_stage_tts_config): + from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled + + single_stage_tts_config.model_config.supports_running_prefix_cache_reset = True + with pytest.raises(ValueError, match="must disable running prefix-cache reset"): + talker_stream_decode_enabled(single_stage_tts_config) diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index b19c97144..e99da6ceb 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -1046,6 +1046,7 @@ def test_sub_config_fields_match_structured_scopes(): "interleave_mm_strings", "media_io_kwargs", "final_output", + "supports_running_prefix_cache_reset", "active_stream_window", "session_mode", "duplex_max_sessions", diff --git a/tests/core/sched/test_omni_ar_scheduler_stale_drain.py b/tests/core/sched/test_omni_ar_scheduler_stale_drain.py index 182857ea9..b4eb839a2 100644 --- a/tests/core/sched/test_omni_ar_scheduler_stale_drain.py +++ b/tests/core/sched/test_omni_ar_scheduler_stale_drain.py @@ -106,7 +106,7 @@ def _make_drain_sched(session: Request) -> MagicMock: return sched -def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: int) -> bool: +def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: int, spec=None) -> bool: """Feed one model-runner frame through update_from_output. Returns True when the frame's tokens were delivered (the append ran), @@ -114,7 +114,7 @@ def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: """ scheduler_output = MagicMock(spec=SchedulerOutput) scheduler_output.num_scheduled_tokens = {session.request_id: num_scheduled} - scheduler_output.scheduled_spec_decode_tokens = {} + scheduler_output.scheduled_spec_decode_tokens = {session.request_id: spec} if spec else {} scheduler_output.num_invalid_spec_tokens = 0 model_runner_output = MagicMock(spec=ModelRunnerOutput) @@ -134,6 +134,27 @@ def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: return sched._update_request_with_output.called +def test_connector_prompt_replacement_drops_old_frame_and_delivers_new_frame() -> None: + from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin + + session = _make_session() + session.num_computed_tokens = 6 + session.num_in_flight_tokens = 1 + session.num_output_placeholders = 1 + sched = _make_drain_sched(session) + sched.chunk_transfer_adapter.replaced_streaming_prompt_ids = {session.request_id} + sched.chunk_transfer_adapter.requests_with_ready_chunks = {session.request_id} + sched.chunk_transfer_adapter.requests_num_chunks_sent = {session.external_req_id: 2} + + OmniSchedulerMixin._reset_ready_async_chunk_replacements(sched) + assert session.num_stale_output_tokens == 1 and session.drop_stale_output + assert session.num_output_placeholders == 0 + sched._release_replaced_streaming_prompt_cache.assert_called_once_with(session) + assert _run_step(sched, session, num_scheduled=1, token=42) is False + assert session.num_stale_output_tokens == 0 + assert _run_step(sched, session, num_scheduled=1, token=43) is True + + def test_exact_drain_delivers_new_segment_frame() -> None: """One in-flight decode frame at replacement: its late output drains the counter exactly and the new segment's first frame is delivered.""" @@ -225,3 +246,80 @@ def test_in_flight_prefill_chunk_drains_exactly_without_underflow() -> None: assert _run_step(sched, session, num_scheduled=5, token=42) is False # late prefill chunk dropped assert session.num_stale_output_tokens == 0 assert _run_step(sched, session, num_scheduled=1, token=43) is True # new segment survives + + +@pytest.mark.parametrize("drop", [False, True]) +def test_preemption_stale_output_follows_upstream_delivery_policy(drop): + request = _make_session() + request.status = RequestStatus.PREEMPTED + request.num_stale_output_tokens = 1 + request.num_in_flight_tokens = 1 + request.drop_stale_output = drop + sched = _make_drain_sched(request) + assert _run_step(sched, request, num_scheduled=1, token=42) is (not drop) + assert request.num_stale_output_tokens == 0 + if not drop: + sched._update_request_with_output.assert_called_once_with(request, [42], is_stale=True) + + +def test_preemption_stale_spec_rejection_does_not_roll_back_resumed_counters(): + request = _make_session() + request.num_stale_output_tokens = 2 + request.num_in_flight_tokens = 2 + request.num_computed_tokens = 3 + request.num_output_placeholders = 2 + request.drop_stale_output = False + sched = _make_drain_sched(request) + assert _run_step(sched, request, num_scheduled=2, token=42, spec=[41]) + assert request.num_computed_tokens == 3 + assert request.num_output_placeholders == 2 + + +@pytest.mark.parametrize( + "arch,supports_reset,running,reset_running,blocked", + [ + ("Qwen3TTSTalkerForConditionalGeneration", False, True, True, True), + ("Qwen3TTSTalkerForConditionalGeneration", False, False, True, False), + ("Qwen3TTSTalkerForConditionalGeneration", False, True, False, False), + ("Qwen3TTSTalkerForConditionalGeneration", True, True, True, False), + ("AnotherStatefulAudioModel", False, True, True, True), + ("AnotherStatefulAudioModel", True, True, True, False), + ], +) +def test_reset_running_streaming_codec_is_rejected_before_preemption( + mocker, arch, supports_reset, running, reset_running, blocked +): + from vllm.config import VllmConfig + from vllm.v1.core.sched.scheduler import Scheduler + + from vllm_omni.config.model import OmniModelConfig + + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler.vllm_config = mocker.Mock( + spec=VllmConfig, + model_config=mocker.Mock( + spec=OmniModelConfig, + model_arch=arch, + supports_running_prefix_cache_reset=supports_reset, + ), + ) + request = _make_session() + request.num_in_flight_tokens = 1 + request.num_output_placeholders = 1 + scheduler.running = [request] if running else [] + + def reset_cache(_self, reset_running_requests=False, reset_connector=False): + assert reset_running_requests is reset_running and reset_connector is True + return True + + reset = mocker.patch.object(Scheduler, "reset_prefix_cache", autospec=True, side_effect=reset_cache) + + assert scheduler.reset_prefix_cache(reset_running, reset_connector=True) is (not blocked) + if blocked: + reset.assert_not_called() + assert scheduler.running == [request] + assert request.status == RequestStatus.RUNNING + assert request.num_in_flight_tokens == request.num_output_placeholders == 1 + assert not request.drop_stale_output + else: + reset.assert_called_once() diff --git a/tests/engine/test_orchestrator_first_audio.py b/tests/engine/test_orchestrator_first_audio.py index 5e467254b..bc99cf7fb 100644 --- a/tests/engine/test_orchestrator_first_audio.py +++ b/tests/engine/test_orchestrator_first_audio.py @@ -114,10 +114,10 @@ def _raw(*, first=False, required=False, terminal=False, samples=(1, 1, 1, 1)): ) -async def _codec_output(obj, raw): - replica = obj.stage_pools[1].replica_id - processed = await obj._process_llm_stage_outputs(1, replica, raw, set()) - await obj._handle_processed_outputs(1, replica, processed) +async def _codec_output(obj, raw, stage=1): + replica = obj.stage_pools[stage].replica_id + processed = await obj._process_llm_stage_outputs(stage, replica, raw, set()) + await obj._handle_processed_outputs(stage, replica, processed) def _messages(obj): @@ -132,6 +132,57 @@ def _audio(output): return torch.cat(audio) if isinstance(audio, list) else audio +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", list(RequestOutputKind)) +@pytest.mark.parametrize("suffix_overtook", [False, True]) +async def test_single_stage_first_audio_obeys_output_kind_and_terminal_order(kind, suffix_overtook): + obj = _orchestrator(kind, final_stage_id=0) + obj.stage_pools = [obj.stage_pools[1]] + suffix = _raw(required=True, samples=(20, 21)) + terminal = _raw(required=True, terminal=True, samples=(30, 31)) + if suffix_overtook: + await _codec_output(obj, suffix, stage=0) + await _codec_output(obj, terminal, stage=0) + assert obj.output_async_queue.empty() + await obj._route_upstream_first_audio(0, 2, _raw(first=True, samples=(10, 11))) + if not suffix_overtook: + await _codec_output(obj, suffix, stage=0) + await _codec_output(obj, terminal, stage=0) + messages = _messages(obj) + expected = { + RequestOutputKind.DELTA: [[10, 11], [20, 21], [30, 31]], + RequestOutputKind.CUMULATIVE: [[10, 11], [10, 11, 20, 21], [10, 11, 20, 21, 30, 31]], + RequestOutputKind.FINAL_ONLY: [[10, 11, 20, 21, 30, 31]], + } + assert [_audio(message.engine_outputs).tolist() for message in messages] == expected[kind] + assert all((message.stage_id, message.replica_id) == (0, 2) for message in messages) + assert messages[-1].finished + assert not obj.request_states["r"].pending_first_audio_outputs + + +@pytest.mark.asyncio +async def test_single_stage_cancellation_discards_overtaking_terminal_and_late_first_audio(): + obj = _orchestrator(final_stage_id=0) + obj.stage_pools = [obj.stage_pools[1]] + await _codec_output(obj, _raw(required=True, terminal=True), stage=0) + assert obj.request_states["r"].pending_first_audio_outputs + await obj._cleanup_request_ids(["r"], abort=True) + await obj._route_upstream_first_audio(0, 2, _raw(first=True)) + assert not obj.request_states and obj.output_async_queue.empty() + + +@pytest.mark.asyncio +async def test_single_stage_first_audio_processing_failure_releases_waiting_terminal(): + obj = _orchestrator(final_stage_id=0) + obj.stage_pools = [obj.stage_pools[1]] + await _codec_output(obj, _raw(required=True, terminal=True), stage=0) + obj.stage_pools[0].process_llm_raw_outputs = AsyncMock(side_effect=RuntimeError("processing failed")) + await obj._route_upstream_first_audio(0, 2, _raw(first=True)) + error = obj.output_async_queue.get_nowait() + assert isinstance(error, ErrorMessage) and error.stage_id == 0 + assert not obj.request_states + + @pytest.mark.asyncio @pytest.mark.parametrize("kind", list(RequestOutputKind)) @pytest.mark.parametrize("codec_overtook", [False, True]) diff --git a/tests/engine/test_stage_engine_core_proc.py b/tests/engine/test_stage_engine_core_proc.py index a7104b14d..217f13e53 100644 --- a/tests/engine/test_stage_engine_core_proc.py +++ b/tests/engine/test_stage_engine_core_proc.py @@ -1,9 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from types import SimpleNamespace from unittest.mock import patch +import pytest from vllm.v1.engine.core import EngineCoreProc +from vllm.v1.executor.uniproc_executor import UniProcExecutor -from vllm_omni.engine.stage_engine_core_proc import StageEngineCoreProc +from vllm_omni.engine.stage_engine_core_proc import StageEngineCoreProc, _bind_first_audio_sink + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize( + "first_decoder,stream_decoder,stream_first_audio,expected", + [(True, False, False, True), (False, True, False, False), (False, True, True, True), (False, False, True, False)], +) +def test_first_audio_sink_preserves_first_decoder_and_opts_in_stream_decoder( + mocker, + first_decoder, + stream_decoder, + stream_first_audio, + expected, +): + executor = UniProcExecutor.__new__(UniProcExecutor) + executor.driver_worker = mocker.Mock() + runner = executor.driver_worker.worker.model_runner + runner.model.first_frame_decoder = object() if first_decoder else None + runner.model.stream_decoder = object() if stream_decoder else None + runner.model.stream_first_audio = stream_first_audio + assert _bind_first_audio_sink(executor, mocker.Mock(), mocker.Mock()) is expected + assert runner.model_state.set_first_audio_sink.call_count == int(expected) def test_preprocess_add_request_preserves_omni_fields(): diff --git a/tests/entrypoints/openai_api/test_audio_format.py b/tests/entrypoints/openai_api/test_audio_format.py index 2fb76e3f6..3eda688c0 100644 --- a/tests/entrypoints/openai_api/test_audio_format.py +++ b/tests/entrypoints/openai_api/test_audio_format.py @@ -346,3 +346,19 @@ def test_chat_audio_metadata_survives_response_serialization(stream, fmt): assert audio_file.samplerate == 22050 else: assert len(raw) == frame_count * 2 + + +def test_float32_pcm_fast_path_matches_soundfile(): + from vllm_omni.entrypoints.openai.audio_utils_mixin import _float32_to_pcm16_bytes + + rng = np.random.default_rng(0) + audio = np.concatenate( + [ + rng.uniform(-1.2, 1.2, 4096), + np.array([-1.0, 1.0, 0.0, -0.0, 32767 / 32768, -1 / 32768, 0.5 / 32768, -0.5 / 32768]), + ] + ).astype(np.float32) + with BytesIO() as buffer: + soundfile.write(buffer, audio, 24000, format="RAW", subtype="PCM_16") + expected = buffer.getvalue() + assert _float32_to_pcm16_bytes(audio) == expected diff --git a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py index 52712a29c..8df20d87c 100644 --- a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py +++ b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py @@ -569,6 +569,27 @@ class TestCodePredictorPerRowGenerators: ) +@pytest.mark.parametrize("fused_requested", [False, True]) +@pytest.mark.parametrize("invalid_input", ["uniforms", "generators"]) +def test_sampling_contract_rejected_before_dispatch(mocker, loaded_target_classes, fused_requested, invalid_input): + predictor, talker_config = TestCodePredictorPerRowGenerators()._make_predictor(mocker, loaded_target_classes) + predictor._fused_requested = fused_requested + predictor._fused = mocker.Mock() + setup = mocker.patch.object(predictor, "_setup_compile") + uniforms = torch.full((2, 3, 1 if invalid_input == "uniforms" else 64), 0.5) + generators = [torch.Generator()] if invalid_input == "generators" else None + with pytest.raises(ValueError, match="sample_uniforms must have shape|one entry per row"): + predictor( + torch.zeros(2, dtype=torch.long), + torch.randn(2, talker_config.hidden_size), + torch.randn(2, talker_config.hidden_size), + sample_uniforms=uniforms, + generators=generators, + ) + setup.assert_not_called() + predictor._fused.assert_not_called() + + class TestCodePredictorModelDtype: """Test the inner model forward with different dtypes.""" @@ -663,6 +684,8 @@ class TestCodePredictorGraphReplay: predictor = object.__new__(code_predictor_wrapper) torch.nn.Module.__init__(predictor) + predictor._fused_requested = False + predictor._fused = None predictor._num_groups = 3 predictor._model_dtype = torch.float32 predictor._setup_compile = mocker.Mock() diff --git a/tests/model_executor/models/qwen3_tts/test_code_predictor_sampling.py b/tests/model_executor/models/qwen3_tts/test_code_predictor_sampling.py index 30994e36b..9fa3a76bb 100644 --- a/tests/model_executor/models/qwen3_tts/test_code_predictor_sampling.py +++ b/tests/model_executor/models/qwen3_tts/test_code_predictor_sampling.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Supplied-uniform samples must match the existing ATen path exactly.""" +from types import SimpleNamespace + import pytest import torch @@ -12,6 +14,79 @@ pytestmark = [ ] +@pytest.mark.parametrize("top_k", [0, 1, 50, 2048]) +@torch.inference_mode() +def test_fused_sample_kernel_preserves_signed_zero_cutoff_ties(top_k): + from vllm_omni.model_executor.models.qwen3_tts.fused_code_predictor import _cp_sample_kernel + + vocab, hidden, groups = 2048, 64, 4 + logits = torch.zeros(2, vocab, device="cuda", dtype=torch.bfloat16) + logits[0, 1] = -0.0 + logits[1, ::2] = -0.0 + uniforms = torch.full((2, groups - 1, vocab), 0.01, device="cuda")[:, 1, :] + uniforms[0, 1] = uniforms[1, 2] = 0.99 + codes = torch.full((2, groups), -1, device="cuda", dtype=torch.int64) + table = torch.randn(vocab, hidden, device="cuda", dtype=torch.bfloat16) + next_input = torch.empty(2, hidden, device="cuda", dtype=table.dtype) + + def sample(): + _cp_sample_kernel[(2,)]( + logits, uniforms, uniforms.stride(0), codes, 1, groups, table, next_input, 1.0, + V=vocab, TOPK=top_k, HID=hidden, HAS_NEXT=True, num_warps=4, + ) # fmt: skip + + scaled = logits.float() + if top_k > 0: + scaled = scaled.masked_fill(scaled < scaled.topk(top_k).values[:, -1:], -torch.inf) + expected = (scaled - torch.log(-torch.log(uniforms))).argmax(-1) + assert expected.tolist() == [1, 2] + sample() + torch.testing.assert_close(codes[:, 1], expected) + torch.testing.assert_close(next_input, table[expected]) + assert (codes[:, [0, 2, 3]] == -1).all() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + sample() + codes.fill_(-1) + graph.replay() + torch.testing.assert_close(codes[:, 1], expected) + torch.testing.assert_close(next_input, table[expected]) + + +@torch.inference_mode() +def test_fused_predictor_warmup_with_capacity_one(): + from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import ( + Qwen3TTSTalkerCodePredictorConfig, + Qwen3TTSTalkerConfig, + ) + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import ( + Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM, + ) + + cp = Qwen3TTSTalkerCodePredictorConfig( + vocab_size=64, hidden_size=64, intermediate_size=128, num_hidden_layers=1, + num_attention_heads=4, num_key_value_heads=2, head_dim=16, num_code_groups=4, + ) # fmt: skip + talker = Qwen3TTSTalkerConfig(hidden_size=64, num_code_groups=4) + config = SimpleNamespace( + model_config=SimpleNamespace(stage_connector_config={}), + additional_config={"code_predictor_kv_cache": True, "code_predictor_fused": True}, + scheduler_config=SimpleNamespace(max_num_seqs=1), + ) + predictor = Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM( + vllm_config=config, config=cp, talker_config=talker + ).to(device="cuda", dtype=torch.bfloat16) + predictor._setup_compile() + assert predictor._fused is not None + hidden = torch.randn(1, 1, 64, device="cuda", dtype=torch.bfloat16) + result = predictor( + torch.zeros(1, 1, device="cuda", dtype=torch.long), hidden, hidden, + sample_uniforms=torch.full((1, 3, 64), 0.5, device="cuda"), + ) # fmt: skip + assert result.shape == (1, 4) + assert ((result >= 0) & (result < 64)).all() + + @pytest.mark.parametrize("batch", [1, 64]) @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) @pytest.mark.parametrize("top_k", [-1, 1, 50, 2048]) @@ -66,3 +141,47 @@ def test_topk_gumbel_accepts_strided_supplied_uniforms(top_k): expected = (scaled.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True) actual = sample_code_topk_gumbel(logits, uniforms, top_k, 1 / 0.9) torch.testing.assert_close(actual, expected) + + +@pytest.mark.parametrize("batch", [1, 4]) +@pytest.mark.parametrize("top_k", [1, 50]) +@torch.inference_mode() +def test_fused_predictor_matches_full_reference_with_fixed_uniforms(mocker, batch, top_k): + from vllm.config import VllmConfig + + from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import ( + Qwen3TTSTalkerCodePredictorConfig, + Qwen3TTSTalkerConfig, + ) + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import ( + Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM, + ) + + torch.manual_seed(42) + cp = Qwen3TTSTalkerCodePredictorConfig( + vocab_size=64, hidden_size=64, intermediate_size=128, num_hidden_layers=1, + num_attention_heads=4, num_key_value_heads=2, head_dim=16, num_code_groups=4, + ) # fmt: skip + config = mocker.Mock( + spec=VllmConfig, + model_config=mocker.Mock(stage_connector_config={}), + additional_config={"code_predictor_kv_cache": True, "code_predictor_fused": True}, + scheduler_config=mocker.Mock(max_num_seqs=4), + ) + predictor = Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM( + vllm_config=config, config=cp, talker_config=Qwen3TTSTalkerConfig(hidden_size=64, num_code_groups=4) + ).to(device="cuda", dtype=torch.bfloat16) + predictor._setup_compile() + assert predictor._fused is not None + ids = torch.randint(0, 64, (batch, 1), device="cuda") + embeds = torch.randn(batch, 1, 64, device="cuda", dtype=torch.bfloat16) + hidden = torch.randn_like(embeds) + uniforms = torch.rand(batch, 3, 64, device="cuda").clamp_(1e-6, 1 - 1e-6) + actual = predictor(ids, embeds, hidden, top_k=top_k, sample_uniforms=uniforms).clone() + # Use identical parameters through the full re-prefill and ATen sampling + # path, bypassing both the fused predictor and its frame-local KV cache. + predictor._fused_requested = False + predictor._fused = None + predictor._frame_cache = None + expected = predictor(ids, embeds, hidden, top_k=top_k, sample_uniforms=uniforms) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/tests/model_executor/models/qwen3_tts/test_first_audio_graphs.py b/tests/model_executor/models/qwen3_tts/test_first_audio_graphs.py index ce92f06a6..66fb9c10f 100644 --- a/tests/model_executor/models/qwen3_tts/test_first_audio_graphs.py +++ b/tests/model_executor/models/qwen3_tts/test_first_audio_graphs.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Prefix graphs must follow per-request first-audio delivery, not a capture hint.""" +import math from types import SimpleNamespace import pytest @@ -18,6 +19,38 @@ pytestmark = [pytest.mark.core_model] DEVICE = torch.device("cuda:0") +@pytest.mark.cpu +@pytest.mark.parametrize("capacity,model_type", [(1, "base"), (128, "base"), (128, "custom_voice")]) +def test_stream_priming_graph_workspace_respects_token_budget(monkeypatch, capacity, model_type): + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration + from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz import streaming_decoder + + captures = [] + + def capture(stream, sizes, frames=1): + captures.append((sizes, frames)) + return SimpleNamespace(sizes=sizes, frames=frames) + + monkeypatch.setattr(streaming_decoder, "StreamingDecodeGraphs", capture) + model = SimpleNamespace( + stream_decoder=SimpleNamespace(max_frames=33), + stream_graphs=None, + stream_prime_graphs=None, + stream_ref_context_frames=25, + stream_chunk_frames=25, + config=SimpleNamespace(tts_model_type=model_type), + vllm_config=SimpleNamespace(scheduler_config=SimpleNamespace(max_num_batched_tokens=512)), + ) + sizes = [size for size in [1, 2, 4, 8, 16, 32, 64, 128] if size <= capacity] + Qwen3TTSTalkerForConditionalGeneration.capture_stream_decode_graphs(model, sizes) + expected = [(sizes, 1)] + if model_type == "base": + expected.append(([size for size in sizes if size * 25 <= 512], 25)) + assert captures == expected + Qwen3TTSTalkerForConditionalGeneration.capture_stream_decode_graphs(model, sizes) + assert captures == expected + + @pytest.mark.cpu @pytest.mark.parametrize("skip_flags", [(False, False), (True, True), (False, True), (True, False)]) @pytest.mark.parametrize("initial_frames", [1, 3]) @@ -124,3 +157,134 @@ def test_xvec_graphs_follow_request_first_audio_flags(initial_frames, time_major torch.testing.assert_close(got, want, atol=1e-4, rtol=1e-4) for key in ("ref_hidden", "ref_conv", "prefix_hidden", "suffix_quantized", "suffix_conv"): torch.testing.assert_close(cache[key], expected_cache[key], atol=1e-5, rtol=1e-5) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@torch.inference_mode() +def test_stream_decoder_preemption_restores_all_state_into_a_different_slot(): + from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder + from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import StreamingCodecDecoder + + decoder = _make_decoder().to(device=DEVICE, dtype=torch.bfloat16) + for name, param in decoder.named_parameters(): + if name.endswith("embedding_sum"): + param.normal_() + decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads + stream = StreamingCodecDecoder(decoder, num_slots=2, dtype=torch.bfloat16) + codes = torch.randint(0, 32, (1, 8, 2), device=DEVICE) + slot = torch.tensor([0], device=DEVICE, dtype=torch.int32) + for pos in range(7): + stream(codes[:, pos : pos + 1], slot, torch.tensor([pos], device=DEVICE, dtype=torch.int32)) + saved = stream.save_slot(0) + expected = stream(codes[:, 7:], slot, torch.tensor([7], device=DEVICE, dtype=torch.int32)).clone() + # Both the freed slot and the destination can hold an unrelated request. + other_codes = (codes[:, :3] + 1) % 32 + for pos in range(3): + stream( + other_codes[:, pos : pos + 1].expand(2, -1, -1).contiguous(), + torch.tensor([0, 1], device=DEVICE, dtype=torch.int32), + torch.full((2,), pos, device=DEVICE, dtype=torch.int32), + ) + without_restore = stream( + codes[:, 7:], + torch.tensor([1], device=DEVICE, dtype=torch.int32), + torch.tensor([7], device=DEVICE, dtype=torch.int32), + ).clone() + assert not torch.equal(without_restore, expected) + stream.restore_slot(1, saved) + actual = stream( + codes[:, 7:], + torch.tensor([1], device=DEVICE, dtype=torch.int32), + torch.tensor([7], device=DEVICE, dtype=torch.int32), + ) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("use_graph", [False, True]) +@torch.inference_mode() +def test_batched_reference_priming_matches_serial_decoder_continuation(use_graph): + from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration + from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import ( + StreamingCodecDecoder, + StreamingDecodeGraphs, + ) + + decoder = _make_decoder().to(device=DEVICE, dtype=torch.bfloat16) + for name, param in decoder.named_parameters(): + if name.endswith("embedding_sum"): + param.normal_() + decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads + stream = StreamingCodecDecoder(decoder, num_slots=5, dtype=torch.bfloat16) + graphs = StreamingDecodeGraphs(stream, [2], frames=25) if use_graph else None + model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration) + torch.nn.Module.__init__(model) + model.stream_decoder = stream + model.stream_prime_graphs = graphs + model.stream_chunk_frames = 25 + refs = [torch.randint(0, 32, (length, 2), device=DEVICE) for length in [25, 3, 51, 25]] + slots = [2, 0, 3, 4] + frame = torch.randint(0, 32, (4, 1, 2), device=DEVICE) + slot_tensor = torch.tensor(slots, device=DEVICE, dtype=torch.int32) + positions = torch.tensor([len(ref) for ref in refs], device=DEVICE, dtype=torch.int32) + for slot, ref in zip(slots, refs, strict=True): + for t0 in range(0, len(ref), 25): + stream( + ref[None, t0 : t0 + 25].to(torch.int32).contiguous(), + torch.tensor([slot], device=DEVICE, dtype=torch.int32), + torch.tensor([t0], device=DEVICE, dtype=torch.int32), + ) + expected = stream(frame, slot_tensor, positions).clone() + untouched = stream.save_slot(1) + # Reuse dirty slots; pos=0 must reset every causal layer for each group. + model.prime_stream_decoder(list(zip(slots, refs, strict=True))) + actual = stream(frame, slot_tensor, positions) + torch.testing.assert_close(actual, expected, rtol=0, atol=3e-5) + for got, want in zip(stream.save_slot(1), untouched, strict=True): + torch.testing.assert_close(got, want, rtol=0, atol=0) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("use_graph", [False, True]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@torch.inference_mode() +def test_stream_frames_match_fp32_exact_decoder(use_graph, dtype): + from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder + from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import ( + StreamingCodecDecoder, + StreamingDecodeGraphs, + ) + + decoder = _make_decoder().to(device=DEVICE, dtype=dtype) + # The tiny HF default initialization shrinks the deep conv stack to + # near-zero PCM. Keep the oracle signal large enough to catch lost state. + for module in decoder.modules(): + if isinstance(module, (torch.nn.Conv1d, torch.nn.ConvTranspose1d, torch.nn.Linear)): + torch.nn.init.kaiming_uniform_(module.weight, a=math.sqrt(5)) + if module.bias is not None: + module.bias.uniform_(-0.01, 0.01) + for name, param in decoder.named_parameters(): + if name.endswith("embedding_sum"): + param.normal_() + decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads + reference = _make_decoder().to(device=DEVICE, dtype=torch.float32) + reference.load_state_dict(decoder.state_dict()) + reference.precompute_snake_caches() + # Cross both the attention window and KV ring wrap, checking every frame. + codes = torch.randint(0, 32, (2, 136, 2), device=DEVICE) + expected = reference._forward_exact(codes.transpose(1, 2))[:, 0] + assert expected.square().sum() > 1e-3 + stream = StreamingCodecDecoder(decoder, num_slots=3, dtype=dtype) + decode = StreamingDecodeGraphs(stream, [2]) if use_graph else stream + slots = torch.tensor([2, 0], device=DEVICE, dtype=torch.int32) + spf = int(stream.spf) + for frame in range(codes.shape[1]): + positions = torch.full((2,), frame, device=DEVICE, dtype=torch.int32) + actual = decode(codes[:, frame : frame + 1].contiguous().to(torch.int32), slots, positions).reshape(2, -1) + exact_frame = expected[:, frame * spf : (frame + 1) * spf] + relative_rms = (actual - exact_frame).square().mean(1).sqrt() / exact_frame.square().mean(1).sqrt() + # BF16 kernels are approximate; compare every frame to the independent + # FP32 full decoder, not only one streaming implementation to itself. + tolerance = 0.05 if dtype == torch.bfloat16 else 5e-4 + assert torch.all(relative_rms < tolerance), (frame, relative_rms) diff --git a/tests/model_executor/models/qwen3_tts/test_streaming_col2im_cuda.py b/tests/model_executor/models/qwen3_tts/test_streaming_col2im_cuda.py new file mode 100644 index 000000000..c7f677821 --- /dev/null +++ b/tests/model_executor/models/qwen3_tts/test_streaming_col2im_cuda.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Streaming col2im preserves PCM inputs and ping-pong slot state.""" + +import pytest +import torch +from vllm.triton_utils import triton + +from tests.helpers.mark import hardware_test +from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import _col2im_kernel + +pytestmark = [pytest.mark.core_model, pytest.mark.tts] + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("batch,frames,channels,rate", [(1, 1, 5, 2), (3, 3, 129, 3), (9, 1, 64, 4)]) +def test_col2im_partial_tiles_and_reordered_slot_state(batch: int, frames: int, channels: int, rate: int): + torch.manual_seed(23) + capacity = batch + 2 + z = torch.randn(batch, frames, 2 * rate, channels).to(torch.bfloat16) + bias = torch.randn(channels).to(torch.bfloat16) + state = torch.randn(2, capacity, rate, channels).to(torch.bfloat16) + expected_state = state.clone() + slots = torch.arange(batch - 1, -1, -1, dtype=torch.int32) + positions = torch.arange(batch, dtype=torch.int32) + parity = torch.arange(capacity, dtype=torch.int32) % 2 + expected = torch.empty(batch, frames, rate, channels, dtype=torch.bfloat16) + for b in range(batch): + slot = int(slots[b]) + side = int(parity[slot]) + for t in range(frames): + previous = z[b, t - 1, rate:] if t else state[side, slot] if positions[b] else torch.zeros(rate, channels) + expected[b, t] = (z[b, t, :rate].float() + previous.float() + bias.float()).to(torch.bfloat16) + expected_state[1 - side, slot] = z[b, -1, rate:] + + z_gpu, bias_gpu, state_gpu = z.cuda(), bias.cuda(), state.cuda() + slots_gpu, positions_gpu, parity_gpu = slots.cuda(), positions.cuda(), parity.cuda() + actual = torch.empty_like(expected, device="cuda") + width = min(128, triton.next_power_of_2(channels)) + rows = batch * frames + _col2im_kernel[(rows, rate, triton.cdiv(channels, width))]( + z_gpu, + bias_gpu, + state_gpu, + slots_gpu, + positions_gpu, + parity_gpu, + actual, + frames, + capacity, + R=rate, + C=channels, + BC=width, + num_warps=4, + ) + torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0) + torch.testing.assert_close(state_gpu.cpu(), expected_state, rtol=0, atol=0) diff --git a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py index f929e41c9..8277ba435 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py @@ -594,6 +594,40 @@ def test_streaming_ref_code_context_is_bounded_for_batchable_shapes(): torch.testing.assert_close(frames[:3], ref_code[-3:]) +@pytest.mark.parametrize("tensor_frames", [False, True]) +def test_async_chunks_preserve_codebook_order_with_reference_and_short_tail(tensor_frames): + tm = _tm(left_context=3, initial_chunk_frames=1) + rid = "r-codebook-order" + quantizers = 16 + reference = [[1500 + f * quantizers + q for q in range(quantizers)] for f in range(5)] + tm.request_payload[rid] = torch.tensor(reference, dtype=torch.long) + generated = [[1 + f * quantizers + q for q in range(quantizers)] for f in range(30)] + + for start, end, finished in [(0, 1, False), (1, 26, False), (26, 30, True)]: + rows = generated[:end] + if tensor_frames: + # Non-contiguous int32 rows exercise conversion independently of + # the native producer's contiguous CPU-long representation. + storage = torch.tensor(rows, dtype=torch.int32).repeat_interleave(2, dim=1) + tm.code_prompt_token_ids[rid] = list(storage[:, ::2].unbind()) + else: + tm.code_prompt_token_ids[rid] = rows + payload = talker2code2wav_async_chunk( + transfer_manager=tm, + multimodal_output=None if finished else {}, + request=_req(rid, finished=finished), + is_finished=finished, + ) + assert payload is not None + expected = (reference[-3:] if start == 0 else []) + generated[start:end] + codes = payload.codes.audio + assert codes.dtype == torch.long and codes.device.type == "cpu" and codes.is_contiguous() + torch.testing.assert_close(codes.reshape(quantizers, -1).T, torch.tensor(expected)) + assert payload.meta.left_context_size == (3 if start == 0 else 0) + assert payload.meta.finished.item() == finished + tm.put_req_chunk[rid] += 1 + + def test_ref_code_context_can_be_buffered_before_first_emit(): tm = _tm() rid = "r-ref-buffered" diff --git a/tests/worker_v2/test_index_staging.py b/tests/worker_v2/test_index_staging.py new file mode 100644 index 000000000..81d97701d --- /dev/null +++ b/tests/worker_v2/test_index_staging.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import pytest +import torch + +from vllm_omni.utils.device_copy import index_to_device + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + + +def test_index_staging_survives_later_copies_and_graph_replay(): + stream = torch.cuda.Stream() + graph = torch.cuda.CUDAGraph() + with torch.cuda.stream(stream): + with torch.cuda.graph(graph, stream=stream): + captured = index_to_device([17] * 6, "cuda") + for _ in range(32): + index_to_device([99] * 6, "cuda") + graph.replay() + torch.accelerator.synchronize() + assert captured.tolist() == [17] * 6 diff --git a/tests/worker_v2/test_omni_ar_model_runner.py b/tests/worker_v2/test_omni_ar_model_runner.py index e425cc9c8..0bbebf491 100644 --- a/tests/worker_v2/test_omni_ar_model_runner.py +++ b/tests/worker_v2/test_omni_ar_model_runner.py @@ -25,6 +25,9 @@ class _FakeStream: def wait_stream(self, _stream) -> None: pass + def wait_event(self, _event) -> None: + pass + class _FakeEvent: def record(self, _stream) -> None: @@ -138,6 +141,7 @@ def test_last_pp_rank_orchestration_and_kv_resolver(monkeypatch, needs_history) def test_async_mm_snapshot_owns_output_until_copy_finishes() -> None: runner = OmniARModelRunner.__new__(OmniARModelRunner) runner.model_config = SimpleNamespace(async_chunk=True) + runner.model = SimpleNamespace() runner._async_mm_snapshot_slots, runner._async_mm_snapshot_events = [{}], [None] runner._async_mm_snapshot_pending, runner._async_mm_snapshot_cursor = [False], 0 waited: list[object] = [] @@ -307,6 +311,231 @@ def test_request_reference_codes_preserve_local_axis(prefill_first, padded): assert torch.equal(outputs[i]["codes.audio"], codes[offsets[i] : offsets[i + 1]]) +def test_stream_audio_cancel_rejects_inflight_output_and_id_reuse() -> None: + from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer + + buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3) + cancelled = buffer.add("request") + buffer.finish({"request"}) + replacement = buffer.add("request") + assert cancelled.push(torch.ones(4), False) is None + assert not cancelled.pending + assert replacement.push(torch.full((4,), 2.0), False).tolist() == [2.0] * 4 + buffer.finish({"request"}) + assert not buffer.requests + + +def test_stream_audio_length_end_flushes_first_and_partial_chunks() -> None: + from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput + + buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3) + state = buffer.add("request") + output = StreamingAudioOutput( + torch.ones(1, 4), + torch.tensor([True]), + torch.tensor(24000), + None, + np.array([0, 1]), + [state], + np.array([True]), + ) + assert output.get_output()[0]["model_outputs"].numel() == 4 + assert not state.active + state = buffer.add("another") + assert state.push(torch.ones(4), False).numel() == 4 + assert state.push(torch.ones(4), False) is None + assert state.push(torch.ones(4), True).numel() == 8 + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.float32, torch.float64, torch.bfloat16]) +def test_stream_audio_batch_views_retain_pcm_after_output_release(dtype): + from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput + + buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3) + states = [buffer.add("a"), buffer.add("b")] + expected: list[list[torch.Tensor]] = [[], []] + emitted: list[list[torch.Tensor]] = [[], []] + for step in range(5): + # Mixed prefill/decode spans, invalid rows, and a final partial chunk. + wav = (torch.arange(16).reshape(4, 4) + 100 * step).to(dtype) + valid = torch.tensor([False, True, True, step != 4]) + output = StreamingAudioOutput( + wav, + valid, + torch.tensor(24000), + None, + np.array([0, 3, 4]), + states, + np.array([step == 4, step == 4]), + ) + for i, (start, end) in enumerate([(0, 3), (3, 4)]): + expected[i].append(wav[start:end][valid[start:end]].reshape(-1).clone()) + for i, payload in enumerate(output.get_output()): + if payload is not None: + emitted[i].append(payload["model_outputs"]) + del output, wav # Pending views must keep the old host allocation alive. + for i in range(2): + assert torch.equal(torch.cat(emitted[i]), torch.cat(expected[i])) + assert torch.cat(emitted[i]).dtype == dtype + assert not states[i].active and not states[i].pending + assert states[i].pending_samples == 0 + + +def test_stream_audio_numpy_direct_first_cancel_and_mixed_fallback(): + from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer + + buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3) + state = buffer.add("r") + state.accept_first_audio() + assert state.push(np.arange(8, dtype=np.float32), False) is None + assert state.pending_samples == 4 + # A dtype fallback can arrive while a NumPy frame is pending. + chunk = state.push(torch.arange(8, 12, dtype=torch.bfloat16), True) + assert chunk.tolist() == list(range(4, 12)) + assert state.pending_samples == 0 + buffer.finish({"r"}) + replacement = buffer.add("r") + assert state.push(np.ones(4, dtype=np.float32), False) is None + assert replacement.push(np.full(4, 99, dtype=np.float32), False).tolist() == [99] * 4 + replacement.push(np.ones(4, dtype=np.float32), False) + assert replacement.pending_samples == 4 + buffer.finish({"r"}) + assert replacement.pending_samples == 0 and not replacement.pending + + +def test_stream_audio_prefill_length_end_reaches_output(mocker) -> None: + from vllm.v1.worker.gpu.input_batch import InputBatch + from vllm.v1.worker.gpu.states import RequestState + + from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState + from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer + + owner = mocker.Mock() + owner.vllm_config.model_config.max_model_len = 4096 + owner._stream_decode_event = None + eager = EagerMTPState(owner) + eager._audio_buffer = StreamingAudioBuffer(4, 25) + eager._audio_buffer.add("request") + eager._frame_requests = {"request"} + batch = InputBatch.__new__(InputBatch) + batch.num_reqs, batch.req_ids = 1, ["request"] + batch.num_computed_tokens_np = np.array([0]) + batch.num_scheduled_tokens = np.array([3]) + batch.idx_mapping_np = np.array([0]) + batch.query_start_loc_np = np.array([0, 3]) + batch.is_prefilling_np = np.array([True]) + states = RequestState.__new__(RequestState) + states.max_seq_len = np.array([4]) + mm = { + "model_outputs": torch.ones(3, 4), + "meta": {"codec_frame_valid": torch.tensor([0, 0, 1])}, + "sr": [torch.tensor(24000)], + } + output = eager.prepare_audio_output(batch, states, mm) + assert output.get_output()[0]["model_outputs"].numel() == 4 + assert "model_outputs" not in mm + + +def test_stream_audio_preemption_preserves_position_rng_and_pending_pcm(mocker) -> None: + from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState, TalkerInputs + from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer + + owner = mocker.Mock() + owner._stream_pos = {"request": 37} + owner._eager_embeds = torch.arange(18).reshape(6, 3).float() + owner._eager_ready = {} + owner.intermediate_buffer.buffers = [{} for _ in range(6)] + generator = torch.Generator().manual_seed(42) + owner._mtp_generators = {"request": generator} + saved = [torch.arange(6)] + owner.model.stream_decoder.save_slot.return_value = saved + eager = EagerMTPState(owner) + eager._talker_inputs["request"] = TalkerInputs(torch.ones(20, 3), 10) + eager._side_stream = mocker.Mock() + eager._audio_buffer = StreamingAudioBuffer(4, 25) + request = eager._audio_buffer.add("request") + request.push(torch.ones(4), False) + request.push(torch.ones(4), False) + eager.suspend_audio("request", 2) + owner._stream_pos.clear() + owner._mtp_generators.clear() + eager.resume_audio("request", 5) + assert owner._stream_pos["request"] == 37 + assert owner._mtp_generators["request"] is generator + assert eager._restore_audio["request"] is saved + assert sum(part.numel() for part in request.pending) == 4 + # A second preemption can happen before the restored slot is used. + eager.suspend_audio("request", 5) + owner.model.stream_decoder.save_slot.assert_called_once_with(2) + eager.finish_audio({"request"}) + assert not eager._suspended_audio and not eager._restore_audio + assert not eager._audio_buffer.requests and not request.active + + +def test_stream_audio_replay_uses_exact_inputs_and_consumes_last_frame_once(mocker): + from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState, TalkerInputs + + owner = mocker.Mock() + owner._eager_embeds = torch.tensor([[20.0, 30.0]]) + owner.model.preprocess.return_value = (None, None, {"mtp_inputs": (None, torch.tensor([[2.0, 3.0]]))}) + owner.intermediate_buffer.buffers = [{"req_id": "request"}] + eager = EagerMTPState(owner) + eager._restore_audio["request"] = [] + history = torch.arange(8).reshape(4, 2).float() + eager._talker_inputs["request"] = TalkerInputs(history, 4) + first = torch.empty(2, 2) + assert eager.replay_inputs("request", 0, 0, torch.ones(2), first) + torch.testing.assert_close(first, history[:2]) + owner.model.preprocess.assert_not_called() + last = torch.empty(3, 2) + assert eager.replay_inputs("request", 0, 2, torch.ones(3), last) + torch.testing.assert_close(last, torch.cat((history[2:], torch.tensor([[22.0, 33.0]])))) + owner.model.preprocess.assert_called_once() + with pytest.raises(RuntimeError, match="exceeds saved inputs"): + eager.replay_inputs("request", 0, 4, torch.ones(2), torch.empty(2, 2)) + + +def test_stream_reference_priming_groups_lengths_without_padding_or_slot_aliasing(mocker): + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration + + primes = [(4, torch.full((26, 2), 4)), (1, torch.full((3, 2), 1)), (7, torch.full((26, 2), 7))] + calls = [] + + def stream(codes, slots, pos): + calls.append((codes.clone(), slots.clone(), pos.clone())) + + decoder = mocker.Mock(side_effect=stream, device=torch.device("cpu")) + model = mocker.Mock(stream_decoder=decoder, stream_prime_graphs=None, stream_chunk_frames=25) + Qwen3TTSTalkerForConditionalGeneration.prime_stream_decoder(model, primes) + assert len(calls) == 3 + for (codes, slots, pos), frames, indices, position in zip( + calls, [25, 1, 3], [[4, 7], [4, 7], [1]], [0, 25, 0], strict=True + ): + assert codes.shape == (len(indices), frames, 2) + assert codes.dtype == slots.dtype == pos.dtype == torch.int32 + assert slots.tolist() == indices and pos.tolist() == [position] * len(indices) + for row, idx in enumerate(indices): + assert torch.all(codes[row] == idx) + Qwen3TTSTalkerForConditionalGeneration.prime_stream_decoder(model, []) + assert len(calls) == 3 + + +@pytest.mark.parametrize("context_frames", [25, 72]) +def test_model_owned_reference_context_respects_configured_boundary(mocker, context_frames): + from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration + + model = mocker.Mock(stream_ref_context_frames=context_frames, talker_config=mocker.Mock(num_code_groups=16)) + reference = torch.arange(100 * 16).reshape(100, 16) + info = {"codes": {"ref": reference}} + got = Qwen3TTSTalkerForConditionalGeneration.get_stream_ref_context(model, info) + torch.testing.assert_close(got, reference[-context_frames:]) + assert info["codes"]["ref"] is reference + assert Qwen3TTSTalkerForConditionalGeneration.get_stream_ref_context(model, {}) is None + assert ( + Qwen3TTSTalkerForConditionalGeneration.get_stream_ref_context(model, {"codes": {"ref": torch.empty(0)}}) is None + ) + + @pytest.mark.parametrize("streaming", [False, True]) def test_model_owned_audio_finalizer_runs_after_copy_without_hidden(monkeypatch, streaming): monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None) @@ -340,6 +569,136 @@ def test_model_owned_audio_finalizer_runs_after_copy_without_hidden(monkeypatch, assert (result.pooler_output is None) == streaming +def test_stream_audio_history_replays_exact_inputs_after_batch_reordering(mocker): + from vllm.v1.worker.gpu.input_batch import InputBatch + + from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState + + owner = mocker.Mock() + owner.vllm_config.model_config.max_model_len = 12 + eager = EagerMTPState(owner) + first = torch.arange(9, dtype=torch.float32).reshape(3, 3) + second = torch.arange(9, 18, dtype=torch.float32).reshape(3, 3) + expected = {"a": torch.cat((first[:2], second[2:])), "b": torch.cat((first[2:], second[:2]))} + batch = InputBatch.__new__(InputBatch) + batch.req_ids = ["a", "b"] + batch.query_start_loc_np = np.array([0, 2, 3]) + batch.num_computed_tokens_np = np.array([0, 0]) + eager.record_inputs(batch, first) + batch.req_ids = ["b", "a"] + batch.num_computed_tokens_np = np.array([1, 2]) + eager.record_inputs(batch, second) + # Later input-buffer reuse must not change the conditioned replay history. + first.fill_(-1) + second.fill_(-2) + for request_id, wanted in expected.items(): + eager._restore_audio[request_id] = [] + replayed = torch.empty_like(wanted) + assert eager.replay_inputs(request_id, 0, 0, torch.zeros(3, dtype=torch.long), replayed) + torch.testing.assert_close(replayed, wanted) + + +@pytest.mark.parametrize("valid", [False, True]) +def test_stream_audio_direct_first_requires_only_real_pcm_on_terminal(valid): + from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY + from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput + + buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3) + state = buffer.add("r") + state.accept_first_audio() + first = StreamingAudioOutput( + torch.arange(4).reshape(1, 4), + torch.tensor([valid]), + torch.tensor(24000), + None, + np.array([0, 1]), + [state], + np.array([True]), + ) + payload = first.get_output()[0] + if valid: + assert payload == {FIRST_AUDIO_REQUIRED_KEY: torch.tensor(True)} + else: + assert payload is None # EOS has no promised first PCM to wait for. + assert not state.active and not state.pending + buffer.finish({"r"}) + replacement = buffer.add("r") + assert replacement.push(torch.full((4,), 99), False).tolist() == [99] * 4 + + +def test_stream_audio_direct_first_drops_one_frame_and_retains_partial_suffix(): + from vllm_omni.worker_v2.streaming_audio import AudioRequest + + state = AudioRequest(samples_per_frame=4, chunk_frames=3) + state.accept_first_audio() + assert state.push(torch.arange(4), False) is None + assert state.first_audio_required and state.emitted + assert state.push(torch.arange(4, 8), False) is None + chunk = state.push(torch.arange(8, 12), True) + assert chunk.tolist() == list(range(4, 12)) + assert not state.active and not state.pending + + +@pytest.mark.parametrize("consumer", ["pcm_only", "finalizer", "extra"]) +def test_stream_audio_snapshot_skips_discarded_code_partition(monkeypatch, mocker, consumer): + from vllm.v1.worker.gpu.input_batch import InputBatch + + from vllm_omni.worker_v2.output_snapshot import RequestOutputSnapshot + from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput + + monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None) + monkeypatch.setattr( + OmniARModelRunner, + "_build_async_chunk_outputs_from_mm", + lambda *args: pytest.fail("in-stage PCM unnecessarily partitioned code outputs"), + ) + batch = InputBatch.__new__(InputBatch) + batch.num_reqs = 3 + batch.query_start_loc_np = np.array([0, 1, 2, 4]) + batch.num_scheduled_tokens = np.array([1, 1, 2]) + batch.num_tokens_after_padding = 4 + buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3) + states = [buffer.add(request_id) for request_id in ["frame", "eos", "prefill"]] + wav = torch.arange(16, dtype=torch.float32).reshape(4, 4) + streaming_audio = StreamingAudioOutput( + wav, + torch.tensor([True, False, True, False]), + torch.tensor(24000), + None, + batch.query_start_loc_np.copy(), + states, + np.array([False, False, True]), + ) + copy_mm = mocker.spy(omni_ar_model_runner, "_async_copy_mm") + + def finalize(payload, counts): + assert payload["codes"]["audio"].shape == (4, 2) + return RequestOutputSnapshot([None] * batch.num_reqs) + + finalizer = mocker.Mock(side_effect=finalize) if consumer == "finalizer" else None + extra = ({"extra": torch.tensor([7])}, _FakeEvent()) if consumer == "extra" else None + result = _async_output( + req_ids=["frame", "eos", "prefill"], + sampler_output=SamplerOutput(torch.tensor([[1], [2], [3]]), None, None, None), + num_sampled_tokens=torch.tensor([1, 1, 1]), + multimodal_outputs={"codes": {"audio": torch.ones(4, 2)}}, + input_batch=batch, + async_chunk=True, + streaming_audio=streaming_audio, + finalize_multimodal=finalizer, + extra_multimodal_outputs=extra, + ).get_output() + assert copy_mm.call_count == {"pcm_only": 0, "finalizer": 1, "extra": 2}[consumer] + if finalizer is not None: + finalizer.assert_called_once() + wav.fill_(-1) + assert result.inter_stage_outputs is None + assert result.multimodal_outputs[0]["model_outputs"].tolist() == [0, 1, 2, 3] + assert result.multimodal_outputs[1] == {} + assert result.multimodal_outputs[2]["model_outputs"].tolist() == [8, 9, 10, 11] + assert not states[1].active and not states[2].active + + @pytest.mark.parametrize("streaming", [False, True]) def test_request_owned_snapshot_skips_generic_partition(monkeypatch, streaming): from vllm_omni.worker_v2.output_snapshot import RequestOutputSnapshot diff --git a/tests/worker_v2/test_omni_model_state.py b/tests/worker_v2/test_omni_model_state.py index 4edfed603..b36dbb139 100644 --- a/tests/worker_v2/test_omni_model_state.py +++ b/tests/worker_v2/test_omni_model_state.py @@ -36,6 +36,7 @@ class _DummyInputBatch: def _make_state(max_num_reqs=4, has_preprocess=False, has_postprocess=False, have_multimodal_outputs=False): state = object.__new__(OmniModelState) model = MagicMock() + model.stream_decoder = None model.has_preprocess = has_preprocess model.has_postprocess = has_postprocess model.have_multimodal_outputs = have_multimodal_outputs @@ -62,6 +63,10 @@ def _make_state(max_num_reqs=4, has_preprocess=False, has_postprocess=False, hav state.intermediate_buffer = OmniIntermediateBuffer(max_num_reqs) state._static_inputs_embeds = None + from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState + + state._eager_state = EagerMTPState(state) + state._stream_pos = {} state._mtp_generators = {} state._mtp_runner = None for name in ("_mtp_input_ids", "_mtp_input_embeds", "_mtp_hidden", "_mtp_text_step", "_mtp_offsets"): @@ -264,6 +269,7 @@ def test_seed_independence_resolve_once_and_sampling_kwargs(): # vLLM sampling seed must not produce a talker generator. cpu = torch.device("cpu") assert state._get_mtp_generator("r1", SimpleNamespace(extra_args={}, seed=42), cpu) is None + state._mtp_generators.clear() # The following rows are new requests. # Same model-local seed reproduces identical uniforms regardless of batch makeup. state._mtp_sample_uniforms = torch.empty((2, 2, 4)) assert torch.equal( diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index a56e6148c..13115f00e 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -5,6 +5,7 @@ from dataclasses import MISSING, field from typing import Any from pydantic import ConfigDict, TypeAdapter +from transformers import PretrainedConfig from vllm.config import ModelConfig from vllm.config.utils import config from vllm.logger import init_logger @@ -124,9 +125,12 @@ class OmniModelConfig(ModelConfig): """ stage_id: int = 0 + # Match the upstream derived field while its value is replaced per stage. + hf_text_config: PretrainedConfig = field(init=False) async_chunk: bool = False session_mode: str = "turn" retains_state_across_chunks: bool = False + supports_running_prefix_cache_reset: bool = True use_v2_model_runner: bool = False supports_native_mrv2_data_plane: bool = False # Stage-1 active stream slots; 0 keeps legacy chunk-level round-robin. @@ -218,7 +222,7 @@ class OmniModelConfig(ModelConfig): return convertor.convert() return super().get_model_arch_config() - def draw_hf_text_config(self): + def draw_hf_text_config(self) -> PretrainedConfig: # transformers' get_text_config method is used to get the text config from thinker_config. # to handle the case that each model stage has their own text config, # we need to draw the text config from the corresponding model stage. diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index d77e44194..19d5ae71c 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -482,6 +482,7 @@ class OmniStageModelConfig(_TrackExplicitConfigFields): interleave_mm_strings: bool | None = None media_io_kwargs: dict[str, Any] | None = None final_output: bool = False + supports_running_prefix_cache_reset: bool = True active_stream_window: int = Field(default=0, ge=0) session_mode: str = "turn" duplex_max_sessions: int = Field(default=1, ge=1) @@ -1325,6 +1326,7 @@ _DIFFUSION_STAGE_METADATA_FIELDS = frozenset( "model_arch", "model_stage", "retains_state_across_chunks", + "supports_running_prefix_cache_reset", "scheduler_cls", "stage_connector_spec", "worker_type", @@ -1972,6 +1974,8 @@ def _build_model_config( if "active_stream_window" not in kwargs: kwargs["active_stream_window"] = _copy_value(deploy.active_stream_window) kwargs["final_output"] = topology.final_output + if not topology.supports_running_prefix_cache_reset: + kwargs["supports_running_prefix_cache_reset"] = False if "custom_voice_dir" not in kwargs and deploy.custom_voice_dir is not None: kwargs["custom_voice_dir"] = _copy_value(deploy.custom_voice_dir) stage_runner = resolve_stage_model_runner(deploy, stage_deploy) diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index f016b7f7d..cbcb94cdb 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -117,7 +117,7 @@ from vllm_omni.model_executor.models.qwen3_omni.pipeline import ( QWEN3_OMNI_THINKER_ONLY_PIPELINE, resolve_qwen3_omni_pipeline, ) -from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE +from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_FUSED_PIPELINE, QWEN3_TTS_PIPELINE from vllm_omni.model_executor.models.step_audio2.pipeline import ( STEP_AUDIO2_ASR_PIPELINE, STEP_AUDIO2_PIPELINE, @@ -204,6 +204,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "qwen3_omni_moe": resolve_qwen3_omni_pipeline, "qwen3_omni_moe_thinker_only": QWEN3_OMNI_THINKER_ONLY_PIPELINE, "qwen3_tts": QWEN3_TTS_PIPELINE, + "qwen3_tts_fused": QWEN3_TTS_FUSED_PIPELINE, "step_audio_2": STEP_AUDIO2_PIPELINE, "step_audio_2_asr": STEP_AUDIO2_ASR_PIPELINE, "voxcpm2": VOXCPM2_PIPELINE, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 5f959050a..eabea35c2 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -245,6 +245,8 @@ class StagePipelineConfig: # The model keeps per-request execution state while awaiting the next # async chunk, so the parked request continues to consume model capacity. retains_state_across_chunks: bool = False + # Some stateful audio stages cannot roll back already consumed frames. + supports_running_prefix_cache_reset: bool = True sampling_constraints: dict[str, Any] = field(default_factory=dict) custom_process_input_func: str | None = None custom_process_next_stage_input_func: str | None = None @@ -295,6 +297,9 @@ class PipelineConfig: model_type: str model_arch: str = "" stages: tuple[StagePipelineConfig, ...] = () + # A single stage that streams its own final output in async-chunk mode + # (e.g. a Talker decoding audio in-stage) keeps deploy.async_chunk. + single_stage_async_chunk: bool = False # HF architecture aliases: used by StageConfigFactory when the model's # HF config reports a generic model_type that collides with a different # model (e.g. MiMo Audio reports model_type="qwen2"). The factory @@ -1107,6 +1112,8 @@ def _build_engine_args( if ps.omni_kv_config: engine_args["omni_kv_config"] = dict(ps.omni_kv_config) engine_args["requires_full_payload_input"] = ps.requires_full_payload_input + if not ps.supports_running_prefix_cache_reset: + engine_args["supports_running_prefix_cache_reset"] = False return engine_args @@ -1172,8 +1179,11 @@ def _resolve_pipeline_async_chunk_enabled( and False otherwise. If the user tried to enable async chunk through the deploy config, but it's inapplicable or unsupported, it will be disabled with a warning. """ - # Single stage should never use async chunk if len(pipeline.stages) <= 1: + if pipeline.single_stage_async_chunk: + if deploy.async_chunk is False: + raise ValueError(f"Pipeline {pipeline.model_type!r} requires async_chunk=True") + return True if deploy.async_chunk: logger.warning( "Deploy config set async_chunk=True, but async chunk is inapplicable " @@ -1350,6 +1360,8 @@ class StageConfig: # Terminal-stage ownership comes from topology, not engine overrides. engine_args["final_output"] = self.final_output + if self.yaml_engine_args.get("supports_running_prefix_cache_reset") is False: + engine_args["supports_running_prefix_cache_reset"] = False # Build runtime config from YAML defaults + CLI overrides runtime: dict[str, Any] = dict(self.yaml_runtime) diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 0138d365e..7dc0a5d7c 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -101,6 +101,23 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): max_num_running_reqs: int + def reset_prefix_cache(self, reset_running_requests: bool = False, reset_connector: bool = False) -> bool: + model_config = self.vllm_config.model_config + if ( + reset_running_requests + and self.running + and not getattr(model_config, "supports_running_prefix_cache_reset", True) + ): + # Reset preempts and resumes in the same step while discarding + # in-flight tokens. The stateful codec and queued PCM have already + # consumed those frames and cannot roll back to that boundary. + logger.warning( + "This stage cannot reset running requests; wait for " + "completion or abort them before resetting the prefix cache." + ) + return False + return super().reset_prefix_cache(reset_running_requests, reset_connector) + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) # Track requests that need KV cache transfer when finished @@ -517,7 +534,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # domains must consume the old frame before new output passes. request.async_tokens_to_discard = max(0, stale_async_tokens - len(generated_token_ids)) - if output_is_stale or async_output_is_stale: + if async_output_is_stale or (output_is_stale and request.drop_stale_output): # Output of a step scheduled before the request's in-flight # tokens were discarded (segment stop / session replacement). # num_computed_tokens was rolled back at the discard site, so @@ -556,12 +573,13 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # tokens and rejections. If some tokens are rejected, # num_computed_tokens is decreased by the number of rejected # tokens. - if request.num_computed_tokens > 0: - request.num_computed_tokens -= num_rejected - # If async scheduling, num_output_placeholders also includes - # the scheduled spec tokens count and so is similarly adjusted. - if request.num_output_placeholders > 0: - request.num_output_placeholders -= num_rejected + if not output_is_stale: + if request.num_computed_tokens > 0: + request.num_computed_tokens -= num_rejected + # If async scheduling, num_output_placeholders also includes + # the scheduled spec tokens count and so is similarly adjusted. + if request.num_output_placeholders > 0: + request.num_output_placeholders -= num_rejected spec_decoding_stats = self.make_spec_decoding_stats( spec_decoding_stats, num_draft_tokens=num_draft_tokens, @@ -600,7 +618,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # Check for stop and update request status. if new_token_ids: num_sampled_tokens = len(new_token_ids) - new_token_ids, stopped = self._update_request_with_output(request, new_token_ids) + if output_is_stale: + new_token_ids, stopped = self._update_request_with_output(request, new_token_ids, is_stale=True) + else: + new_token_ids, stopped = self._update_request_with_output(request, new_token_ids) if new_logprobs is not None and len(new_token_ids) < num_sampled_tokens: # A mid-step stop (e.g. spec-decode tokens sampled past # EOS) trims new_token_ids after the validation slice @@ -704,6 +725,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # a queued streaming update, and adding twice swallows the # next duplex unit's listen/speak under async scheduling. if request.num_in_flight_tokens > 0: + request.drop_stale_output = True request.num_stale_output_tokens = request.num_in_flight_tokens if outstanding_async_tokens > 0: # Discard only outputs that are already in flight and @@ -884,6 +906,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # a resumable stop applies a queued update through this helper. in_flight_tokens = int(getattr(session, "num_in_flight_tokens", 0) or 0) if in_flight_tokens > 0: + session.drop_stale_output = True session.num_stale_output_tokens = in_flight_tokens if outstanding_async_tokens > 0: # Async scheduling may already have sampled the previous @@ -1405,6 +1428,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if status is not None: request.status = status if native_transfer and connector_delay_free_blocks: + assert transfer_params is not None kv_xfer_params = { **(kv_xfer_params or {}), "transfer_id": transfer_params["transfer_id"], diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 1a1a46276..77a568d14 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -222,6 +222,7 @@ class OmniSchedulerMixin(_SchedulerMixinBase): # num_in_flight_tokens already includes any undrained stale share. # Assign instead of accumulating so callers that fenced the same # rollover before entering this helper do not count it twice. + session.drop_stale_output = True session.num_stale_output_tokens = int(getattr(session, "num_in_flight_tokens", 0) or 0) session.num_output_placeholders = 0 session.spec_token_ids = [] @@ -316,6 +317,7 @@ class OmniSchedulerMixin(_SchedulerMixinBase): # The streaming update may already have fenced this same in-flight # frame. Seed idempotently so the replacement does not count it twice. request.num_stale_output_tokens = int(getattr(request, "num_in_flight_tokens", 0) or 0) + request.drop_stale_output = True request.num_output_placeholders = 0 request.spec_token_ids = [] self._release_replaced_streaming_prompt_cache(request) diff --git a/vllm_omni/deploy/qwen3_tts_fused_single_gpu.yaml b/vllm_omni/deploy/qwen3_tts_fused_single_gpu.yaml new file mode 100644 index 000000000..c2492f68c --- /dev/null +++ b/vllm_omni/deploy/qwen3_tts_fused_single_gpu.yaml @@ -0,0 +1,54 @@ +# Single-stage Qwen3-TTS: the Talker decodes every frame with the stateful +# streaming codec decoder and returns PCM itself (no Code2Wav stage). +# reset_prefix_cache(reset_running_requests=True) returns False while requests +# are running; wait for completion or abort them before resetting the cache. +pipeline: qwen3_tts_fused +model_runner: v2 +async_chunk: true + +stages: + - stage_id: 0 + devices: "0" + max_num_seqs: 128 + # Exact conditioned Talker inputs for preemption replay are allocated + # after KV profiling: 16 MiB per active request at 1.7B/max_model_len=4096 + # (~2 GiB at 128 requests), in addition to codec state and output buffers. + # Keep headroom when increasing gpu_memory_utilization or capacity. + gpu_memory_utilization: 0.5 + trust_remote_code: true + enable_prefix_caching: false + # Model options that multi-stage deployments put on the connector: a + # single stage has no connector edge. + additional_config: + talker_stream_decode: true + # Opt in to first-frame delivery on its own queue for lower first-audio + # latency; leave disabled when prioritizing peak throughput. + talker_stream_first_audio: false + # Match the standard two-stage profile's voice-clone context boundary. + # A shorter context (e.g. 25) can change timbre continuity. + ref_code_context_frames: 72 + code_predictor_kv_cache: true + code_predictor_fused_sampling: true + code_predictor_fused: true + # Talker decode contexts are short (~100-700 tokens): FA3 split-KV plus + # its combine kernel costs more than one pass per head. + attention_config: + flash_attn_max_num_splits_for_cuda_graph: 1 + async_scheduling: true + max_num_batched_tokens: 512 + max_model_len: 4096 + compilation_config: + cudagraph_mode: FULL + cudagraph_capture_sizes: [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 144, 160, 176, 192, 208, 224, 240, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512] + max_cudagraph_capture_size: 512 + default_sampling_params: + temperature: 0.9 + top_k: 50 + max_tokens: 4096 + min_tokens: 2 + repetition_penalty: 1.05 + subtalker_sampling_params: + do_sample: true + temperature: 0.9 + top_k: 50 + top_p: 1.0 diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 458a6237d..bad555704 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -198,6 +198,7 @@ class OmniEngineArgs(EngineArgs): async_chunk: bool = False session_mode: str = "turn" retains_state_across_chunks: bool = False + supports_running_prefix_cache_reset: bool = True use_v2_model_runner: bool = False supports_native_mrv2_data_plane: bool = False # WS-A: Stage-1 active stream slots. 0 = legacy preempt-everything. @@ -430,6 +431,7 @@ class OmniEngineArgs(EngineArgs): async_chunk=self.async_chunk, session_mode=self.session_mode, retains_state_across_chunks=self.retains_state_across_chunks, + supports_running_prefix_cache_reset=self.supports_running_prefix_cache_reset, use_v2_model_runner=self.use_v2_model_runner, supports_native_mrv2_data_plane=self.supports_native_mrv2_data_plane, active_stream_window=self.active_stream_window, diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 4047a7147..8f55280cd 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -898,7 +898,7 @@ class OrchestratorBase: req_state is None or req_state.upstream_first_audio or req_state.pending_upstream_first_audio is not None - or stage_id + 1 > req_state.final_stage_id + or (stage_id if final_output else stage_id + 1) > req_state.final_stage_id ): continue audio = mm.get("model_outputs") @@ -935,7 +935,7 @@ class OrchestratorBase: if pending is None or self.request_states.get(req_state.request_id) is not req_state: return source_stage, first_output = pending - codec_stage = source_stage + 1 + codec_stage = source_stage if self.stage_pools[source_stage].final_output else source_stage + 1 pool = self.stage_pools[codec_stage] replica_id = pool.get_bound_replica_id(req_state.request_id) if replica_id is None or req_state.request_id not in pool.output_processor.request_states: diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index ecc9576bc..c19a5cf75 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -87,7 +87,10 @@ def _bind_first_audio_sink(model_executor: Any, output_queue: Any, scheduler: An model_runner = getattr(worker, "model_runner", None) model_state = getattr(model_runner, "model_state", None) model = getattr(model_runner, "model", None) - if getattr(model, "first_frame_decoder", None) is None or not hasattr(model_state, "set_first_audio_sink"): + decodes_audio = getattr(model, "first_frame_decoder", None) is not None or ( + getattr(model, "stream_decoder", None) is not None and bool(getattr(model, "stream_first_audio", False)) + ) + if not decodes_audio or not hasattr(model_state, "set_first_audio_sink"): return False from vllm_omni.worker_v2.first_audio_sender import engine_output_queue_sink diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index f9f3849f4..92e04e946 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -249,6 +249,33 @@ async def omni_run_server(args, **uvicorn_kwargs) -> None: await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs) +def _own_reuseport_socket(sock: socket.socket, client_count: int) -> socket.socket: + """Give this API server its own SO_REUSEPORT listener on the shared address. + + All API servers accepting from the one inherited socket lets whichever + wakes first take a whole burst of connections (the event loop accepts + in a loop), and keep-alive then pins that imbalance for the connections' + lifetime. Separate reuseport listeners let the kernel spread connections + by hash. The inherited socket is bound but never listened on, so it + receives nothing. + """ + if client_count <= 1: + return sock + if sock.family not in (socket.AF_INET, socket.AF_INET6): + return sock + try: + if not sock.getsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT): + return sock + own = socket.socket(family=sock.family, type=socket.SOCK_STREAM) + own.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + own.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) + own.bind(sock.getsockname()[:2]) + except OSError as e: + logger.warning("Per-server reuseport listener unavailable (%s); sharing the inherited socket", e) + return sock + return own + + def run_omni_api_server_worker_proc( listen_address: str, sock: socket.socket, @@ -278,6 +305,7 @@ def run_omni_api_server_worker_proc( set_process_title("APIServer", str(client_index)) decorate_logs("APIServer", skip_if_decorated=True) + sock = _own_reuseport_socket(sock, int(omni_client_config["client_count"])) uvloop.run( omni_run_server_worker( listen_address, @@ -512,7 +540,7 @@ async def build_async_omni( # Ensures everything is shutdown and cleaned up on error/exit async with build_async_omni_from_stage_config( args, - disable_frontend_multiprocessing=disable_frontend_multiprocessing, + disable_frontend_multiprocessing=bool(disable_frontend_multiprocessing), client_config=client_config, ) as async_omni: yield async_omni diff --git a/vllm_omni/entrypoints/openai/audio_utils_mixin.py b/vllm_omni/entrypoints/openai/audio_utils_mixin.py index b82842c42..f8af4d9cd 100644 --- a/vllm_omni/entrypoints/openai/audio_utils_mixin.py +++ b/vllm_omni/entrypoints/openai/audio_utils_mixin.py @@ -46,6 +46,15 @@ except ImportError: logger = init_logger(__name__) +def _float32_to_pcm16_bytes(audio: np.ndarray) -> bytes: + """PCM_16 bytes of float32 samples, bit-identical to libsndfile's RAW writer. + + libsndfile (1.2) converts with ``floor(x * 0x8000)`` clipped to int16. + """ + scaled = np.floor(audio.astype(np.float64) * 32768.0) + return np.clip(scaled, -32768, 32767).astype(" None: + server = self.ctx.server + if server._tts_tokenizer is not None: + return + import transformers + from transformers import AutoTokenizer + + # Same construction as the engine's text tokenizer + # (Qwen3TTSPromptEmbedsBuilder.get_text_tokenizer), so the ids computed + # here can be handed to the engine. + kwargs: dict[str, Any] = dict(trust_remote_code=True, use_fast=True, padding_side="left") + if transformers.__version__ < "5": + kwargs["fix_mistral_regex"] = True + assert self.ctx.engine_client is not None + model_name = self.ctx.engine_client.model_config.model + server._tts_tokenizer = AutoTokenizer.from_pretrained(model_name, **kwargs) + def _estimate_prompt_len(self, tts_params: dict[str, Any]) -> int: """Estimate prompt length so the placeholder matches model-side embeddings.""" try: from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import ( + PRECOMPUTED_TEXT_IDS_KEY, Qwen3TTSPromptEmbedsBuilder, + build_assistant_text, ) server = self.ctx.server - if server._tts_tokenizer is None: - from transformers import AutoTokenizer - - model_name = self.ctx.engine_client.model_config.model - server._tts_tokenizer = AutoTokenizer.from_pretrained( - model_name, - trust_remote_code=True, - padding_side="left", - ) + self._load_text_tokenizer() + assert self.ctx.engine_client is not None hf_config = self.ctx.engine_client.model_config.hf_config talker_config = hf_config.talker_config task_type = (tts_params.get("task_type") or ["CustomVoice"])[0] - return Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information( + text = (tts_params.get("text") or [""])[0] + assistant_text = build_assistant_text(text) if isinstance(text, str) else None + assistant_ids: list[list[int]] = [] + + def tokenize_prompt(t: str) -> list[int]: + ids = server._tts_tokenizer(t, padding=False)["input_ids"] + if t == assistant_text: + assistant_ids.append(ids) + return ids + + prompt_len = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information( additional_information=tts_params, task_type=task_type, - tokenize_prompt=lambda t: server._tts_tokenizer(t, padding=False)["input_ids"], + tokenize_prompt=tokenize_prompt, codec_language_id=getattr(talker_config, "codec_language_id", None), spk_is_dialect=getattr(talker_config, "spk_is_dialect", None), estimate_ref_code_len=self._estimate_ref_code_len, ) + if assistant_ids: + # Spares the engine tokenizing the same text on its step thread. + tts_params[PRECOMPUTED_TEXT_IDS_KEY] = [assistant_ids[0]] + return prompt_len except Exception as e: logger.warning("Failed to estimate TTS prompt length, using fallback 2048: %s", e) return 2048 diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index 37a9c2388..88d3427bb 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -954,7 +954,13 @@ class CodePredictorWrapper(nn.Module): extra_cfg = connector_cfg.get("extra", connector_cfg) else: extra_cfg = getattr(connector_cfg, "extra", None) - return extra_cfg if isinstance(extra_cfg, dict) else {} + extra_cfg = extra_cfg if isinstance(extra_cfg, dict) else {} + # A single-stage deployment has no connector edge to carry model + # options; it sets them in the stage's ``additional_config``. + additional = getattr(vllm_config, "additional_config", None) + if isinstance(additional, dict) and additional: + return {**additional, **extra_cfg} + return extra_cfg @staticmethod def _parse_bool_config(value: object) -> bool: @@ -1332,6 +1338,21 @@ class CodePredictorWrapper(nn.Module): scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf")) return self._sample_codes_gumbel(scaled, generator=generator, uniforms=uniforms) + def _validate_sampling_inputs( + self, + bsz: int, + generators: Sequence[torch.Generator | None] | None, + sample_uniforms: torch.Tensor | None, + ) -> None: + if generators is not None and len(generators) != bsz: + raise ValueError(f"generators must have one entry per row: got {len(generators)} for batch {bsz}") + if sample_uniforms is not None: + expected_shape = (bsz, self._num_groups - 1, int(self.config.vocab_size)) + if tuple(sample_uniforms.shape) != expected_shape: + raise ValueError( + f"sample_uniforms must have shape {expected_shape}, got {tuple(sample_uniforms.shape)}" + ) + @torch.inference_mode() def forward( self, @@ -1349,14 +1370,7 @@ class CodePredictorWrapper(nn.Module): """Predict residual codebooks 1..G-1 autoregressively via re-prefill.""" bsz = int(layer0_code.shape[0]) num_groups = self._num_groups - if generators is not None and len(generators) != bsz: - raise ValueError(f"generators must have one entry per row: got {len(generators)} for batch {bsz}") - if sample_uniforms is not None: - expected_shape = (bsz, num_groups - 1, int(self.config.vocab_size)) - if tuple(sample_uniforms.shape) != expected_shape: - raise ValueError( - f"sample_uniforms must have shape {expected_shape}, got {tuple(sample_uniforms.shape)}" - ) + self._validate_sampling_inputs(bsz, generators, sample_uniforms) sample_generator: _GeneratorLike = generators if generators is not None else generator device = layer0_code.device diff --git a/vllm_omni/model_executor/models/qwen3_tts/fused_code_predictor.py b/vllm_omni/model_executor/models/qwen3_tts/fused_code_predictor.py new file mode 100644 index 000000000..5133dd85a --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/fused_code_predictor.py @@ -0,0 +1,312 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# ruff: noqa: N803 - Triton constexpr parameters use kernel-style capitals. +"""Fused residual-codebook predictor for Qwen3-TTS (CUDA, BF16). + +Same computation as the frame-local-KV predictor with per-call top-k Gumbel +sampling, restructured to issue few kernels per residual step, since at serving +batch sizes each of the 15 dependent steps is launch- and latency-bound: + +- Per layer: RMSNorm, one QKV GEMM, one attention kernel (q/k RMSNorm, RoPE, + this call's K/V write and attention over the frame's keys), O projection + accumulated into the residual stream (``addmm_``), RMSNorm, gate/up GEMM, + SiLU-mul, down projection accumulated into the stream. +- The embedding of each sampled code followed by ``small_to_mtp_projection`` + is linear, so it folds at load into one ``[vocab, hidden]`` table per + codebook; the sampling kernel writes the next step's input row directly. +- Sampling finds the top-k threshold by bisection over the BF16 order keys + (exact k-th largest, ties kept) instead of a sort. +""" + +from __future__ import annotations + +import torch +from vllm.triton_utils import tl, tldevice, triton + +from .tokenizer_12hz.streaming_decoder import _silu_mul_kernel + + +@triton.jit +def _rms_rows_kernel(x_ptr, x_stride, w_ptr, out_ptr, eps, N: tl.constexpr): + row = tl.program_id(0).to(tl.int64) + offs = tl.arange(0, N) + x = tl.load(x_ptr + row * x_stride + offs).to(tl.float32) + var = tl.sum(x * x, axis=0) / N + dt = out_ptr.dtype.element_ty + y = (x * tl.rsqrt(var + eps)).to(dt).to(tl.float32) + w = tl.load(w_ptr + offs).to(tl.float32) + tl.store(out_ptr + row * N + offs, (w * y).to(dt)) + + +@triton.jit +def _bf(x): + return x.to(tl.bfloat16).to(tl.float32) + + +@triton.jit +def _head_norm_rope(x1, x2, w1, w2, c1, c2, s1, s2, eps, HD: tl.constexpr): + """HF RMSNorm (``weight * normed.to(bf16)``) then RoPE in BF16 elementwise steps, on the two halves.""" + var = (tl.sum(x1 * x1, axis=0) + tl.sum(x2 * x2, axis=0)) / HD + r = tl.rsqrt(var + eps) + n1 = _bf(w1 * _bf(x1 * r)) + n2 = _bf(w2 * _bf(x2 * r)) + # rotate_half(n) = [-n2, n1] + o1 = _bf(_bf(n1 * c1) + _bf(-n2 * s1)) + o2 = _bf(_bf(n2 * c2) + _bf(n1 * s2)) + return o1, o2 + + +@triton.jit +def _cp_attention_kernel( + qkv_ptr, # [B * NQ, (H + 2 * KVH) * HD] + kc_ptr, # [Bmax, KVH, MAXP, HD] + vc_ptr, + qw_ptr, + kw_ptr, + cos_ptr, # [MAXP, HD] bf16 + sin_ptr, + out_ptr, # [B * NQ, H * HD] + first, + eps, + scale, + NQ: tl.constexpr, + H: tl.constexpr, + KVH: tl.constexpr, + HD: tl.constexpr, + MAXP: tl.constexpr, + BK: tl.constexpr, +): + b = tl.program_id(0) + h = tl.program_id(1) + qi = tl.program_id(2) + kvh = h // (H // KVH) + W: tl.constexpr = (H + 2 * KVH) * HD + HALF: tl.constexpr = HD // 2 + d = tl.arange(0, HALF) + p = first + qi + row = (b * NQ + qi).to(tl.int64) + qw1 = tl.load(qw_ptr + d).to(tl.float32) + qw2 = tl.load(qw_ptr + HALF + d).to(tl.float32) + kw1 = tl.load(kw_ptr + d).to(tl.float32) + kw2 = tl.load(kw_ptr + HALF + d).to(tl.float32) + qb = qkv_ptr + row * W + h * HD + q1, q2 = _head_norm_rope( + tl.load(qb + d).to(tl.float32), tl.load(qb + HALF + d).to(tl.float32), qw1, qw2, + tl.load(cos_ptr + p * HD + d).to(tl.float32), tl.load(cos_ptr + p * HD + HALF + d).to(tl.float32), + tl.load(sin_ptr + p * HD + d).to(tl.float32), tl.load(sin_ptr + p * HD + HALF + d).to(tl.float32), + eps, HD, + ) # fmt: skip + j = tl.arange(0, BK) + cache = ((b * KVH + kvh) * MAXP).to(tl.int64) * HD + old = (j < first)[:, None] + kp = kc_ptr + cache + j[:, None] * HD + d[None, :] + vp = vc_ptr + cache + j[:, None] * HD + d[None, :] + k1 = tl.load(kp, mask=old, other=0.0).to(tl.float32) + k2 = tl.load(kp + HALF, mask=old, other=0.0).to(tl.float32) + v1 = tl.load(vp, mask=old, other=0.0).to(tl.float32) + v2 = tl.load(vp + HALF, mask=old, other=0.0).to(tl.float32) + for n in tl.static_range(NQ): + # this call's position first + n, computed from its own QKV row + on = first + n + kb = qkv_ptr + (b * NQ + n).to(tl.int64) * W + (H + kvh) * HD + vb = qkv_ptr + (b * NQ + n).to(tl.int64) * W + (H + KVH + kvh) * HD + kn1, kn2 = _head_norm_rope( + tl.load(kb + d).to(tl.float32), tl.load(kb + HALF + d).to(tl.float32), kw1, kw2, + tl.load(cos_ptr + on * HD + d).to(tl.float32), tl.load(cos_ptr + on * HD + HALF + d).to(tl.float32), + tl.load(sin_ptr + on * HD + d).to(tl.float32), tl.load(sin_ptr + on * HD + HALF + d).to(tl.float32), + eps, HD, + ) # fmt: skip + vn1 = tl.load(vb + d).to(tl.float32) + vn2 = tl.load(vb + HALF + d).to(tl.float32) + sel = (j == on)[:, None] + k1 = tl.where(sel, kn1[None, :], k1) + k2 = tl.where(sel, kn2[None, :], k2) + v1 = tl.where(sel, vn1[None, :], v1) + v2 = tl.where(sel, vn2[None, :], v2) + wm = (d < HALF) & (qi == n) & (h % (H // KVH) == 0) + dtc = kc_ptr.dtype.element_ty + tl.store(kc_ptr + cache + on * HD + d, kn1.to(dtc), mask=wm) + tl.store(kc_ptr + cache + on * HD + HALF + d, kn2.to(dtc), mask=wm) + tl.store(vc_ptr + cache + on * HD + d, vn1.to(dtc), mask=wm) + tl.store(vc_ptr + cache + on * HD + HALF + d, vn2.to(dtc), mask=wm) + km = j <= p + s = (tl.sum(k1 * q1[None, :], axis=1) + tl.sum(k2 * q2[None, :], axis=1)) * scale + s = tl.where(km, s, -float("inf")) + e = tl.exp(s - tl.max(s, axis=0)) + e = tl.where(km, e, 0.0) + w = e / tl.sum(e, axis=0) + ob = out_ptr + row * (H * HD) + h * HD + dto = out_ptr.dtype.element_ty + tl.store(ob + d, tl.sum(v1 * w[:, None], axis=0).to(dto)) + tl.store(ob + HALF + d, tl.sum(v2 * w[:, None], axis=0).to(dto)) + + +@triton.jit +def _cp_sample_kernel( + logits_ptr, # [B, V] bf16 + u_ptr, # uniforms, row stride u_stride, contiguous over V + u_stride, + codes_ptr, # [B, G] int64 + step, + G: tl.constexpr, + table_ptr, # [V, HID] next-step input table, or dummy + next_ptr, # [B, HID] + inv_temperature, + V: tl.constexpr, + TOPK: tl.constexpr, + HID: tl.constexpr, + HAS_NEXT: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + idx = tl.arange(0, V) + value = tl.load(logits_ptr + row * V + idx) + scaled = (value.to(tl.float32) * inv_temperature).to(value.dtype) + if TOPK > 0: + # order-preserving 16-bit keys of the BF16 values + bits = scaled.to(tl.int16, bitcast=True).to(tl.int32) & 0xFFFF + # Top-k keeps all numeric ties at the cutoff, including both zeros. + bits = tl.where((bits & 0x7FFF) == 0, 0, bits) + key = tl.where(bits >= 0x8000, 0xFFFF - bits, bits | 0x8000) + the = 0 + for bit in tl.static_range(15, -1, -1): + cand = the | (1 << bit) + cnt = tl.sum((key >= cand).to(tl.int32), axis=0) + the = tl.where(cnt >= TOPK, cand, the) + sf = tl.where(key >= the, scaled.to(tl.float32), -float("inf")) + else: + sf = scaled.to(tl.float32) + u = tl.load(u_ptr + row * u_stride + idx) + score = sf - tldevice.log(-tldevice.log(u)) + max_score = tl.max(score, 0) + chosen = tl.min(tl.where(score == max_score, idx, V), 0) + first_nan = tl.min(tl.where(score != score, idx, V), 0) + code = tl.where(first_nan < V, first_nan, chosen) + tl.store(codes_ptr + row * G + step, code.to(tl.int64)) + if HAS_NEXT: + c = tl.arange(0, HID) + tl.store(next_ptr + row * HID + c, tl.load(table_ptr + code.to(tl.int64) * HID + c)) + + +class FusedCodePredictor: + """Execution object built from a loaded ``Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM``.""" + + def __init__(self, predictor, max_batch: int) -> None: + model = predictor.model + weight = next(model.parameters()) + dev, dt = weight.device, weight.dtype + if dt != torch.bfloat16: + raise ValueError("fused code predictor requires BF16 weights") + cfg = predictor.config + self.groups = int(cfg.num_code_groups) + self.vocab = int(cfg.vocab_size) + self.hidden = int(cfg.hidden_size) + attn0 = model.layers[0].self_attn + self.heads, self.kv_heads, self.head_dim = attn0.num_heads, attn0.num_kv_heads, attn0.head_dim + self.scale = attn0.scaling + self.eps = float(model.norm.variance_epsilon) + self.max_pos = attn0.max_seq + self.max_batch = max_batch + f32 = lambda t: t.detach().to(device=dev, dtype=torch.float32) # noqa: E731 + c = lambda t: t.to(dt).contiguous() # noqa: E731 + proj = predictor.small_to_mtp_projection + if isinstance(proj, torch.nn.Identity): + pw, pb = None, None + else: + pw, pb = f32(proj.weight), f32(proj.bias) + self.proj_wt = c(pw.t()) if pw is not None else None + self.proj_b = c(pb) if pb is not None else None + tables = [] + for emb in model.codec_embedding: + e = f32(emb.weight) + tables.append(c(e @ pw.t() + pb) if pw is not None else c(e)) + self.tables = tables + self.layers = [] + for layer in model.layers: + a = layer.self_attn + self.layers.append( + dict( + ln1=c(f32(layer.input_layernorm.weight)), + ln2=c(f32(layer.post_attention_layernorm.weight)), + eps1=float(layer.input_layernorm.variance_epsilon), + eps2=float(layer.post_attention_layernorm.variance_epsilon), + qkv=c(f32(a.qkv_proj.weight).t()), + qn=c(f32(a.q_norm.weight)), + kn=c(f32(a.k_norm.weight)), + head_eps=float(a.q_norm.variance_epsilon), + o=c(f32(a.o_proj.weight).t()), + gu=c(f32(layer.mlp.gate_up_proj.weight).t()), + down=c(f32(layer.mlp.down_proj.weight).t()), + ) + ) + if a.qkv_proj.bias is not None: + raise ValueError("fused code predictor assumes bias-free attention") + self.norm_w = c(f32(model.norm.weight)) + self.lm_heads = [c(f32(h.weight).t()) for h in predictor.lm_head] + rot = model.rotary_emb + self.cos = rot.cos_cached.to(device=dev, dtype=dt).contiguous() + self.sin = rot.sin_cached.to(device=dev, dtype=dt).contiguous() + kv = (len(self.layers), 2, max_batch, self.kv_heads, self.max_pos, self.head_dim) + self.cache = torch.zeros(kv, device=dev, dtype=dt) + self._dummy = torch.zeros(1, device=dev, dtype=dt) + + def _layer(self, x: torch.Tensor, li: int, B: int, nq: int, first: int) -> None: + L = self.layers[li] + rows = x.shape[0] + h = torch.empty_like(x) + _rms_rows_kernel[(rows,)](x, x.stride(0), L["ln1"], h, L["eps1"], N=self.hidden) + qkv = torch.mm(h, L["qkv"]) + att = torch.empty(rows, self.heads * self.head_dim, device=x.device, dtype=x.dtype) + _cp_attention_kernel[(B, self.heads, nq)]( + qkv, self.cache[li, 0], self.cache[li, 1], L["qn"], L["kn"], self.cos, self.sin, att, + first, L["head_eps"], self.scale, + NQ=nq, H=self.heads, KVH=self.kv_heads, HD=self.head_dim, MAXP=self.max_pos, + BK=triton.next_power_of_2(self.max_pos), num_warps=1, + ) # fmt: skip + x.addmm_(att, L["o"]) + _rms_rows_kernel[(rows,)](x, x.stride(0), L["ln2"], h, L["eps2"], N=self.hidden) + gu = torch.mm(h, L["gu"]) + inter = gu.shape[1] // 2 + a = torch.empty(rows, inter, device=x.device, dtype=x.dtype) + _silu_mul_kernel[(rows, triton.cdiv(inter, 1024))](gu, a, N=inter, BN=1024) + x.addmm_(a, L["down"]) + + @torch.inference_mode() + def __call__( + self, + layer0_code: torch.Tensor, + layer0_embed: torch.Tensor, + last_talker_hidden: torch.Tensor, + inv_temperature: float, + top_k: int, + sample_uniforms: torch.Tensor, + ) -> torch.Tensor: + B = int(layer0_code.shape[0]) + if B > self.max_batch: + raise ValueError("fused code predictor batch exceeds its capacity") + dev, dt = layer0_embed.device, self.norm_w.dtype + inp = torch.stack([last_talker_hidden.reshape(B, -1), layer0_embed.reshape(B, -1)], 1).to(dt) + inp = inp.reshape(2 * B, -1) + x = torch.addmm(self.proj_b, inp, self.proj_wt) if self.proj_wt is not None else inp.clone() + codes = torch.empty(B, self.groups, device=dev, dtype=torch.int64) + codes[:, 0] = layer0_code.reshape(B) + u = sample_uniforms + nq, first = 2, 0 + for step in range(1, self.groups): + for li in range(len(self.layers)): + self._layer(x, li, B, nq, first) + last = x[nq - 1 :: nq] if nq > 1 else x + hf = torch.empty(B, self.hidden, device=dev, dtype=dt) + _rms_rows_kernel[(B,)](last, last.stride(0), self.norm_w, hf, self.eps, N=self.hidden) + logits = torch.mm(hf, self.lm_heads[step - 1]) + has_next = step < self.groups - 1 + nxt = torch.empty(B, self.hidden, device=dev, dtype=dt) if has_next else self._dummy + us = u[:, step - 1] + _cp_sample_kernel[(B,)]( + logits, us, us.stride(0), codes, step, self.groups, + self.tables[step - 1] if has_next else self._dummy, nxt, inv_temperature, + V=self.vocab, TOPK=top_k, HID=self.hidden, HAS_NEXT=has_next, num_warps=4, + ) # fmt: skip + x = nxt + first = step + 1 + nq = 1 + return codes diff --git a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py index dd753861d..08c1f7c3e 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py @@ -72,3 +72,33 @@ QWEN3_TTS_PIPELINE = PipelineConfig( ), ), ) + + +# Single-stage variant: the Talker decodes each frame with the stateful +# streaming codec decoder and emits PCM as the final output (enable with the +# ``talker_stream_decode`` model option). No Code2Wav stage, no connector. +QWEN3_TTS_FUSED_PIPELINE = PipelineConfig( + model_type="qwen3_tts_fused", + default_deploy_config_name="qwen3_tts_fused_single_gpu.yaml", + single_stage_async_chunk=True, + model_arch="Qwen3TTSTalkerForConditionalGeneration", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="qwen3_tts", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + final_output=True, + final_output_type="audio", + # The API output processor tags ``model_outputs`` by this modality. + engine_output_type="audio", + supports_running_prefix_cache_reset=False, + supports_native_mrv2_data_plane=True, + sampling_constraints={ + "detokenize": False, + "stop_token_ids": [2150], + }, + ), + ), +) diff --git a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py index 45d887ee2..7a63a3072 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py +++ b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py @@ -42,6 +42,7 @@ from vllm.logger import init_logger from vllm.multimodal.audio import AudioResampler from vllm_omni.utils.audio import mel_filter_bank +from vllm_omni.utils.device_copy import index_to_device, to_device_nonblocking if TYPE_CHECKING: from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSTalkerConfig @@ -160,7 +161,7 @@ def coerce_token_ids(value: object, *, device: torch.device) -> torch.Tensor | N ids = ids.unsqueeze(0) if ids.ndim != 2 or ids.numel() == 0: return None - return ids.to(device=device, dtype=torch.long).contiguous() + return to_device_nonblocking(ids.to(dtype=torch.long).contiguous(), device) # --------------------------------------------------------------------------- @@ -366,6 +367,9 @@ class Qwen3TTSPromptEmbedsBuilder: # Reusing them avoids launching the embedding/projection pair once per # request during prompt construction. self._projected_token_cache: dict[tuple[str, tuple[int, ...]], torch.Tensor] = {} + # CustomVoice prompt pieces that do not depend on the text, per + # (device, speaker, language): bounded by the model's voices. + self._custom_voice_prompt_cache: dict[tuple, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {} self._ref_audio_artifact_cache_max_entries = int(ref_audio_artifact_cache_max_entries) self._ref_audio_artifact_cache: OrderedDict[str, dict[str, torch.Tensor | bool]] = OrderedDict() @@ -1024,7 +1028,9 @@ class Qwen3TTSPromptEmbedsBuilder: dev = self._device() input_ids = coerce_token_ids(info_dict.pop(PRECOMPUTED_TEXT_IDS_KEY, None), device=dev) if input_ids is None: - input_ids = tok(build_assistant_text(text), return_tensors="pt", padding=False)["input_ids"].to(device=dev) + input_ids = to_device_nonblocking( + tok(build_assistant_text(text), return_tensors="pt", padding=False)["input_ids"], dev + ) # Optional instruct prefix. instruct = (info_dict.get("instruct") or [""])[0] @@ -1411,35 +1417,44 @@ class Qwen3TTSPromptEmbedsBuilder: raise ValueError(f"Unsupported speaker: {speaker}") spk_id = spk_id_map[speaker] # Keep it at least 1D; embedding on a 0-d tensor can return 1D. - spk_tensor = torch.tensor([spk_id], device=input_ids.device, dtype=torch.long) + spk_tensor = index_to_device([spk_id], input_ids.device) spk_embed = codec_embed(spk_tensor) if spk_embed.ndim in (1, 2): spk_embed = spk_embed.view(1, 1, -1) speaker_embed = spk_embed codec_input = torch.cat([codec_input_0, speaker_embed, codec_input_1], dim=1) - role_embed = text_projection(text_embedding(input_ids[:, :3])) - codec_prefix = torch.cat((tts_pad_embed.expand(-1, codec_input.shape[1] - 2, -1), tts_bos_embed), dim=1) - codec_prefix = codec_prefix + codec_input[:, :-1] - talker_prompt = torch.cat((role_embed, codec_prefix), dim=1) + # The role tokens, speaker/language codec prefix, codec pad row and + # closing row do not depend on the text: cache them per speaker and + # language so a request only embeds and projects its own text. + # A text starting with whitespace could merge with the template's + # newline token, so it keeps the per-request role embedding. + cache_key = None + if non_streaming_mode and instruct_embed is None and not text[:1].isspace(): + cache_key = ( + str(input_ids.device), + int(spk_id), + codec_prefill_list[0][2] if len(codec_prefill_list[0]) == 4 else None, + ) + cv_cache = self._custom_voice_prompt_cache + cached = cv_cache.get(cache_key) if cache_key is not None else None + if cached is None: + role_embed = text_projection(text_embedding(input_ids[:, :3])) + codec_prefix = torch.cat((tts_pad_embed.expand(-1, codec_input.shape[1] - 2, -1), tts_bos_embed), dim=1) + codec_prefix = codec_prefix + codec_input[:, :-1] + prefix = torch.cat((role_embed, codec_prefix), dim=1) + pad_row = codec_embed(self._long_tensor([talker_config.codec_pad_id], input_ids.device)).view(1, 1, -1) + tail = tts_pad_embed + codec_embed(self._long_tensor([talker_config.codec_bos_id], input_ids.device)) + if cache_key is not None: + cv_cache[cache_key] = (prefix, pad_row, tail) + else: + prefix, pad_row, tail = cached + talker_prompt = prefix if non_streaming_mode: text_all = text_projection(text_embedding(input_ids[:, 3:-5])) text_all = torch.cat([text_all, tts_eos_embed], dim=1) - pad_ids = torch.full( - (1, int(text_all.shape[1])), - int(talker_config.codec_pad_id), - device=input_ids.device, - dtype=torch.long, - ) - talker_prompt = torch.cat( - [ - talker_prompt, - text_all + codec_embed(pad_ids), - tts_pad_embed + codec_embed(self._long_tensor([talker_config.codec_bos_id], input_ids.device)), - ], - dim=1, - ) + talker_prompt = torch.cat([talker_prompt, text_all + pad_row, tail], dim=1) trailing_text_hidden = tts_pad_embed else: first_text = text_projection(text_embedding(input_ids[:, 3:4])) + codec_input[:, -1:] diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py index 9620d0d86..493b3b2a2 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py @@ -65,6 +65,8 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapp self._kv_requested = self._parse_bool_config(extra.get("code_predictor_kv_cache")) self._fused_sampling = self._parse_bool_config(extra.get("code_predictor_fused_sampling")) self._frame_cache: FrameLocalKVCache | None = None + self._fused_requested = self._parse_bool_config(extra.get("code_predictor_fused")) + self._fused = None def _setup_compile(self) -> None: if self._compiled_model_fwd is not None: @@ -100,11 +102,73 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapp cache(self._proj_buf, batch, step) self._synchronize_warmup(weight.device) self._frame_cache = cache + if self._fused_requested: + from .fused_code_predictor import FusedCodePredictor + + self._fused = FusedCodePredictor(self, max_batch) + # Compile every kernel variant outside any graph capture. + hidden = int(self.talker_config.hidden_size) + for batch in range(1, min(2, max_batch) + 1): + zeros = torch.zeros(batch, 1, hidden, device=weight.device, dtype=weight.dtype) + uniforms = torch.full( + (batch, self._num_groups - 1, int(self.config.vocab_size)), 0.5, device=weight.device + ) + codes = torch.zeros(batch, 1, device=weight.device, dtype=torch.long) + self._fused(codes, zeros, zeros, 1.0, 50, uniforms) + self._synchronize_warmup(weight.device) + logger.info("Qwen3-TTS fused code predictor enabled (capacity %d)", max_batch) # Preserve the shared wrapper's initialization sentinel. Its full # re-prefill callable is unused while _frame_cache is active. self._compiled_model_fwd = self.model.forward logger.info("Qwen3-TTS frame-local KV warmed for batches %s", self._bucket_sizes) + @torch.inference_mode() + def forward( + self, + layer0_code: torch.Tensor, + layer0_embed: torch.Tensor, + last_talker_hidden: torch.Tensor, + do_sample: bool = True, + temperature: float = 0.9, + top_k: int = 50, + top_p: float = 1.0, + generator: torch.Generator | None = None, + generators: Sequence[torch.Generator | None] | None = None, + sample_uniforms: torch.Tensor | None = None, + ): + self._validate_sampling_inputs(int(layer0_code.shape[0]), generators, sample_uniforms) + if self._fused_requested: + self._setup_compile() + if ( + self._fused is not None + and do_sample + and temperature > 0 + and top_p == 1.0 + and sample_uniforms is not None + and sample_uniforms.dtype == torch.float32 + and sample_uniforms.is_contiguous() + ): + return self._fused( + layer0_code, + layer0_embed, + last_talker_hidden, + 1.0 / temperature, + int(top_k), + sample_uniforms, + ) + return super().forward( + layer0_code, + layer0_embed, + last_talker_hidden, + do_sample=do_sample, + temperature=temperature, + top_k=top_k, + top_p=top_p, + generator=generator, + generators=generators, + sample_uniforms=sample_uniforms, + ) + def _predict_step_logits( self, proj_buf: torch.Tensor, diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index e05ac3871..6e199454f 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py @@ -27,6 +27,7 @@ from vllm.sequence import IntermediateTensors from vllm_omni.data_entry_keys import OmniPayload from vllm_omni.model_executor.models.output_templates import OmniOutput, OwnedBatchTensor +from vllm_omni.utils.device_copy import index_to_device, to_device_nonblocking from vllm_omni.utils.speaker_cache import ( get_speaker_cache, iter_custom_voice_profiles, @@ -39,6 +40,7 @@ from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSSpeakerEncoderConfi from .first_audio import talker_first_audio_enabled from .prompt_embeds_builder import PRECOMPUTED_TEXT_IDS_KEY, Qwen3TTSPromptEmbedsBuilder, resolve_x_vector_only from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM +from .stream_decode import stream_ref_context_frames, talker_stream_decode_enabled from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Encoder @@ -411,7 +413,21 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): # A frame's residual codebooks depend only on the Talker hidden and # CB0 of the step that sampled it, so MRV2 may complete the frame at # the end of that step instead of in the next step's preprocess. - self.mtp_eager_frames = talker_first_audio_enabled(vllm_config) + self.stream_decode = talker_stream_decode_enabled(vllm_config) + predictor = Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM + extra = predictor._stage_connector_extra_config(vllm_config) + self.stream_first_audio = self.stream_decode and predictor._parse_bool_config( + extra.get("talker_stream_first_audio") + ) + # Model-local voice-clone priming boundary; each deployment selects + # its reference context independently of the two-stage codec profile. + self.stream_ref_context_frames = stream_ref_context_frames(vllm_config) if self.stream_decode else 0 + self.mtp_eager_frames = talker_first_audio_enabled(vllm_config) or self.stream_decode + self.stream_decoder = None + self.stream_graphs = None + self.stream_prime_graphs = None + self.stream_sample_rate = 0 + self.stream_chunk_frames = 25 # The runners bypass only the outer whole-MTP graph when explicit # generators are present, so seeded requests can still share one raw # batched MTP call with independent per-row streams. @@ -718,6 +734,11 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): # -------------------- Omni multimodal output plumbing -------------------- + @property + def mm_outputs_fresh_per_step(self) -> bool: + """Stream-decode outputs come from :meth:`_make_eager_omni_output`, allocated per step.""" + return self.stream_decoder is not None and self.eager_frames_active + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: if isinstance(model_outputs, OmniOutput): return model_outputs @@ -895,6 +916,13 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): "first_audio": torch.zeros((num_tokens,), dtype=torch.int8, device=hidden.device), }, } + if self.stream_decoder is not None: + # Filled per sampled row by the runner's eager MTP; token-major so + # each request's slice is its frames of this step. + spf = int(self.stream_decoder.spf) + mm["model_outputs"] = torch.zeros((num_tokens, spf), dtype=torch.float32, device=hidden.device) + sr = torch.tensor(self.stream_sample_rate, dtype=torch.int32) + mm["sr"] = [sr] * len(info_dicts) ref_rows = [ index for index, info in enumerate(info_dicts) @@ -1442,8 +1470,72 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): self.first_frame_decoder = Qwen3TTSFirstFrameDecoder(self.model_path) decoder_loaded = self.first_frame_decoder.load(self.vllm_config) loaded = set(loaded) | {f"first_frame_decoder.{name}" for name in decoder_loaded} + elif self.stream_decode: + # The in-stage decoder owns every frame. Optional first-audio + # delivery uses its PCM; later chunks use regular step outputs. + from .first_frame_decoder import Qwen3TTSFirstFrameDecoder + from .tokenizer_12hz.streaming_decoder import StreamingCodecDecoder + + holder = Qwen3TTSFirstFrameDecoder(self.model_path) + decoder_loaded = holder.load(self.vllm_config) + loaded = set(loaded) | {f"stream_decoder.{name}" for name in decoder_loaded} + slots = int(self.vllm_config.scheduler_config.max_num_seqs) + self.stream_decoder = StreamingCodecDecoder(holder.decoder, num_slots=slots, dtype=torch.bfloat16) + self.stream_sample_rate = int(holder.sample_rate) + logger.info("Qwen3-TTS Talker stream decode enabled (%d slots)", slots) return loaded + def capture_stream_decode_graphs(self, batch_sizes: list[int]) -> None: + if self.stream_decoder is None or self.stream_graphs is not None: + return + from .tokenizer_12hz.streaming_decoder import StreamingDecodeGraphs + + with torch.inference_mode(): + self.stream_graphs = StreamingDecodeGraphs(self.stream_decoder, batch_sizes) + frames = min(self.stream_ref_context_frames, self.stream_chunk_frames, self.stream_decoder.max_frames) + if self.config.tts_model_type == "base" and frames > 1: + # Bound the multi-frame graph workspace by the configured + # token budget; larger priming groups keep the eager path. + budget = max(1, self.vllm_config.scheduler_config.max_num_batched_tokens // frames) + prime_sizes = sorted({1, *(size for size in batch_sizes if size <= budget)}) + self.stream_prime_graphs = StreamingDecodeGraphs(self.stream_decoder, prime_sizes, frames=frames) + logger.info("Captured Talker stream decode graphs for batch sizes %s", self.stream_graphs.sizes) + if self.stream_prime_graphs is not None: + logger.info( + "Captured Talker reference priming graphs for %d frames and batch sizes %s", + self.stream_prime_graphs.frames, + self.stream_prime_graphs.sizes, + ) + + def get_stream_ref_context(self, info: OmniPayload) -> torch.Tensor | None: + """Model-owned reference-code layout and context boundary for a PCM slot.""" + codes = info.get("codes") + ref = codes.get("ref") if isinstance(codes, dict) else None + if not isinstance(ref, torch.Tensor) or not ref.numel(): + return None + return ref.reshape(-1, int(self.talker_config.num_code_groups))[-self.stream_ref_context_frames :] + + def prime_stream_decoder(self, primes: list[tuple[int, torch.Tensor]]) -> None: + """Prime equal-length reference groups without padding or aliasing live slots.""" + if not primes: + return + stream = self.stream_decoder + if stream is None: + raise RuntimeError("Reference priming requires the single-stage PCM decoder") + groups: dict[int, list[tuple[int, torch.Tensor]]] = {} + for req_idx, ref in primes: + groups.setdefault(int(ref.shape[0]), []).append((req_idx, ref)) + for items in groups.values(): + codes = to_device_nonblocking(torch.stack([ref for _idx, ref in items]), stream.device).to(torch.int32) + n = len(items) + slots = index_to_device([idx for idx, _ref in items], stream.device, dtype=torch.int32) + for t0 in range(0, int(codes.shape[1]), self.stream_chunk_frames): + pos = index_to_device([t0] * n, stream.device, dtype=torch.int32) + chunk = codes[:, t0 : t0 + self.stream_chunk_frames].contiguous() + graphs = self.stream_prime_graphs + decode = graphs if graphs is not None and chunk.shape[1] == graphs.frames else stream + decode(chunk, slots, pos) + def _build_stacked_codec_embed(self) -> None: embeds = self.code_predictor.get_input_embeddings() if not embeds: diff --git a/vllm_omni/model_executor/models/qwen3_tts/stream_decode.py b/vllm_omni/model_executor/models/qwen3_tts/stream_decode.py new file mode 100644 index 000000000..747e1f0b4 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/stream_decode.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Eligibility and model options for the single-stage CUDA PCM pipeline.""" + +import torch +from vllm.config import VllmConfig + +from vllm_omni.platforms import current_omni_platform + + +def talker_stream_decode_enabled(vllm_config: VllmConfig) -> bool: + """The Talker decodes every frame itself and emits PCM as its final output.""" + from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor + + extra = Predictor._stage_connector_extra_config(vllm_config) + model = vllm_config.model_config + parallel = vllm_config.parallel_config + requested = Predictor._parse_bool_config(extra.get("talker_stream_decode")) + if not requested: + if getattr(model, "engine_output_type", None) == "audio": + raise ValueError("The single-stage Qwen3-TTS pipeline requires talker_stream_decode=True") + if Predictor._parse_bool_config(extra.get("talker_stream_first_audio")): + raise ValueError("talker_stream_first_audio requires the single-stage talker_stream_decode path") + return False + if getattr(model, "engine_output_type", None) != "audio" or not getattr(model, "final_output", False): + raise ValueError( + "talker_stream_decode requires a final audio-output Talker; " + "select qwen3_tts_fused_single_gpu.yaml instead of enabling it on the two-stage pipeline" + ) + if getattr(model, "supports_running_prefix_cache_reset", True): + raise ValueError( + "The single-stage PCM pipeline must disable running prefix-cache reset; " + "select qwen3_tts_fused_single_gpu.yaml" + ) + if not ( + current_omni_platform.is_cuda() + and torch.device(vllm_config.device_config.device).type == "cuda" + and bool(getattr(model, "use_v2_model_runner", False)) + and bool(getattr(model, "async_chunk", False)) + and parallel.tensor_parallel_size == 1 + and parallel.pipeline_parallel_size == 1 + and parallel.distributed_executor_backend in (None, "uni") + and not vllm_config.cache_config.enable_prefix_caching + ): + raise ValueError( + "talker_stream_decode requires CUDA, MRv2, async_chunk=True, TP=PP=1, " + "an in-process worker and enable_prefix_caching=False; " + "use qwen3_tts_fused_single_gpu.yaml" + ) + if Predictor._parse_bool_config(extra.get("talker_first_audio")): + raise ValueError("talker_stream_decode and talker_first_audio are mutually exclusive") + return True + + +def stream_ref_context_frames(vllm_config: VllmConfig) -> int: + """Reference-code context frames for voice-clone priming in the Talker.""" + from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor + + extra = Predictor._stage_connector_extra_config(vllm_config) + return int(extra.get("ref_code_context_frames") or extra.get("codec_left_context_frames", 25)) diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/streaming_decoder.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/streaming_decoder.py new file mode 100644 index 000000000..b1dffed00 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/streaming_decoder.py @@ -0,0 +1,799 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# ruff: noqa: N803 - Triton constexpr parameters use kernel-style capitals. +"""Stateful, fused streaming decoder for the Qwen3-TTS 12 Hz codec. + +Every layer of the codec decoder is causal, so what a new frame needs from the +past is small and fixed: each causal conv's last ``(k - 1) * dilation`` input +rows, each overlapping transposed conv's last GEMM row, and each attention +layer's last ``sliding_window`` keys and values. A request owns a slot of +per-layer state holding exactly that, and a call decodes ``T`` new frames for +each of ``B`` requests, reading and advancing the state. The output equals the +whole-utterance decode (``Qwen3TTSTokenizerV2Decoder._forward_exact``) frame by +frame, up to floating point rounding: nothing is recomputed and nothing is +truncated. + +Layout and fusion: + +- Activations are time-major ``[rows, C]``; every dense conv is a tap-major + im2col (one Triton kernel that also applies the pending bias and SnakeBeta, + and reads/writes the history) followed by one cuBLAS GEMM. +- Linear chains fold at load: the RVQ output projections, ``pre_conv`` and the + transformer's ``input_proj`` become one ``k = 3`` conv over the concatenated + codebook embeddings; LayerScale folds into ``o_proj``/``down_proj``; ConvNeXt + ``gamma`` folds into ``pwconv2``. +- A bias that feeds a residual stream is carried to that stream's consumer + (the next im2col or transposed-conv GEMM) instead of being added in place. +- Residual branches accumulate into their stream with in-place ``addmm_``. + +History buffers are double-buffered per slot: a call reads half ``parity[slot]`` +and writes the other half, then flips the parity, so no kernel overwrites rows +another program of the same launch still reads. A row whose first frame index +``pos`` is 0 reads zeros, which are the convs' causal padding, and attention +only looks back ``min(pos + 1, window)`` positions, so a new request needs no +state reset. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +from torch import nn +from vllm.triton_utils import tl, triton + +from vllm_omni.platforms import current_omni_platform + +_KERNEL = 7 + + +# --------------------------------------------------------------------------- kernels + + +@triton.jit +def _rvq_gather_kernel(codes_ptr, books_ptr, out_ptr, NQ: tl.constexpr, CB: tl.constexpr, HALF: tl.constexpr): + row = tl.program_id(0) + offs = tl.arange(0, HALF) + c0 = tl.load(codes_ptr + row * NQ).to(tl.int64) + c0 = tl.where((c0 >= 0) & (c0 < CB), c0, 0) # EOS/special rows decode as code 0; callers drop them + first = tl.load(books_ptr + c0 * HALF + offs).to(tl.float32) + acc = tl.zeros([HALF], dtype=tl.float32) + for q in range(1, NQ): + c = tl.load(codes_ptr + row * NQ + q).to(tl.int64) + c = tl.where((c >= 0) & (c < CB), c, 0) + acc += tl.load(books_ptr + (q * CB + c) * HALF + offs).to(tl.float32) + dt = out_ptr.dtype.element_ty + tl.store(out_ptr + row * 2 * HALF + offs, first.to(dt)) + tl.store(out_ptr + row * 2 * HALF + HALF + offs, acc.to(dt)) + + +@triton.jit +def _fast_sin(y): + # Two-step reduction to [-pi, pi], then the hardware approximation + # (absolute error ~4e-7 there): accurate sin's instruction count, not + # memory, bounds every kernel that applies SnakeBeta. + k = tl.extra.cuda.libdevice.rint(y * 0.15915494309189535) + r = tl.fma(-k, -1.7484555314695172e-07, tl.fma(-k, 6.2831854820251465, y)) + return tl.inline_asm_elementwise("sin.approx.f32 $0, $1;", "=r,r", [r], dtype=tl.float32, is_pure=True, pack=1) + + +@triton.jit +def _snake(v, a, ib): + s = _fast_sin(v * a) + return v + ib * s * s + + +@triton.jit +def _im2col_kernel( + x_ptr, # [B * T, C] layer input rows + bias_ptr, # [C] pending bias (added to x rows only) or dummy + alpha_ptr, # [C] exp(alpha) or dummy + ibeta_ptr, # [C] 1 / (exp(beta) + eps) or dummy + hist_ptr, # [2, S, H, C] transformed history rows + slots_ptr, + pos_ptr, + par_ptr, + col_ptr, # [B * T, K * C] + T, + S, + C: tl.constexpr, + K: tl.constexpr, + D: tl.constexpr, + H: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_SNAKE: tl.constexpr, + BU: tl.constexpr, + BC: tl.constexpr, + EXT: tl.constexpr = False, +): + """Source row ``u`` of ``[history (H) | x (T)]``: scatter to every tap that reads it, save the last H.""" + b = tl.program_id(0) + u = tl.program_id(1) * BU + tl.arange(0, BU) + c = tl.program_id(2) * BC + tl.arange(0, BC) + cm = c < C + n_src = H + T + um = u < n_src + slot = tl.load(slots_ptr + b).to(tl.int64) + live = tl.load(pos_ptr + b) != 0 + par = tl.load(par_ptr + slot).to(tl.int64) + # history rows + from_hist = u < H + hrow = tl.where(from_hist, u, 0) + hv = tl.load( + hist_ptr + ((par * S + slot) * H + hrow)[:, None] * C + c[None, :], + mask=(from_hist & um & live)[:, None] & cm[None, :], + other=0.0, + ).to(tl.float32) + xrow = tl.where(from_hist, 0, u - H) + xv = tl.load( + x_ptr + (b * T + xrow).to(tl.int64)[:, None] * C + c[None, :], + mask=((~from_hist) & um)[:, None] & cm[None, :], + other=0.0, + ).to(tl.float32) + if HAS_BIAS: + xv += tl.load(bias_ptr + c, mask=cm, other=0.0).to(tl.float32)[None, :] + if HAS_SNAKE: + a = tl.load(alpha_ptr + c, mask=cm, other=0.0).to(tl.float32)[None, :] + ib = tl.load(ibeta_ptr + c, mask=cm, other=0.0).to(tl.float32)[None, :] + xv = _snake(xv, a, ib) + dt = col_ptr.dtype.element_ty + v = tl.where(from_hist[:, None], hv, xv.to(dt).to(tl.float32)).to(dt) + if EXT: + tl.store(col_ptr + (b * n_src + u).to(tl.int64)[:, None] * C + c[None, :], v, mask=um[:, None] & cm[None, :]) + else: + for j in tl.static_range(K): + t = u - j * D + tm = (t >= 0) & (t < T) & um + tl.store( + col_ptr + (b * T + tl.where(tm, t, 0)).to(tl.int64)[:, None] * (K * C) + j * C + c[None, :], + v, + mask=tm[:, None] & cm[None, :], + ) + if H > 0: + hm = (u >= T) & um + tl.store( + hist_ptr + (((1 - par) * S + slot) * H + tl.where(hm, u - T, 0))[:, None] * C + c[None, :], + v, + mask=hm[:, None] & cm[None, :], + ) + + +@triton.jit +def _rms_norm_kernel(x_ptr, w_ptr, out_ptr, eps, N: tl.constexpr): + row = tl.program_id(0).to(tl.int64) + offs = tl.arange(0, N) + x = tl.load(x_ptr + row * N + offs).to(tl.float32) + var = tl.sum(x * x, axis=0) / N + dt = out_ptr.dtype.element_ty + y = (x * tl.rsqrt(var + eps)).to(dt).to(tl.float32) + w = tl.load(w_ptr + offs).to(tl.float32) + tl.store(out_ptr + row * N + offs, (w * y).to(dt)) + + +@triton.jit +def _rope_kv_write_kernel( + qkv_ptr, # [B * T, 3 * NH * HD] + q_ptr, # [B * T, NH * HD] + k_ring_ptr, # [S, RING, NH * HD] + v_ring_ptr, + slots_ptr, + pos_ptr, + T, + theta, + NH: tl.constexpr, + HD: tl.constexpr, + RING: tl.constexpr, +): + r = tl.program_id(0) + h = tl.program_id(1) + b = r // T + p = tl.load(pos_ptr + b) + r % T + slot = tl.load(slots_ptr + b).to(tl.int64) + half: tl.constexpr = HD // 2 + i = tl.arange(0, half) + inv_freq = tl.exp(-(2.0 * i.to(tl.float32) / HD) * tl.log(theta)) + ang = p.to(tl.float32) * inv_freq + dt = q_ptr.dtype.element_ty + cos = tl.cos(ang).to(dt).to(tl.float32) + sin = tl.sin(ang).to(dt).to(tl.float32) + W: tl.constexpr = NH * HD + base = qkv_ptr + r.to(tl.int64) * 3 * W + h * HD + q1 = tl.load(base + i).to(tl.float32) + q2 = tl.load(base + half + i).to(tl.float32) + k1 = tl.load(base + W + i).to(tl.float32) + k2 = tl.load(base + W + half + i).to(tl.float32) + vv = tl.arange(0, HD) + v = tl.load(base + 2 * W + vv) + qo = q_ptr + r.to(tl.int64) * W + h * HD + tl.store(qo + i, (q1 * cos - q2 * sin).to(dt)) + tl.store(qo + half + i, (q2 * cos + q1 * sin).to(dt)) + ring = (slot * RING + p % RING) * W + h * HD + tl.store(k_ring_ptr + ring + i, (k1 * cos - k2 * sin).to(dt)) + tl.store(k_ring_ptr + ring + half + i, (k2 * cos + k1 * sin).to(dt)) + tl.store(v_ring_ptr + ring + vv, v) + + +@triton.jit +def _attend_kernel( + q_ptr, # [B * T, NH * HD] + k_ring_ptr, + v_ring_ptr, + out_ptr, # [B * T, NH * HD] + slots_ptr, + pos_ptr, + T, + scale, + NH: tl.constexpr, + HD: tl.constexpr, + RING: tl.constexpr, + WINDOW: tl.constexpr, + BK: tl.constexpr, +): + r = tl.program_id(0) + h = tl.program_id(1) + b = r // T + p = tl.load(pos_ptr + b) + r % T + slot = tl.load(slots_ptr + b).to(tl.int64) + W: tl.constexpr = NH * HD + d = tl.arange(0, HD) + q = tl.load(q_ptr + r.to(tl.int64) * W + h * HD + d).to(tl.float32) + j = tl.arange(0, BK) + n = tl.minimum(p + 1, WINDOW) + km = j < n + kp = p - j # key positions p, p-1, ... + idx = (slot * RING + tl.where(km, kp, 0) % RING) * W + h * HD + k = tl.load(k_ring_ptr + idx[:, None] + d[None, :], mask=km[:, None], other=0.0).to(tl.float32) + s = tl.sum(k * q[None, :], axis=1) * scale + s = tl.where(km, s, -float("inf")) + m = tl.max(s, axis=0) + e = tl.exp(s - m) + e = tl.where(km, e, 0.0) + w = e / tl.sum(e, axis=0) + dt = out_ptr.dtype.element_ty + w = w.to(dt).to(tl.float32) + v = tl.load(v_ring_ptr + idx[:, None] + d[None, :], mask=km[:, None], other=0.0).to(tl.float32) + o = tl.sum(v * w[:, None], axis=0) + tl.store(out_ptr + r.to(tl.int64) * W + h * HD + d, o.to(dt)) + + +@triton.jit +def _silu_mul_kernel(gu_ptr, out_ptr, N: tl.constexpr, BN: tl.constexpr): + row = tl.program_id(0).to(tl.int64) + offs = tl.program_id(1) * BN + tl.arange(0, BN) + m = offs < N + g = tl.load(gu_ptr + row * 2 * N + offs, mask=m).to(tl.float32) + u = tl.load(gu_ptr + row * 2 * N + N + offs, mask=m).to(tl.float32) + dt = out_ptr.dtype.element_ty + sg = (g / (1.0 + tl.exp(-g))).to(dt).to(tl.float32) + tl.store(out_ptr + row * N + offs, (sg * u).to(dt), mask=m) + + +@triton.jit +def _dwconv_ln_kernel( + z_ptr, # [B * T, C] ConvNeXt input rows (true values) + w_ptr, # [K, C] + b_ptr, # [C] + lnw_ptr, + lnb_ptr, + hist_ptr, # [2, S, K - 1, C] + slots_ptr, + pos_ptr, + par_ptr, + out_ptr, # [B * T, C] LayerNorm output + T, + S, + eps, + C: tl.constexpr, + K: tl.constexpr, +): + b = tl.program_id(0) + t = tl.program_id(1) # 0 .. max(T, K - 1) + H: tl.constexpr = K - 1 + c = tl.arange(0, C) + slot = tl.load(slots_ptr + b).to(tl.int64) + live = tl.load(pos_ptr + b) != 0 + par = tl.load(par_ptr + slot).to(tl.int64) + dt = out_ptr.dtype.element_ty + hbase = hist_ptr + (par * S + slot) * H * C + acc = tl.load(b_ptr + c).to(tl.float32) + for j in tl.static_range(K): + u = t + j # index into [hist (H) | z (T)] + fh = u < H + hv = tl.load(hbase + tl.where(fh, u, 0) * C + c, mask=(c < C) & fh & live, other=0.0).to(tl.float32) + zrow = (b * T + tl.where(fh, 0, u - H)).to(tl.int64) + zv = tl.load(z_ptr + zrow * C + c, mask=(c < C) & (~fh) & (t < T), other=0.0).to(tl.float32) + acc += tl.where(fh, hv, zv) * tl.load(w_ptr + j * C + c).to(tl.float32) + h = acc.to(dt).to(tl.float32) + mean = tl.sum(h, axis=0) / C + hc = h - mean + var = tl.sum(hc * hc, axis=0) / C + y = hc * tl.rsqrt(var + eps) * tl.load(lnw_ptr + c).to(tl.float32) + tl.load(lnb_ptr + c).to(tl.float32) + tl.store(out_ptr + (b * T + t).to(tl.int64) * C + c, y.to(dt), mask=(c < C) & (t < T)) + # new history row t: source row T + t of [hist | z] + u = T + t + fh = u < H + hv = tl.load(hbase + tl.where(fh, u, 0) * C + c, mask=(c < C) & fh & live & (t < H), other=0.0) + zrow = (b * T + tl.where(fh, 0, u - H)).to(tl.int64) + zv = tl.load(z_ptr + zrow * C + c, mask=(c < C) & (~fh) & (t < H), other=0.0) + tl.store( + hist_ptr + (((1 - par) * S + slot) * H + tl.where(t < H, t, 0)) * C + c, + tl.where(fh, hv, zv).to(dt), + mask=(c < C) & (t < H), + ) + + +@triton.jit +def _bias_gelu_kernel(x_ptr, b_ptr, N: tl.constexpr, BN: tl.constexpr): + row = tl.program_id(0).to(tl.int64) + offs = tl.program_id(1) * BN + tl.arange(0, BN) + m = offs < N + x = tl.load(x_ptr + row * N + offs, mask=m).to(tl.float32) + tl.load(b_ptr + offs, mask=m).to(tl.float32) + dt = x_ptr.dtype.element_ty + x = x.to(dt).to(tl.float32) + y = 0.5 * x * (1.0 + tl.erf(x * 0.7071067811865476)) + tl.store(x_ptr + row * N + offs, y.to(dt), mask=m) + + +@triton.jit +def _col2im_kernel( + z_ptr, # [B * T, 2 * R, C] transposed-conv GEMM rows + bias_ptr, # [C] + prev_ptr, # [2, S, R, C] previous call's last GEMM row, second half + slots_ptr, + pos_ptr, + par_ptr, + out_ptr, # [B * T * R, C] + T, + S, + R: tl.constexpr, + C: tl.constexpr, + BC: tl.constexpr, +): + row = tl.program_id(0) # b * T + t + j = tl.program_id(1) # 0 .. R - 1 + c = tl.program_id(2) * BC + tl.arange(0, BC) + cm = c < C + b = row // T + t = row % T + slot = tl.load(slots_ptr + b).to(tl.int64) + live = tl.load(pos_ptr + b) != 0 + par = tl.load(par_ptr + slot).to(tl.int64) + zr = z_ptr + row.to(tl.int64) * 2 * R * C + cur = tl.load(zr + j * C + c, mask=cm).to(tl.float32) + prev_in = tl.load(zr - 2 * R * C + (R + j) * C + c, mask=cm & (t > 0), other=0.0).to(tl.float32) + prev_st = tl.load(prev_ptr + ((par * S + slot) * R + j) * C + c, mask=cm & (t == 0) & live, other=0.0).to( + tl.float32 + ) + y = cur + tl.where(t > 0, prev_in, prev_st) + tl.load(bias_ptr + c, mask=cm).to(tl.float32) + dt = out_ptr.dtype.element_ty + tl.store(out_ptr + (row.to(tl.int64) * R + j) * C + c, y.to(dt), mask=cm) + last = tl.load(zr + (R + j) * C + c, mask=cm & (t == T - 1)) + tl.store(prev_ptr + (((1 - par) * S + slot) * R + j) * C + c, last, mask=cm & (t == T - 1)) + + +@triton.jit +def _conv_out_kernel( + ext_ptr, # [B, H + T, C] transformed rows: history then this call's + w_ptr, # [K, C] + wb, # conv bias (float) + out_ptr, # [B, T] float32 waveform + T, + C: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + CP: tl.constexpr, +): + b = tl.program_id(0) + t = tl.program_id(1) * BT + tl.arange(0, BT) + c = tl.arange(0, CP) + cm = c < C + tm = t < T + base = ext_ptr + b.to(tl.int64) * (T + K - 1) * C + acc = tl.zeros([BT], dtype=tl.float32) + for j in tl.static_range(K): + v = tl.load(base + (t + j)[:, None] * C + c[None, :], mask=tm[:, None] & cm[None, :], other=0.0).to(tl.float32) + acc += tl.sum(v * tl.load(w_ptr + j * C + c, mask=cm, other=0.0).to(tl.float32)[None, :], axis=1) + dt = ext_ptr.dtype.element_ty + y = acc.to(dt).to(tl.float32) + wb + y = tl.minimum(tl.maximum(y.to(dt).to(tl.float32), -1.0), 1.0) + tl.store(out_ptr + b.to(tl.int64) * T + t, y, mask=tm) + + +@triton.jit +def _flip_kernel(par_ptr, slots_ptr): + slot = tl.load(slots_ptr + tl.program_id(0)).to(tl.int64) + tl.store(par_ptr + slot, 1 - tl.load(par_ptr + slot)) + + +# --------------------------------------------------------------------------- module + + +@dataclass +class _ConvSpec: + weight: torch.Tensor # [K * C_in, C_out] tap-major + c_in: int + k: int + d: int + hist: torch.Tensor | None = None # [2, S, H, C_in] + + +def _conv_taps(w: torch.Tensor) -> torch.Tensor: + """``[C_out, C_in, K]`` -> ``[K * C_in, C_out]``.""" + c_out, c_in, k = w.shape + return w.permute(2, 1, 0).reshape(k * c_in, c_out) + + +def _snake_params(module) -> tuple[torch.Tensor, torch.Tensor]: + # SnakeBeta: x + 1 / (exp(beta) + 1e-9) * sin(exp(alpha) * x) ** 2 + alpha = module.alpha.detach().float() + beta = module.beta.detach().float() + return torch.exp(alpha).contiguous(), (1.0 / (torch.exp(beta) + 1e-9)).contiguous() + + +class StreamingCodecDecoder(nn.Module): + """Frame-exact stateful decoder built from a loaded ``Qwen3TTSTokenizerV2Decoder``.""" + + RING = 128 # KV ring positions; a call may carry up to RING - window + 1 frames + + def __init__(self, decoder: nn.Module, num_slots: int, dtype: torch.dtype | None = None) -> None: + super().__init__() + cfg = decoder.config + dev = next(decoder.parameters()).device + dtype = dtype or next(decoder.parameters()).dtype + self.dtype, self.device = dtype, dev + self.num_slots = num_slots + 1 # last slot: scratch for padding rows + self.scratch_slot = num_slots + S = self.num_slots + self.nq = int(cfg.num_quantizers) + self.cb = int(cfg.codebook_size) + self.half = int(cfg.codebook_dim) // 2 + self.hidden = int(cfg.hidden_size) + self.heads = int(cfg.num_attention_heads) + self.head_dim = int(cfg.head_dim) + self.window = int(cfg.sliding_window) + self.theta = float(cfg.rope_theta) + self.eps = float(cfg.rms_norm_eps) + self.max_frames = self.RING - self.window + 1 + self.spf = int(decoder.total_upsample) + if self.heads != int(cfg.num_key_value_heads): + raise ValueError("streaming codec decoder assumes no grouped KV") + + f32 = lambda t: t.detach().to(device=dev, dtype=torch.float32) # noqa: E731 + cast = lambda t: t.to(dtype).contiguous() # noqa: E731 + + def hist(h: int, c: int) -> torch.Tensor: + return torch.zeros(2, S, h, c, device=dev, dtype=dtype) + + # RVQ books, pre-divided by usage: [NQ * CB, HALF] + q = decoder.quantizer + books = [] + for vq in [q.rvq_first.vq.layers[0]] + list(q.rvq_rest.vq.layers): + cbk = vq._codebook + books.append(f32(cbk.embedding_sum) / f32(cbk.cluster_usage).clamp(min=cbk.epsilon)[:, None]) + self.books = cast(torch.cat(books)) + w_rvq = torch.cat( + [f32(q.rvq_first.output_proj.weight)[:, :, 0], f32(q.rvq_rest.output_proj.weight)[:, :, 0]], 1 + ) + # pre_conv (k=3) and input_proj folded onto the embedding concat + pre_w, pre_b = f32(decoder.pre_conv.conv.weight), f32(decoder.pre_conv.conv.bias) + pt = decoder.pre_transformer + in_w, in_b = f32(pt.input_proj.weight), f32(pt.input_proj.bias) + taps = torch.stack([in_w @ pre_w[:, :, j] @ w_rvq for j in range(pre_w.shape[2])], dim=2) # [hid, 2H, 3] + self.pre = _ConvSpec(cast(_conv_taps(taps)), 2 * self.half, 3, 1, hist(2, 2 * self.half)) + self.pre_b = cast(in_w @ pre_b + in_b) + + # transformer + self.layers = [] + qd = self.heads * self.head_dim + for layer in pt.layers: + a = layer.self_attn + qkv = torch.cat([f32(a.q_proj.weight), f32(a.k_proj.weight), f32(a.v_proj.weight)], 0) + o = f32(a.o_proj.weight) * f32(layer.self_attn_layer_scale.scale)[:, None] + gu = torch.cat([f32(layer.mlp.gate_proj.weight), f32(layer.mlp.up_proj.weight)], 0) + down = f32(layer.mlp.down_proj.weight) * f32(layer.mlp_layer_scale.scale)[:, None] + self.layers.append( + dict( + ln1=cast(f32(layer.input_layernorm.weight)), + ln2=cast(f32(layer.post_attention_layernorm.weight)), + qkv=cast(qkv.t()), + o=cast(o.t()), + gu=cast(gu.t()), + down=cast(down.t()), + k_ring=torch.zeros(S, self.RING, qd, device=dev, dtype=dtype), + v_ring=torch.zeros(S, self.RING, qd, device=dev, dtype=dtype), + ) + ) + self.norm_w = cast(f32(pt.norm.weight)) + self.out_w = cast(f32(pt.output_proj.weight).t()) + self.out_b = cast(f32(pt.output_proj.bias)) + + # upsample: transposed conv (kernel = stride, no overlap) + ConvNeXt + self.ups = [] + latent = int(cfg.latent_dim) + pending = torch.zeros(latent, device=dev, dtype=torch.float32) + for trans, block in decoder.upsample: + tw = f32(trans.conv.weight) # [C_in, C_out, r] + r = tw.shape[2] + w = tw.permute(0, 2, 1).reshape(latent, r * latent) # [C_in, r * C_out] + bias = (pending @ w).view(r, latent) + f32(trans.conv.bias)[None, :] + dw = block.dwconv.conv + gamma = f32(block.gamma) + self.ups.append( + dict( + r=r, + w=cast(w), + b=cast(bias.reshape(-1)), + dw_w=cast(f32(dw.weight)[:, 0, :].t()), + dw_b=cast(f32(dw.bias)), + ln_w=cast(f32(block.norm.weight)), + ln_b=cast(f32(block.norm.bias)), + ln_eps=float(block.norm.eps), + pw1=cast(f32(block.pwconv1.weight).t()), + pw1_b=cast(f32(block.pwconv1.bias)), + pw2=cast((f32(block.pwconv2.weight) * gamma[:, None]).t()), + hist=hist(_KERNEL - 1, latent), + ) + ) + pending = f32(block.pwconv2.bias) * gamma + + # conv_in (pending bias of the last ConvNeXt added on load) + dec = decoder.decoder + cin = dec[0].conv + self.conv_in = _ConvSpec(cast(_conv_taps(f32(cin.weight))), latent, _KERNEL, 1, hist(_KERNEL - 1, latent)) + self.conv_in_pb = cast(pending) + pending = f32(cin.bias) + + # decoder blocks + self.blocks = [] + for i, blk in enumerate(dec[1:-2]): + snake, trans, *units = blk.block + c_in = trans.conv.in_channels + c_out = trans.conv.out_channels + rate = trans.conv.stride[0] + tw = f32(trans.conv.weight) # [C_in, C_out, 2 * rate] + entry = dict( + pb=cast(pending), + snake=_snake_params(snake), + rate=rate, + c_in=c_in, + c_out=c_out, + w=cast(tw.permute(0, 2, 1).reshape(c_in, 2 * rate * c_out)), + b=cast(f32(trans.conv.bias)), + prev=torch.zeros(2, S, rate, c_out, device=dev, dtype=dtype), + units=[], + ) + pending = torch.zeros(c_out, device=dev, dtype=torch.float32) + for u in units: + c1, c2 = u.conv1.conv, u.conv2.conv + entry["units"].append( + dict( + pb=cast(pending), + act1=_snake_params(u.act1), + conv1=_ConvSpec( + cast(_conv_taps(f32(c1.weight))), + c_out, + _KERNEL, + c1.dilation[0], + hist((_KERNEL - 1) * c1.dilation[0], c_out), + ), + b1=cast(f32(c1.bias)), + conv1_w4=( + cast(f32(c1.weight)).unsqueeze(2).contiguous(memory_format=torch.channels_last) + if c1.dilation[0] == 1 + else None + ), + act2=_snake_params(u.act2), + w2=cast(f32(c2.weight)[:, :, 0].t()), + ) + ) + pending = pending + f32(c2.bias) + self.blocks.append(entry) + self.out_pb = cast(pending) + self.out_snake = _snake_params(dec[-2]) + co = dec[-1].conv + self.out_c = co.in_channels + self.conv_out_w = cast(f32(co.weight)[0].t()) # [K, C] + self.conv_out_b = float(f32(co.bias)[0]) + self.conv_out = _ConvSpec(self.conv_out_w, self.out_c, _KERNEL, 1, hist(_KERNEL - 1, self.out_c)) + self.parity = torch.zeros(S, device=dev, dtype=torch.int32) + self._dummy = torch.zeros(1, device=dev, dtype=dtype) + + def _slot_state(self, slot: int) -> list[torch.Tensor]: + state = [self.parity[slot : slot + 1]] + for layer in self.layers: + state.extend((layer["k_ring"][slot], layer["v_ring"][slot])) + for spec in (self.pre, self.conv_in, self.conv_out): + assert spec.hist is not None + state.append(spec.hist[:, slot]) + for up in self.ups: + state.append(up["hist"][:, slot]) + for block in self.blocks: + state.append(block["prev"][:, slot]) + for unit in block["units"]: + state.append(unit["conv1"].hist[:, slot]) + return state + + def save_slot(self, slot: int) -> list[torch.Tensor]: + """Offload a preempted request before its runner slot is recycled.""" + return [tensor.to(device="cpu", copy=True) for tensor in self._slot_state(slot)] + + def restore_slot(self, slot: int, state: list[torch.Tensor]) -> None: + for target, saved in zip(self._slot_state(slot), state, strict=True): + target.copy_(saved) + + # ------------------------------------------------------------------ helpers + def _im2col(self, x, spec: _ConvSpec, T, slots, pos, bias=None, snake=None, ext=False): + k = spec.k + d = spec.d + c = x.shape[1] + h = (k - 1) * d + B = slots.shape[0] + if ext: + col = torch.empty(B * (h + T), c, device=x.device, dtype=self.dtype) + else: + col = torch.empty(B * T, k * c, device=x.device, dtype=self.dtype) + bu = 32 + bc = min(128, triton.next_power_of_2(c)) + grid = (B, triton.cdiv(h + T, bu), triton.cdiv(c, bc)) + a, ib = snake if snake is not None else (self._dummy, self._dummy) + _im2col_kernel[grid]( + x, bias if bias is not None else self._dummy, a, ib, + spec.hist if h > 0 else self._dummy, slots, pos, self.parity, col, + T, self.num_slots, C=c, K=k, D=d, H=h, + HAS_BIAS=bias is not None, HAS_SNAKE=snake is not None, BU=bu, BC=bc, EXT=ext, + ) # fmt: skip + return col + + def _conv_cudnn(self, x, u, T, slots, pos): + """Undilated unit conv: one source-row pass, then cuDNN on its channels-last view. + + The im2col path writes a ``K``-times patch matrix; cuDNN reads the + ``[history | new]`` rows directly (same result up to summation order). + """ + spec = u["conv1"] + ext = self._im2col(x, spec, T, slots, pos, bias=u["pb"], snake=u["act1"], ext=True) + b = slots.shape[0] + c = x.shape[1] + x4 = ext.view(b, -1, c).permute(0, 2, 1).unsqueeze(2) # [B, C, 1, H + T], channels-last memory + y = torch.nn.functional.conv2d(x4, u["conv1_w4"]) # [B, N, 1, T], channels-last memory + return y.permute(0, 2, 3, 1).reshape(b * T, -1) + + def _snake(self, x, snake, bias): + n, c = x.shape + out = torch.empty_like(x) + bu, bc = 32, min(128, triton.next_power_of_2(c)) + one = torch.zeros(1, device=x.device, dtype=torch.int32) + _im2col_kernel[(1, triton.cdiv(n, bu), triton.cdiv(c, bc))]( + x, bias, snake[0], snake[1], self._dummy, one, one, self.parity, out, + n, self.num_slots, C=c, K=1, D=1, H=0, HAS_BIAS=True, HAS_SNAKE=True, BU=bu, BC=bc, + ) # fmt: skip + return out + + # ------------------------------------------------------------------ forward + @torch.no_grad() + def forward(self, codes: torch.Tensor, slots: torch.Tensor, pos: torch.Tensor) -> torch.Tensor: + """``codes`` [B, T, NQ] int, ``slots``/``pos`` [B] int32 -> waveform [B, T * spf] float32. + + ``pos[b]`` is the index of the row's first frame in its stream; 0 starts + a new stream in ``slots[b]``. Frames beyond ``max_frames`` per call must + be split by the caller. + """ + B, T, _ = codes.shape + if T > self.max_frames: + raise ValueError(f"at most {self.max_frames} frames per call, got {T}") + dev, dt = codes.device, self.dtype + codes = codes.reshape(B * T, self.nq).to(torch.int32).contiguous() + emb = torch.empty(B * T, 2 * self.half, device=dev, dtype=dt) + _rvq_gather_kernel[(B * T,)](codes, self.books, emb, NQ=self.nq, CB=self.cb, HALF=self.half) + x = torch.addmm(self.pre_b, self._im2col(emb, self.pre, T, slots, pos), self.pre.weight) + + hid, qd = self.hidden, self.heads * self.head_dim + rows = B * T + h = torch.empty(rows, hid, device=dev, dtype=dt) + q = torch.empty(rows, qd, device=dev, dtype=dt) + att = torch.empty(rows, qd, device=dev, dtype=dt) + a = torch.empty(rows, self.layers[0]["down"].shape[0], device=dev, dtype=dt) + inter = a.shape[1] + scale = self.head_dim**-0.5 + for L in self.layers: + _rms_norm_kernel[(rows,)](x, L["ln1"], h, self.eps, N=hid) + qkv = torch.mm(h, L["qkv"]) + _rope_kv_write_kernel[(rows, self.heads)]( + qkv, q, L["k_ring"], L["v_ring"], slots, pos, T, self.theta, + NH=self.heads, HD=self.head_dim, RING=self.RING, + ) # fmt: skip + _attend_kernel[(rows, self.heads)]( + q, L["k_ring"], L["v_ring"], att, slots, pos, T, scale, + NH=self.heads, HD=self.head_dim, RING=self.RING, WINDOW=self.window, + BK=triton.next_power_of_2(self.window), + ) # fmt: skip + x.addmm_(att, L["o"]) + _rms_norm_kernel[(rows,)](x, L["ln2"], h, self.eps, N=hid) + gu = torch.mm(h, L["gu"]) + _silu_mul_kernel[(rows, triton.cdiv(inter, 1024))](gu, a, N=inter, BN=1024) + x.addmm_(a, L["down"]) + _rms_norm_kernel[(rows,)](x, self.norm_w, h, self.eps, N=hid) + y = torch.addmm(self.out_b, h, self.out_w) # [rows, latent] + + t_rows = T + for U in self.ups: + r = U["r"] + z = torch.addmm(U["b"], y, U["w"]).view(rows * r, -1) + rows, t_rows = rows * r, t_rows * r + c = z.shape[1] + n = torch.empty_like(z) + _dwconv_ln_kernel[(B, max(t_rows, _KERNEL - 1))]( + z, U["dw_w"], U["dw_b"], U["ln_w"], U["ln_b"], U["hist"], slots, pos, self.parity, n, + t_rows, self.num_slots, U["ln_eps"], C=c, K=_KERNEL, + ) # fmt: skip + p1 = torch.mm(n, U["pw1"]) + _bias_gelu_kernel[(rows, triton.cdiv(p1.shape[1], 1024))](p1, U["pw1_b"], N=p1.shape[1], BN=1024) + z.addmm_(p1, U["pw2"]) + y = z + + x = torch.mm(self._im2col(y, self.conv_in, t_rows, slots, pos, bias=self.conv_in_pb), self.conv_in.weight) + for blk in self.blocks: + s = self._snake(x, blk["snake"], blk["pb"]) + zt = torch.mm(s, blk["w"]) # [rows, 2 * rate * c_out] + r, c_out = blk["rate"], blk["c_out"] + x = torch.empty(rows * r, c_out, device=dev, dtype=dt) + bc = min(128, triton.next_power_of_2(c_out)) + _col2im_kernel[(rows, r, triton.cdiv(c_out, bc))]( + zt, blk["b"], blk["prev"], slots, pos, self.parity, x, t_rows, self.num_slots, + R=r, C=c_out, BC=bc, + ) # fmt: skip + rows, t_rows = rows * r, t_rows * r + for u in blk["units"]: + if u.get("conv1_w4") is not None: + c1 = self._conv_cudnn(x, u, t_rows, slots, pos) + else: + col = self._im2col(x, u["conv1"], t_rows, slots, pos, bias=u["pb"], snake=u["act1"]) + c1 = torch.mm(col, u["conv1"].weight) + s2 = self._snake(c1, u["act2"], u["b1"]) + x.addmm_(s2, u["w2"]) + + ext = self._im2col(x, self.conv_out, t_rows, slots, pos, bias=self.out_pb, snake=self.out_snake, ext=True) + wav = torch.empty(B, t_rows, device=dev, dtype=torch.float32) + _conv_out_kernel[(B, triton.cdiv(t_rows, 256))]( + ext, self.conv_out_w, self.conv_out_b, wav, t_rows, + C=self.out_c, K=_KERNEL, BT=256, CP=triton.next_power_of_2(self.out_c), + ) # fmt: skip + _flip_kernel[(B,)](self.parity, slots) + return wav + + +class StreamingDecodeGraphs: + """CUDA graphs of ``StreamingCodecDecoder`` for a fixed frame count per batch bucket.""" + + def __init__(self, sd: StreamingCodecDecoder, batch_sizes: list[int], frames: int = 1) -> None: + self.sd, self.frames = sd, frames + self.graphs: dict[int, tuple] = {} + dev = sd.device + pool = torch.cuda.graph_pool_handle() + for bsz in sorted(set(batch_sizes)): + codes = torch.zeros(bsz, frames, sd.nq, dtype=torch.int32, device=dev) + slots = torch.full((bsz,), sd.scratch_slot, dtype=torch.int32, device=dev) + pos = torch.zeros(bsz, dtype=torch.int32, device=dev) + sd(codes, slots, pos) + current_omni_platform.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=pool): + out = sd(codes, slots, pos) + self.graphs[bsz] = (graph, codes, slots, pos, out) + self.sizes = sorted(self.graphs) + + def __call__(self, codes: torch.Tensor, slots: torch.Tensor, pos: torch.Tensor) -> torch.Tensor: + """``codes`` [n, frames, NQ], ``slots``/``pos`` [n] on device -> waveform view [n, frames * spf].""" + n = int(codes.shape[0]) + size = next((b for b in self.sizes if b >= n), None) + if size is None: + return self.sd(codes, slots.to(torch.int32), pos.to(torch.int32)) + graph, s_codes, s_slots, s_pos, out = self.graphs[size] + s_codes[:n].copy_(codes) + s_slots[:n].copy_(slots) + s_pos[:n].copy_(pos) + if n < size: + s_slots[n:].fill_(self.sd.scratch_slot) + s_pos[n:].zero_() + graph.replay() + return out[:n] diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py index 482514c04..33c0da7bd 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py @@ -390,16 +390,17 @@ def talker2code2wav_async_chunk( if ref_context_size > 0: if emitted_chunks <= 0: ref_context_request_id = request_id - ref_frames = ref_context.tolist() + ref_frames = list(ref_context.to(device="cpu").unbind(0)) window_frames = ref_frames + window_frames ref_context_included = True left_context_size = ref_context_size - num_quantizers = len(window_frames[0]) - num_frames = len(window_frames) - code_predictor_codes = torch.tensor( - [window_frames[f][q] for q in range(num_quantizers) for f in range(num_frames)], - dtype=torch.long, + # Keep the codebook-major CPU wire format without extracting every + # tensor element as a Python scalar. Legacy buffered frames can be lists. + code_predictor_codes = ( + torch.stack([torch.as_tensor(frame, dtype=torch.long, device="cpu") for frame in window_frames]) + .transpose(0, 1) + .reshape(-1) ) meta = MetaStruct( diff --git a/vllm_omni/worker_v2/model_states/eager_mtp.py b/vllm_omni/worker_v2/model_states/eager_mtp.py index ba347206a..03fb8c0f1 100644 --- a/vllm_omni/worker_v2/model_states/eager_mtp.py +++ b/vllm_omni/worker_v2/model_states/eager_mtp.py @@ -5,13 +5,17 @@ from __future__ import annotations from collections.abc import Callable +from dataclasses import dataclass from typing import TYPE_CHECKING, Any +import numpy as np import torch from vllm.logger import init_logger from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.states import RequestState from vllm_omni.utils.device_copy import index_to_device +from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput if TYPE_CHECKING: from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState @@ -24,12 +28,175 @@ def _has_ref_codes(buffer: dict[str, Any]) -> bool: return isinstance(ref, torch.Tensor) and ref.numel() > 0 +def _snapshot_runtime(value: Any) -> Any: + if isinstance(value, torch.Tensor): + return value.to("cpu", copy=True) + if isinstance(value, dict): + return {key: _snapshot_runtime(item) for key, item in value.items()} + if isinstance(value, list): + return [_snapshot_runtime(item) for item in value] + if isinstance(value, tuple): + return tuple(_snapshot_runtime(item) for item in value) + return value + + +@dataclass +class TalkerInputs: + embeds: torch.Tensor + length: int = 0 + + +@dataclass +class SuspendedAudio: + position: int + decoder_state: list[torch.Tensor] + generator: torch.Generator | None + runtime: dict[str, Any] + frame_embed: torch.Tensor + + class EagerMTPState: """Used only by models explicitly declaring ``mtp_eager_frames``.""" def __init__(self, owner: OmniModelState) -> None: self.owner = owner self._first_audio_valid: torch.Tensor | None = None + self._side_stream: torch.cuda.Stream | None = None + self._fast_ok: bool | None = None + self._audio_buffer: StreamingAudioBuffer | None = None + self._frame_requests: set[str] = set() + self._suspended_audio: dict[str, SuspendedAudio] = {} + self._restore_audio: dict[str, list[torch.Tensor]] = {} + self._talker_inputs: dict[str, TalkerInputs] = {} + + def suspend_audio(self, req_id: str, req_idx: int) -> None: + decoder = getattr(self.owner.model, "stream_decoder", None) + position = self.owner._stream_pos.get(req_id) + if decoder is None or position is None: + return + # Slot reuse is rare and must wait for the last side-stream write. + saved = self._restore_audio.pop(req_id, None) + if saved is None: + self._decode_stream.synchronize() + saved = decoder.save_slot(req_idx) + history = self._talker_inputs[req_id] + history.embeds = history.embeds[: history.length].to("cpu", copy=True) + runtime = self.owner.intermediate_buffer.buffers[req_idx] + assert self.owner._eager_embeds is not None + self._suspended_audio[req_id] = SuspendedAudio( + position, + saved, + self.owner._mtp_generators.get(req_id), + _snapshot_runtime(runtime), + self.owner._eager_embeds[req_idx].to("cpu", copy=True), + ) + + def resume_audio(self, req_id: str, req_idx: int) -> None: + saved = self._suspended_audio.pop(req_id, None) + if saved is not None: + self.owner._stream_pos[req_id] = saved.position + self.owner._mtp_generators[req_id] = saved.generator + self._restore_audio[req_id] = saved.decoder_state + self.owner.intermediate_buffer.buffers[req_idx] = saved.runtime + assert self.owner._eager_embeds is not None + self.owner._eager_embeds[req_idx].copy_(saved.frame_embed) + self.owner._eager_ready[req_idx] = req_id + + def finish_audio(self, req_ids: set[str]) -> None: + for req_id in req_ids: + self._suspended_audio.pop(req_id, None) + self._restore_audio.pop(req_id, None) + self.owner._stream_pos.pop(req_id, None) + self._talker_inputs.pop(req_id, None) + if self._audio_buffer is not None: + self._audio_buffer.finish(req_ids) + + def replay_inputs(self, req_id: str, req_idx: int, offset: int, ids: torch.Tensor, embeds: torch.Tensor) -> bool: + """Rebuild freed Talker KV from the exact conditioned inputs, without rerunning MTP.""" + if req_id not in self._restore_audio: + return False + history = self._talker_inputs[req_id] + count = min(len(ids), max(0, history.length - offset)) + if count: + embeds[:count].copy_(history.embeds[offset : offset + count]) + if count < len(ids): + # Recompute ends with the last accepted CB0, whose full eager frame + # was saved but has not yet been consumed as a Talker input. + if len(ids) - count != 1 or offset + count != history.length: + raise RuntimeError(f"Talker replay exceeds saved inputs for {req_id}") + info = self.owner.intermediate_buffer.buffers[req_idx] + _, _, updates = self.owner.model.preprocess(ids[count:], embeds[count:], **info, _omni_is_prefill=False) + mtp_inputs = updates.pop("mtp_inputs", None) + if mtp_inputs is None: + raise RuntimeError(f"Talker replay has no text conditioning for {req_id}") + assert self.owner._eager_embeds is not None + embeds[count:].copy_(self.owner._eager_embeds[req_idx] + mtp_inputs[1].reshape(1, -1)) + self.owner.intermediate_buffer.update(req_idx, updates, self.owner.model.gpu_resident_buffer_keys) + return True + + def record_inputs(self, input_batch: InputBatch, embeds: torch.Tensor) -> None: + if getattr(self.owner.model, "stream_decoder", None) is None: + return + history_slices: list[torch.Tensor] = [] + input_slices: list[torch.Tensor] = [] + for i, req_id in enumerate(input_batch.req_ids): + start, end = input_batch.query_start_loc_np[i : i + 2] + offset = int(input_batch.num_computed_tokens_np[i]) + length = offset + int(end - start) + history = self._talker_inputs.get(req_id) + if history is None or history.embeds.device != embeds.device: + storage = torch.empty( + self.owner.vllm_config.model_config.max_model_len, + embeds.shape[-1], + device=embeds.device, + dtype=embeds.dtype, + ) + if history is not None: + storage[: history.length].copy_(history.embeds) + history = TalkerInputs(storage, history.length if history else 0) + self._talker_inputs[req_id] = history + history_slices.append(history.embeds[offset:length]) + input_slices.append(embeds[start:end]) + history.length = max(history.length, length) + if history_slices: + # Keep exact inputs for replay without launching one copy per row. + torch._foreach_copy_(history_slices, input_slices) + + def prepare_audio_output( + self, input_batch: InputBatch, req_states: RequestState, outputs: dict[str, Any] + ) -> StreamingAudioOutput | None: + if self._audio_buffer is None: + return None + wav = outputs.pop("model_outputs", None) + if not isinstance(wav, torch.Tensor): + return None + n = input_batch.num_reqs + total_after = input_batch.num_computed_tokens_np[:n] + np.asarray(input_batch.num_scheduled_tokens[:n]) + 1 + limits = np.minimum( + req_states.max_seq_len[input_batch.idx_mapping_np[:n]], self.owner.vllm_config.model_config.max_model_len + ) + requests = [ + self._audio_buffer.requests.get(rid) if rid in self._frame_requests else None for rid in input_batch.req_ids + ] + self._frame_requests.clear() + event = self.owner._stream_decode_event + self.owner._stream_decode_event = None + return StreamingAudioOutput( + wav, + outputs["meta"]["codec_frame_valid"], + outputs["sr"][0], + event, + input_batch.query_start_loc_np[: n + 1].copy(), + requests, + total_after >= limits, + ) + + @property + def _decode_stream(self) -> torch.cuda.Stream: + """Side stream for Talker stream decode, so the codec overlaps the next step.""" + if self._side_stream is None: + self._side_stream = torch.cuda.Stream() + return self._side_stream def set_first_audio_sink(self, sink: Any) -> None: from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender @@ -106,6 +273,14 @@ class EagerMTPState: bsz = len(entries) device = text_hidden.device + if self._fast_path_ok(codes_out, text_hidden): + self._run_eager_mtp_fused( + input_batch, text_hidden, sampled_token_ids, multimodal_outputs, entries, input_ids, + codes_out, valid_out, mtp_batch_descriptor_dispatcher, + ) # fmt: skip + return + if getattr(self.owner.model, "stream_decoder", None) is not None: + raise RuntimeError("Talker stream decode requires the fused eager-MTP path") batch_rows = index_to_device([i for i, _req_idx, _req_id, _prefill in entries], device) req_indices = [req_idx for _i, req_idx, _req_id, _prefill in entries] last_tokens = input_batch.query_start_loc.index_select(0, batch_rows + 1).long() - 1 @@ -162,6 +337,168 @@ class EagerMTPState: scheduled *= self._first_audio_valid.index_select(0, rows) first_audio.index_copy_(0, last_tokens, scheduled.to(first_audio.dtype)) + def _embedding_weight(self) -> torch.Tensor | None: + inner = getattr(self.owner.model, "model", None) + embed = getattr(inner, "embed_tokens", None) + weight = getattr(embed, "weight", None) + method = getattr(embed, "quant_method", None) + unquantized = method is None or type(method).__name__ == "UnquantizedEmbeddingMethod" + tp = getattr(embed, "tp_size", 1) + if isinstance(weight, torch.Tensor) and weight.ndim == 2 and unquantized and tp == 1: + return weight + return None + + def _fast_path_ok(self, codes_out: torch.Tensor, text_hidden: torch.Tensor) -> bool: + if self._fast_ok is None: + self._fast_ok = ( + text_hidden.is_cuda + and not self.owner.vllm_config.cache_config.enable_prefix_caching + and self._embedding_weight() is not None + and hasattr(self.owner.model, "_codebook_vocab_size") + ) + return self._fast_ok and codes_out.ndim == 2 and text_hidden.is_contiguous() + + def _run_eager_mtp_fused( + self, + input_batch: InputBatch, + text_hidden: torch.Tensor, + sampled_token_ids: torch.Tensor, + multimodal_outputs: dict[str, Any], + entries: list[tuple[int, int, str, bool]], + input_ids: torch.Tensor, + codes_out: torch.Tensor, + valid_out: torch.Tensor, + mtp_batch_descriptor_dispatcher: Callable[[int], Any] | None, + ) -> None: + from vllm_omni.worker_v2.model_states.eager_mtp_kernels import eager_post, eager_pre + + owner = self.owner + bsz = len(entries) + device = text_hidden.device + first_audio = multimodal_outputs.get("meta", {}).get("first_audio") + fa_requests = owner._first_audio_requests + model = owner.model + stream = getattr(model, "stream_decoder", None) + stream_out = multimodal_outputs.get("model_outputs") if stream is not None else None + pos_list: list[int] = [] + first_rows: list[int] = [] + primes: list[tuple[int, torch.Tensor]] = [] + if isinstance(stream_out, torch.Tensor): + assert stream is not None + # Frame index of each row's frame in its request's stream; 0 starts + # a fresh decoder state in the request's slot. A voice-clone + # request's stream starts after its reference codes, which prime + # the decoder state first (as Code2Wav's first-chunk context does). + if self._audio_buffer is None: + self._audio_buffer = StreamingAudioBuffer(int(stream.spf), int(model.stream_chunk_frames)) + self._frame_requests = {req_id for _i, _idx, req_id, _p in entries} + positions = owner._stream_pos + for row, (_i, req_idx, req_id, _p) in enumerate(entries): + self._audio_buffer.add(req_id) + p = positions.get(req_id) + if p is None: + first_rows.append(row) + ref = model.get_stream_ref_context(owner.intermediate_buffer.buffers[req_idx]) + p = 0 + if ref is not None: + primes.append((req_idx, ref)) + p = int(ref.shape[0]) + pos_list.append(p) + positions[req_id] = p + 1 + meta_list = ( + [i for i, _r, _q, _p in entries] + + [r for _i, r, _q, _p in entries] + + [int(p) for _i, _r, _q, p in entries] + + [int(q in fa_requests) for _i, _r, q, _p in entries] + + pos_list + ) + meta = index_to_device(meta_list, device) + rows = meta[bsz : 2 * bsz] + sampled = sampled_token_ids.reshape(input_batch.num_reqs, -1) + assert owner._mtp_input_ids is not None and owner._mtp_input_embeds is not None + assert owner._mtp_hidden is not None and owner._mtp_text_step is not None + batch_ids = owner._mtp_input_ids[:bsz] + batch_emb = owner._mtp_input_embeds[:bsz] + batch_hidden = owner._mtp_hidden[:bsz] + batch_step = owner._mtp_text_step[:bsz] + last_tokens, layer0 = eager_pre( + meta, bsz, input_batch.query_start_loc, sampled, text_hidden, self._embedding_weight(), + batch_ids, batch_emb, batch_hidden, batch_step, + ) # fmt: skip + req_indices = meta_list[bsz : 2 * bsz] + frame_embeds, codes = owner._mtp_forward( + req_indices, batch_ids, batch_emb, batch_hidden, batch_step, mtp_batch_descriptor_dispatcher + ) + assert codes is not None + has_prefill = any(p for _i, _r, _q, p in entries) + fa_in_kernel = isinstance(first_audio, torch.Tensor) and not has_prefill + valid = eager_post( + meta, bsz, last_tokens, layer0, frame_embeds.reshape(frame_embeds.shape[0], -1).contiguous(), + owner._eager_embeds, codes.contiguous(), codes_out, input_ids, valid_out, + first_audio if fa_in_kernel else None, self._first_audio_valid if fa_in_kernel else None, + owner.model._codebook_vocab_size, + ) # fmt: skip + for _i, req_idx, req_id, _prefill in entries: + owner._eager_ready[req_idx] = req_id + if pos_list: + assert stream is not None + assert isinstance(stream_out, torch.Tensor) + decode = model.stream_graphs if model.stream_graphs is not None else stream + frame_codes = codes[:bsz].reshape(bsz, 1, -1).to(torch.int32) + # The next Talker step does not read the PCM, so the codec runs + # beside it; only the output copies wait for it. + side = self._decode_stream + side.wait_stream(torch.cuda.current_stream()) + for t in (frame_codes, meta, last_tokens, valid, stream_out): + t.record_stream(side) + with torch.cuda.stream(side): + for _i, req_idx, req_id, _p in entries: + saved = self._restore_audio.pop(req_id, None) + if saved is not None: + stream.restore_slot(req_idx, saved) + model.prime_stream_decoder(primes) + pcm = decode(frame_codes, rows, meta[4 * bsz : 5 * bsz]) + stream_out.index_copy_(0, last_tokens, pcm.reshape(bsz, -1).to(stream_out.dtype)) + self._send_stream_first_frames(entries, pcm, valid, first_rows) + done = torch.cuda.Event() + done.record(side) + owner._stream_decode_event = done + if has_prefill: + self._publish_first_audio(input_batch, entries, codes[:bsz], valid) + if isinstance(first_audio, torch.Tensor): + # Requests accepted by this step's publish count as scheduled. + scheduled = index_to_device( + [int(q in owner._first_audio_requests) for _i, _r, q, _p in entries], device + ) + if self._first_audio_valid is None: + scheduled.zero_() + else: + scheduled *= self._first_audio_valid.index_select(0, rows) + first_audio.index_copy_(0, last_tokens, scheduled.to(first_audio.dtype)) + + def _send_stream_first_frames( + self, + entries: list[tuple[int, int, str, bool]], + pcm: torch.Tensor, + valid: torch.Tensor, + first_rows: list[int], + ) -> None: + sender = self.owner._first_audio_sender + if sender is None or not first_rows: + return + rows = index_to_device(first_rows, pcm.device) + request_ids = [entries[row][2] for row in first_rows] + sample_rate = torch.tensor(int(self.owner.model.stream_sample_rate), dtype=torch.int32) + accepted = sender.submit( + request_ids, + pcm.reshape(pcm.shape[0], -1).index_select(0, rows).float(), + sample_rate, + valid=valid.index_select(0, rows), + ) + assert self._audio_buffer is not None + for request_id in accepted: + self._audio_buffer.requests[request_id].accept_first_audio() + def _publish_first_audio( self, input_batch: InputBatch, diff --git a/vllm_omni/worker_v2/model_states/eager_mtp_kernels.py b/vllm_omni/worker_v2/model_states/eager_mtp_kernels.py new file mode 100644 index 000000000..c48f87875 --- /dev/null +++ b/vllm_omni/worker_v2/model_states/eager_mtp_kernels.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# ruff: noqa: N803 - Triton constexpr parameters use kernel-style capitals. +"""Gather/scatter kernels around the eager Talker-MTP call. + +Each engine step, the rows sampled this step need their CB0, codec embedding +and last hidden state gathered into the MTP inputs, and the MTP results +scattered back into the eager-frame embeddings and the retained outputs. As +separate torch ops that is a few dozen small launches per step, and the engine +loop is host-bound, so each costs the whole batch. These two kernels do the +same indexing in one launch each; the row metadata arrives in one H2D copy. +""" + +from __future__ import annotations + +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def _pre_kernel( + meta_ptr, # int64 [4, n]: batch row, request index, prefill flag, first-audio scheduled + n, + qsl_ptr, # query_start_loc + sampled_ptr, # [num_reqs, s_stride] + s_stride, + hidden_ptr, # [T, H] + emb_w_ptr, # [V, H] + ids_ptr, # [n] MTP input ids + emb_ptr, # [n, H] + mtp_hidden_ptr, # [n, H] + step_ptr, # [n, H] + last_ptr, # int64 [n] out + layer0_ptr, # int64 [n] out + H: tl.constexpr, + BH: tl.constexpr, +): + b = tl.program_id(0) + row = tl.load(meta_ptr + b) + last = tl.load(qsl_ptr + row + 1).to(tl.int64) - 1 + l0 = tl.load(sampled_ptr + row * s_stride).to(tl.int64) + tl.store(ids_ptr + b, l0.to(ids_ptr.dtype.element_ty)) + tl.store(last_ptr + b, last) + tl.store(layer0_ptr + b, l0) + for h0 in tl.static_range(0, H, BH): + c = h0 + tl.arange(0, BH) + m = c < H + tl.store(emb_ptr + b * H + c, tl.load(emb_w_ptr + l0 * H + c, mask=m).to(emb_ptr.dtype.element_ty), mask=m) + tl.store(mtp_hidden_ptr + b * H + c, tl.load(hidden_ptr + last * H + c, mask=m), mask=m) + tl.store(step_ptr + b * H + c, tl.zeros([BH], dtype=step_ptr.dtype.element_ty), mask=m) + + +@triton.jit +def _post_kernel( + meta_ptr, + n, + last_ptr, + layer0_ptr, + frame_ptr, # [>=n, H] MTP frame embeddings + eager_ptr, # [R, H] + codes_ptr, # [>=n, Q] + codes_out_ptr, # [T, Q] + input_ids_ptr, # [T] this step's input ids + valid_out_ptr, # [T] + valid_ptr, # [n] bool out + fa_ptr, # [T] first_audio output or dummy + fa_valid_ptr, # [R] bool or dummy + vocab, + H: tl.constexpr, + Q: tl.constexpr, + QP: tl.constexpr, + BH: tl.constexpr, + WRITE_FA: tl.constexpr, + HAS_FA_VALID: tl.constexpr, +): + b = tl.program_id(0) + req = tl.load(meta_ptr + n + b) + prefill = tl.load(meta_ptr + 2 * n + b) != 0 + last = tl.load(last_ptr + b) + l0 = tl.load(layer0_ptr + b) + for h0 in tl.static_range(0, H, BH): + c = h0 + tl.arange(0, BH) + m = c < H + v = tl.load(frame_ptr + b * H + c, mask=m) + tl.store(eager_ptr + req * H + c, v.to(eager_ptr.dtype.element_ty), mask=m) + q = tl.arange(0, QP) + qm = q < Q + cv = tl.load(codes_ptr + b * Q + q, mask=qm) + tl.store(codes_out_ptr + last * Q + q, cv.to(codes_out_ptr.dtype.element_ty), mask=qm) + inp = tl.load(input_ids_ptr + last).to(tl.int64) + in_valid = ((inp >= 0) & (inp < vocab)) | prefill + valid = (l0 >= 0) & (l0 < vocab) & in_valid + tl.store(valid_out_ptr + last, valid.to(valid_out_ptr.dtype.element_ty)) + tl.store(valid_ptr + b, valid) + if WRITE_FA: + scheduled = tl.load(meta_ptr + 3 * n + b) != 0 + if HAS_FA_VALID: + scheduled = scheduled & (tl.load(fa_valid_ptr + req) != 0) + else: + scheduled = scheduled & False + tl.store(fa_ptr + last, scheduled.to(fa_ptr.dtype.element_ty)) + + +def eager_pre(meta, n, qsl, sampled, hidden, emb_w, ids, emb, mtp_hidden, step): + last = torch.empty(n, dtype=torch.long, device=hidden.device) + layer0 = torch.empty(n, dtype=torch.long, device=hidden.device) + H = hidden.shape[-1] + _pre_kernel[(n,)]( + meta, n, qsl, sampled, sampled.stride(0), hidden, emb_w, ids, emb, mtp_hidden, step, last, layer0, + H=H, BH=min(1024, triton.next_power_of_2(H)), + ) # fmt: skip + return last, layer0 + + +def eager_post(meta, n, last, layer0, frames, eager, codes, codes_out, input_ids, valid_out, fa, fa_valid, vocab): + valid = torch.empty(n, dtype=torch.bool, device=last.device) + H = eager.shape[-1] + Q = codes_out.shape[-1] + dummy = valid + _post_kernel[(n,)]( + meta, n, last, layer0, frames, eager, codes, codes_out, input_ids, valid_out, valid, + fa if fa is not None else dummy, fa_valid if fa_valid is not None else dummy, int(vocab), + H=H, Q=Q, QP=triton.next_power_of_2(Q), BH=min(1024, triton.next_power_of_2(H)), + WRITE_FA=fa is not None, HAS_FA_VALID=fa_valid is not None, + ) # fmt: skip + return valid diff --git a/vllm_omni/worker_v2/model_states/omni_model_state.py b/vllm_omni/worker_v2/model_states/omni_model_state.py index f20a0f2db..f10149f46 100644 --- a/vllm_omni/worker_v2/model_states/omni_model_state.py +++ b/vllm_omni/worker_v2/model_states/omni_model_state.py @@ -41,6 +41,7 @@ from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState from vllm_omni.worker_v2.model_states.intermediate_buffer import ( OmniIntermediateBuffer, ) +from vllm_omni.worker_v2.streaming_audio import StreamingAudioOutput logger = init_logger(__name__) @@ -85,6 +86,9 @@ def _make_safe_get_rope(orig_get_rope): return _safe_get_rope +_NO_SEED = object() + + class OmniModelState(DefaultModelState): """Generic Omni ``ModelState`` — works for **all** Omni model stages. @@ -146,7 +150,7 @@ class OmniModelState(DefaultModelState): self.have_multimodal_outputs: bool = getattr(model, "have_multimodal_outputs", False) self._decode_preprocess = self._resolve_decode_preprocess(model) self._decode_preprocess_is_identity = bool(getattr(model, "mrv2_decode_preprocess_is_identity", False)) - self._mtp_generators: dict[str, torch.Generator] = {} + self._mtp_generators: dict[str, torch.Generator | None] = {} # Talker's codec_embedding dim may differ from hf_text_config.hidden_size; probe real dim. self._embed_dim = self._get_embed_dim(model, device) if self.has_preprocess else 0 @@ -167,6 +171,10 @@ class OmniModelState(DefaultModelState): self._mtp_offsets: torch.Tensor | None = None self._mtp_sample_uniforms: torch.Tensor | None = None self._mtp_runner: Any | None = None + # Talker stream decode: each request's next frame position, and the + # side-stream event the step's PCM output copy must wait on. + self._stream_pos: dict[str, int] = {} + self._stream_decode_event: torch.cuda.Event | None = None if self._embed_dim > 0 and hasattr(model, "mtp"): max_bs = max_num_reqs self._mtp_input_ids = torch.zeros(max_bs, dtype=torch.long, device=device) @@ -218,6 +226,17 @@ class OmniModelState(DefaultModelState): input_batch, text_hidden, sampled_token_ids, multimodal_outputs, mtp_batch_descriptor_dispatcher ) + def prepare_streaming_audio_output( + self, input_batch: InputBatch, req_states: RequestState, outputs: dict[str, Any] + ) -> StreamingAudioOutput | None: + return self._eager_state.prepare_audio_output(input_batch, req_states, outputs) + + def on_requests_finished(self, req_ids: set[str]) -> None: + self._eager_state.finish_audio(req_ids) + + def on_request_preempted(self, req_id: str, req_index: int) -> None: + self._eager_state.suspend_audio(req_id, req_index) + def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None: """Let the model replace or wrap the MRv2 sampler (``mrv2_custom_sampler``).""" hook = getattr(self.model, "mrv2_custom_sampler", None) @@ -372,6 +391,7 @@ class OmniModelState(DefaultModelState): def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: super().add_request(req_index, new_req_data) self.intermediate_buffer.add_request(req_index, new_req_data) + self._eager_state.resume_audio(new_req_data.req_id, req_index) self._initialize_upstream_warmup_buffer(req_index, new_req_data.req_id) def _initialize_upstream_warmup_buffer(self, req_index: int, req_id: str) -> None: @@ -407,6 +427,7 @@ class OmniModelState(DefaultModelState): req_id = self.intermediate_buffer.buffers[req_index].get("req_id") if req_id is not None: getattr(self, "_mtp_generators", {}).pop(req_id, None) + self._stream_pos.pop(req_id, None) getattr(self, "_first_audio_requests", set()).discard(req_id) getattr(self, "_eager_ready", {}).pop(req_index, None) getattr(self, "_eager_settled", {}).pop(req_index, None) @@ -673,6 +694,7 @@ class OmniModelState(DefaultModelState): preprocess_entries: list[tuple[int, int, int, int, dict[str, Any], bool]] = [] settled = self._eager_settled if self._eager_fastpath else None settled_rows: list[tuple[int, int, int, str]] = [] + replay_rows: list[tuple[int, int, str, bool]] = [] # Declared identity decode has no per-row updates or hook work. skip_decode_rows = self._decode_preprocess_is_identity and not self._eager_mtp is_prefilling_np = getattr(input_batch, "is_prefilling_np", None) if skip_decode_rows else None @@ -684,6 +706,20 @@ class OmniModelState(DefaultModelState): continue start = int(input_batch.query_start_loc_np[i]) n_tok = int(input_batch.num_scheduled_tokens[i]) + if ( + self._eager_mtp + and self.model.stream_decoder is not None + and self._eager_state.replay_inputs( + buf["req_id"], + req_idx, + int(input_batch.num_computed_tokens_np[i]), + input_ids[start : start + n_tok], + embeds[start : start + n_tok], + ) + ): + if input_batch.num_computed_prefill_tokens_np[i] + n_tok >= input_batch.prefill_len_np[i]: + replay_rows.append((i, req_idx, buf["req_id"], True)) + continue if settled and n_tok == 1 and settled.get(req_idx) == buf["req_id"]: settled_rows.append((i, req_idx, start, buf["req_id"])) continue @@ -725,7 +761,8 @@ class OmniModelState(DefaultModelState): if not is_prefill or int(info.get("_omni_num_computed_tokens", 0)) + n_tok >= int(info.get("_omni_prompt_len", 0)) ] - + [(i, req_idx, str(req_id), False) for i, req_idx, _start, req_id in settled_rows], + + [(i, req_idx, str(req_id), False) for i, req_idx, _start, req_id in settled_rows] + + replay_rows, input_ids, ) @@ -760,7 +797,7 @@ class OmniModelState(DefaultModelState): batch_offsets = None else: # A pageable host tensor would block on the stream here. - batch_offsets = torch.as_tensor(starts, device=input_ids.device, dtype=torch.long) + batch_offsets = index_to_device(starts, input_ids.device) batch_ids = input_ids.index_select(0, batch_offsets) batch_embeds = embeds.index_select(0, batch_offsets) @@ -866,6 +903,9 @@ class OmniModelState(DefaultModelState): prepacked_mtp_inputs=prepacked_mtp_inputs, ) + if self._eager_mtp: + self._eager_state.record_inputs(input_batch, embeds) + def publish_sampled_embeddings( self, input_batch: InputBatch, sampled_token_ids: torch.Tensor ) -> tuple[dict[str, Any], torch.cuda.Event] | None: @@ -989,7 +1029,7 @@ class OmniModelState(DefaultModelState): offsets.copy_(query_start_loc[:bsz], non_blocking=True) return offsets - return torch.as_tensor([start for _i, start, _mtp in mtp_batches], device=device, dtype=torch.long) + return index_to_device([start for _i, start, _mtp in mtp_batches], device) def _run_batched_mtp( self, @@ -1180,15 +1220,21 @@ class OmniModelState(DefaultModelState): sampling_params: Any, device: torch.device, ) -> torch.Generator | None: - resolve_seed = getattr(self.model, "get_mtp_seed", None) - seed = resolve_seed(sampling_params) if resolve_seed is not None else None - if seed is None: - return None - cache = getattr(self, "_mtp_generators", None) if cache is None: cache = {} self._mtp_generators = cache + generator = cache.get(req_id, _NO_SEED) + if generator is None: + # Resolved earlier: this request has no seed. + return None + if generator is not _NO_SEED and generator.device == device: + return generator + resolve_seed = getattr(self.model, "get_mtp_seed", None) + seed = resolve_seed(sampling_params) if resolve_seed is not None else None + if seed is None: + cache[req_id] = None + return None generator = cache.get(req_id) if generator is None or generator.device != device: generator = torch.Generator(device=device) diff --git a/vllm_omni/worker_v2/omni_ar_model_runner.py b/vllm_omni/worker_v2/omni_ar_model_runner.py index 32622e8ef..73434b8a1 100644 --- a/vllm_omni/worker_v2/omni_ar_model_runner.py +++ b/vllm_omni/worker_v2/omni_ar_model_runner.py @@ -14,6 +14,7 @@ Extends ``OmniGPUModelRunner`` with: from __future__ import annotations from contextlib import nullcontext +from functools import partial from typing import Any, cast import numpy as np @@ -37,6 +38,7 @@ from vllm_omni.utils.mm_outputs import partition_flat_payload from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner from vllm_omni.worker_v2.omni_sampler import sample_with_output from vllm_omni.worker_v2.output_snapshot import PackedOutputSnapshot, RequestOutputSnapshot, pack_output_snapshot +from vllm_omni.worker_v2.streaming_audio import StreamingAudioOutput logger = init_logger(__name__) _ASYNC_MM_SNAPSHOT_MAX_BUCKETS_PER_SLOT = 64 @@ -250,6 +252,13 @@ class OmniARModelRunner(OmniGPUModelRunner): model_runner_output.kv_extracted_req_ids = kv_extracted model_runner_output._async_chunk = bool(getattr(self.model_config, "async_chunk", False)) + prepare_streaming = getattr(self.model_state, "prepare_streaming_audio_output", None) + streaming_audio = ( + prepare_streaming(input_batch, self.req_states, multimodal_outputs) + if need_pooler and multimodal_outputs and callable(prepare_streaming) + else None + ) + # --- Async D2H via OmniAsyncOutput --- materialize_native = self._uses_native_output_materializer() async_output = OmniAsyncOutput( @@ -266,6 +275,7 @@ class OmniARModelRunner(OmniGPUModelRunner): finalize_multimodal=sampling_output.finalize_multimodal, check_ep_fault=self.check_ep_fault, routed_experts=routed_experts, + streaming_audio=streaming_audio, extra_multimodal_outputs=extra_outputs, ) self._release_multimodal_snapshot(snapshot_slot, async_output.copy_event) @@ -297,6 +307,10 @@ class OmniARModelRunner(OmniGPUModelRunner): def _retain_multimodal_outputs(self, outputs: dict[str, Any]) -> dict[str, Any]: if not bool(getattr(self.model_config, "async_chunk", False)) or not outputs: return outputs + if getattr(self.model, "mm_outputs_fresh_per_step", False): + # Freshly allocated each step and filled after sampling; nothing + # overwrites them before the host copy. + return outputs slot_index = self._async_mm_snapshot_cursor if self._async_mm_snapshot_pending[slot_index]: raise RuntimeError("Async multimodal snapshot ring exhausted before sample_tokens().") @@ -516,6 +530,12 @@ def _async_copy_mm_value( for key, val in value.items() } if isinstance(value, list): + if len(value) > 1 and isinstance(value[0], torch.Tensor) and value[0].device.type == "cpu": + first = value[0] + if all(val is first for val in value): + # One shared per-step host tensor (e.g. every request's sample + # rate): one copy serves all entries. + return [_async_copy_tensor(first)] * len(value) return [ _async_copy_mm_value( val, @@ -679,6 +699,7 @@ class OmniAsyncOutput(AsyncModelRunnerOutput): finalize_multimodal: Any | None = None, check_ep_fault: bool = False, routed_experts: RoutedExpertsTensors | None = None, + streaming_audio: StreamingAudioOutput | None = None, extra_multimodal_outputs: tuple[dict[str, Any], torch.cuda.Event] | None = None, ): self.model_runner_output = model_runner_output @@ -768,17 +789,31 @@ class OmniAsyncOutput(AsyncModelRunnerOutput): self._hidden_cpu: torch.Tensor | None = None self._mm_cpu: dict[str, Any] = {} self._mm_snapshot: dict[str, Any] = {} + self._streaming_audio = ( + streaming_audio.to_cpu( + copy_stream, partial(_async_copy_tensor, copy_stream=copy_stream, pin_memory=pin_memory) + ) + if streaming_audio is not None + else None + ) if self._need_pooler and (self._async_chunk or self._finalize_multimodal is not None): # CUDA graph replay reuses the model's output buffers. Take # ownership directly in pinned host memory on the output copy # stream so deferred finalization never performs a blocking # D2H copy on the runner thread. - self._mm_snapshot = _async_copy_mm( - multimodal_outputs, - self._total_tokens, - copy_stream=copy_stream, - pin_memory=pin_memory, - ) + if ( + streaming_audio is None + or self._finalize_multimodal is not None + or extra_multimodal_outputs is not None + ): + # PCM owns its request partition. Generic codes/meta are + # discarded unless a finalizer or extra payload needs them. + self._mm_snapshot = _async_copy_mm( + multimodal_outputs, + self._total_tokens, + copy_stream=copy_stream, + pin_memory=pin_memory, + ) if extra_multimodal_outputs: # Produced after sampling on the producer stream; copy only # once its completion event has been observed. @@ -856,7 +891,9 @@ class OmniAsyncOutput(AsyncModelRunnerOutput): pooler_client is not None and len(pooler_client) != self._num_reqs ): raise ValueError("Model-owned output snapshot does not match the request batch") - else: + elif self._streaming_audio is None: + # In-stage PCM already owns its request partition. Building + # generic code payloads here would immediately discard them. pooler_inter, pooler_client = OmniARModelRunner._build_async_chunk_outputs_from_mm( self._mm_snapshot, self._query_start_loc_np, @@ -865,6 +902,8 @@ class OmniAsyncOutput(AsyncModelRunnerOutput): self._total_tokens, self._padded_total_tokens, ) + if self._streaming_audio is not None: + pooler_inter, pooler_client = None, self._streaming_audio.get_output() self.model_runner_output.pooler_output = None if self._async_chunk else pooler_inter self.model_runner_output.inter_stage_outputs = pooler_inter self.model_runner_output.multimodal_outputs = ( diff --git a/vllm_omni/worker_v2/omni_model_runner.py b/vllm_omni/worker_v2/omni_model_runner.py index 10f4053cf..9e73201e5 100644 --- a/vllm_omni/worker_v2/omni_model_runner.py +++ b/vllm_omni/worker_v2/omni_model_runner.py @@ -437,6 +437,10 @@ class OmniGPUModelRunner(GPUModelRunner): capture_first_frame = getattr(self.model, "capture_first_frame_graphs", None) if callable(capture_first_frame): capture_first_frame() + capture_stream = getattr(self.model, "capture_stream_decode_graphs", None) + if callable(capture_stream): + sizes = [int(s) for s in self.model_state._get_mtp_capture_sizes()] + capture_stream(sorted(set(sizes + [1]))) return result def _dispatch_mtp_batch_descriptor(self, num_mtp_reqs: int) -> Any: @@ -746,7 +750,16 @@ class OmniGPUModelRunner(GPUModelRunner): # Request-owned state outlives generation slots and chunk boundaries. # Notify even if the last chunk already released its runner slot. on_finished(finished) + finish_state = getattr(self.model_state, "on_requests_finished", None) + if finished and callable(finish_state): + finish_state(finished) preempted = scheduler_output.preempted_req_ids + suspend_state = getattr(self.model_state, "on_request_preempted", None) + if preempted and callable(suspend_state): + for req_id in preempted - finished: + idx = self.req_states.req_id_to_index.get(req_id) + if idx is not None: + suspend_state(req_id, idx) all_done = finished | preempted if preempted else finished for req_id in all_done: idx = self.req_states.req_id_to_index.get(req_id) diff --git a/vllm_omni/worker_v2/streaming_audio.py b/vllm_omni/worker_v2/streaming_audio.py new file mode 100644 index 000000000..bf78de678 --- /dev/null +++ b/vllm_omni/worker_v2/streaming_audio.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Request-owned PCM accumulation for a stage that decodes audio in place.""" + +from __future__ import annotations + +import threading +from collections.abc import Callable, Iterable +from dataclasses import dataclass, field + +import numpy as np +import torch + +from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY + + +@dataclass +class AudioRequest: + samples_per_frame: int + chunk_frames: int + active: bool = True + emitted: bool = False + direct_first_pending: bool = False + first_audio_required: bool = False + pending: list[torch.Tensor | np.ndarray] = field(default_factory=list) + pending_samples: int = 0 + lock: threading.Lock = field(default_factory=threading.Lock) + + def accept_first_audio(self) -> None: + with self.lock: + if self.active: + self.direct_first_pending = True + + def push(self, pcm: torch.Tensor | np.ndarray, ended: bool) -> torch.Tensor | None: + with self.lock: + if not self.active: + return None + size = pcm.numel() if isinstance(pcm, torch.Tensor) else pcm.size + if self.direct_first_pending: + self.direct_first_pending = False + if size: + pcm = pcm[self.samples_per_frame :] + size = max(0, size - self.samples_per_frame) + self.emitted = True + self.first_audio_required = True + if size: + self.pending.append(pcm) + self.pending_samples += size + frames = self.pending_samples // self.samples_per_frame + out = None + if frames and (ended or not self.emitted or frames >= self.chunk_frames): + if all(isinstance(part, np.ndarray) for part in self.pending): + array = np.concatenate(self.pending) if len(self.pending) > 1 else self.pending[0] + out = torch.from_numpy(array) + else: + parts = [torch.from_numpy(part) if isinstance(part, np.ndarray) else part for part in self.pending] + out = torch.cat(parts) if len(parts) > 1 else parts[0] + self.pending.clear() + self.pending_samples = 0 + self.emitted = True + if ended: + self.active = False + self.pending.clear() + self.pending_samples = 0 + return out + + def finish(self) -> None: + with self.lock: + self.active = False + self.pending.clear() + self.pending_samples = 0 + + +class StreamingAudioBuffer: + def __init__(self, samples_per_frame: int, chunk_frames: int) -> None: + self.samples_per_frame = samples_per_frame + self.chunk_frames = chunk_frames + self.requests: dict[str, AudioRequest] = {} + + def add(self, request_id: str) -> AudioRequest: + if request_id not in self.requests: + self.requests[request_id] = AudioRequest(self.samples_per_frame, self.chunk_frames) + return self.requests[request_id] + + def finish(self, request_ids: Iterable[str]) -> None: + for request_id in request_ids: + state = self.requests.pop(request_id, None) + if state is not None: + # Already queued outputs retain this object, never a fresh lookup + # by ID. A late copy cannot resurrect a cancelled request. + state.finish() + + +@dataclass +class StreamingAudioOutput: + wav: torch.Tensor + valid: torch.Tensor + sample_rate: torch.Tensor + event: torch.cuda.Event | None + query_start_loc: np.ndarray + requests: list[AudioRequest | None] + length_end: np.ndarray + + def to_cpu( + self, + copy_stream: torch.cuda.Stream, + copy_tensor: Callable[[torch.Tensor], torch.Tensor], + ) -> StreamingAudioOutput: + if self.event is not None: + copy_stream.wait_event(self.event) + return StreamingAudioOutput( + copy_tensor(self.wav), + copy_tensor(self.valid), + self.sample_rate, + None, + self.query_start_loc, + self.requests, + self.length_end, + ) + + def get_output(self) -> list[dict[str, torch.Tensor] | None]: + valid = self.valid.numpy().astype(bool) + # Each D2H batch owns its host allocation. NumPy slices retain that + # backing tensor while a request accumulates PCM across steps; no + # per-request Torch views are needed until an audio chunk is emitted. + # Keep the tensor path for dtypes NumPy cannot represent (e.g. BF16). + wav = self.wav.numpy() if self.wav.dtype in (torch.float16, torch.float32, torch.float64) else None + output: list[dict[str, torch.Tensor] | None] = [] + required_marker = None + for i, state in enumerate(self.requests): + payload = None + start, end = self.query_start_loc[i : i + 2] + if state is not None and end > start: + if wav is not None: + if end - start == 1: + pcm = wav[start:end].reshape(-1) if valid[start] else wav[0:0].reshape(-1) + else: + pcm = wav[start:end][valid[start:end]].reshape(-1) + elif end - start == 1: + pcm = self.wav[start:end].reshape(-1) if valid[start] else self.wav.new_empty(0) + else: + pcm = self.wav[start:end][torch.from_numpy(valid[start:end])].reshape(-1) + ended = bool(self.length_end[i]) or not valid[end - 1] + chunk = state.push(pcm, ended) + if chunk is not None: + payload = {"model_outputs": chunk, "sr": self.sample_rate} + if state.first_audio_required and (chunk is not None or ended): + if required_marker is None: + required_marker = torch.tensor(True) + if payload is None: + payload = {} + payload[FIRST_AUDIO_REQUIRED_KEY] = required_marker + output.append(payload) + return output From b10270a01ce6a1691120951a3969202e5f785957 Mon Sep 17 00:00:00 2001 From: linzhenpl07 <60651557+linzhenpl07@users.noreply.github.com> Date: Fri, 2 Oct 2026 20:05:04 +0800 Subject: [PATCH 12/22] [AR-Diffusion] Page in kernel-legal blocks so the default resolution can run (#6481) Signed-off-by: linzhenpl07 Co-authored-by: Claude Opus 5 Co-authored-by: Zhou Taichang --- .../ar_diffusion/test_capability_runner.py | 114 ++- tests/diffusion/ar_diffusion/test_kv_cache.py | 68 ++ .../ar_diffusion/test_paged_attention.py | 667 +++++++++++++++++- .../ar_diffusion/kv_cache/manager.py | 47 +- .../ar_diffusion/kv_cache/paged.py | 53 +- .../ar_diffusion/kv_cache/paged_attention.py | 272 ++++++- .../ar_diffusion/kv_cache/state.py | 1 + vllm_omni/experimental/ar_diffusion/runner.py | 57 +- 8 files changed, 1229 insertions(+), 50 deletions(-) diff --git a/tests/diffusion/ar_diffusion/test_capability_runner.py b/tests/diffusion/ar_diffusion/test_capability_runner.py index d753efbf6..01d8b2844 100644 --- a/tests/diffusion/ar_diffusion/test_capability_runner.py +++ b/tests/diffusion/ar_diffusion/test_capability_runner.py @@ -376,10 +376,17 @@ def test_lru_eviction_releases_blocks_and_notifies_pipeline(): def test_budget_reduced_capacity_drives_runner_lru(): pipeline = CapablePipeline(tiny_spec(capacity=2)) - # One tiny LingBot-like session requires 128 bytes; two require 192. + # One tiny LingBot-like session requires 1,808 bytes; two require 2,592. + # + # These read 128 and 192 while the paging unit was the frame. tiny_spec + # declares one token per frame, so it used to page a single token at a + # time -- a block size no attention kernel accepts. Paging at 16 makes + # each page sixteen times larger. The block *counts* are unchanged, and + # so is what this test is about: a budget that fits fewer sessions than + # the pipeline asked for must drive the runner's LRU. runner = make_runner( pipeline, - available_bytes=128, + available_bytes=2048, gpu_memory_fraction=1.0, ) kv = runner.kv_cache @@ -665,3 +672,106 @@ def test_pipeline_without_warmup_provider_is_safely_skipped(monkeypatch): monkeypatch.setattr(runner, "execute_model", fail_if_called) runner._warmup_ar_rollout() assert execute.called is False + + +# ── paging unit vs eviction unit ──────────────────────────────────────────── + + +def test_a_frame_that_is_already_a_legal_block_keeps_its_paging(): + """Every resolution that runs today must page exactly as it does today.""" + from vllm_omni.experimental.ar_diffusion.runner import paging_block_size + + for tokens_per_frame in (16, 320, 1440, 4096): + assert tokens_per_frame % 16 == 0 + assert paging_block_size(tokens_per_frame) == tokens_per_frame + + +def test_a_frame_the_kernel_would_reject_pages_at_the_finest_legal_unit(): + """832x480 is the checkpoint's own default and the kernel rejects it. + + (480/16) x (832/16) = 1560 tokens per frame, and 1560 % 16 == 8, so a + frame-sized block is not something FlashAttention's paged kernel accepts. + """ + from vllm_omni.experimental.ar_diffusion.runner import paging_block_size + + assert 1560 % 16 == 8 + assert paging_block_size(1560) == 16 + # And whatever it returns is always something the kernel takes. + for tokens_per_frame in range(1, 200): + assert paging_block_size(tokens_per_frame) % 16 == 0 + + +# ── the kernel's constraint is data, not a constant ───────────────────────── + + +def _multiple_of(base): + from vllm.v1.attention.backend import MultipleOf + + return MultipleOf(base) + + +def test_the_dispatching_module_answers_in_vllms_own_vocabulary(): + """AR-Diffusion calls flash_attn itself, so no backend can be asked. + + The answer has to come from the module that picks the kernel, and it has to + have the shape vLLM uses, so the policy that consumes it does not care + where it came from. + """ + from vllm.v1.attention.backend import MultipleOf + + from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import supported_kernel_block_sizes + + advertised = supported_kernel_block_sizes() + assert advertised, "a kernel that accepts nothing cannot be paged for" + for entry in advertised: + assert isinstance(entry, int | MultipleOf) + + +def test_a_kernel_that_takes_one_fixed_size_gets_that_size(): + """hpc_attn advertises [64] -- not a multiple, a single legal value. + + A frame of 1560 tokens is not 64, and 1560 % 64 == 24, so the frame cannot + be the block and the only thing left to page at is 64. Assuming multiples + of 16 here would hand the kernel a block size it rejects. + """ + from vllm_omni.experimental.ar_diffusion.runner import paging_block_size + + assert 1560 % 64 == 24 + assert paging_block_size(1560, [64]) == 64 + assert paging_block_size(64, [64]) == 64 + assert paging_block_size(128, [64]) == 64 # a legal multiple is still not the advertised size + + +def test_a_kernel_with_large_pages_still_keeps_a_whole_frame_when_it_fits(): + """FlashInfer offers 128-and-up pages on Blackwell and not on Hopper. + + The same checkpoint on the two cards must therefore page differently, which + is exactly what a constant cannot express. + """ + from vllm_omni.experimental.ar_diffusion.runner import paging_block_size + + hopper = [_multiple_of(16)] + blackwell = [_multiple_of(16), 128, 256] + assert paging_block_size(1440, hopper) == 1440 # a legal multiple of 16 + assert paging_block_size(1440, blackwell) == 1440 + assert paging_block_size(1560, hopper) == 16 + assert paging_block_size(1560, blackwell) == 16 # finest legal unit, not the largest + + +def test_whatever_comes_back_is_always_something_the_kernel_accepts(): + from vllm_omni.experimental.ar_diffusion.runner import paging_block_size + + for advertised, check in ( + ([_multiple_of(16)], lambda n: n % 16 == 0), + ([64], lambda n: n == 64), + ([_multiple_of(32), 48], lambda n: n % 32 == 0 or n == 48), + ): + for tokens_per_frame in range(1, 300): + assert check(paging_block_size(tokens_per_frame, advertised)) + + +def test_a_kernel_that_advertises_nothing_is_an_error_not_a_guess(): + from vllm_omni.experimental.ar_diffusion.runner import paging_block_size + + with pytest.raises(ValueError, match="no legal block size"): + paging_block_size(1560, []) diff --git a/tests/diffusion/ar_diffusion/test_kv_cache.py b/tests/diffusion/ar_diffusion/test_kv_cache.py index 84908fef1..e73713f8e 100644 --- a/tests/diffusion/ar_diffusion/test_kv_cache.py +++ b/tests/diffusion/ar_diffusion/test_kv_cache.py @@ -625,6 +625,74 @@ def test_non_contiguous_branch_indices_rejected(): ) +# --- eviction against the block table --------------------------------------- + + +class _RecordingManager(ChunkWindowManager): + """Capture the block range eviction asks for, without a real block pool. + + ``remove_skipped_blocks`` reads three attributes and calls one method. A + pool would add nothing: the bug under test is the conversion from the + spec's units into the block table's, and that is entirely visible here. + """ + + def __init__(self, spec): + self.kv_cache_spec = spec + self.sliding_window = spec.sliding_window + self.block_size = spec.block_size + self.freed: tuple[int, int] | None = None + + def _remove_blocks_in_range(self, request_id, first_block, last_block): + self.freed = (first_block, last_block) + + +def test_sink_is_preserved_when_a_frame_spans_many_blocks(): + # The shipped geometry: 832x480 is 1560 tokens per frame, 1560 % 16 == 8 is + # not a legal block size, so paging falls back to 16 and a frame spans 98 + # blocks. Every other eviction test uses chunk_size == BLOCK, where a frame + # is a block and a sink counted in frames indexes the block table correctly + # by accident -- which is why this went unnoticed. + chunk_size, sink_chunks, window_chunks = 1560, 9, 9 + spec = ChunkWindowSpec( + block_size=16, + num_kv_heads=4, + head_size=64, + dtype=torch.float16, + sliding_window=window_chunks * chunk_size, + chunk_size=chunk_size, + window_chunks=window_chunks, + sink_chunks=sink_chunks, + ) + manager = _RecordingManager(spec) + + # Far enough past sink + window that eviction has definitely started. + manager.remove_skipped_blocks("req", total_computed_tokens=30 * chunk_size) + assert manager.freed is not None + first_freed, last_freed = manager.freed + + # The assertion is on the token boundary rather than a block count, so it + # stays meaningful if the paging unit changes again. + sink_tokens = sink_chunks * chunk_size + assert first_freed * spec.block_size >= sink_tokens, ( + f"eviction starts at token {first_freed * spec.block_size}, inside the {sink_tokens}-token sink" + ) + # And it must not over-protect: at most one straddling block beyond the sink. + assert (first_freed - 1) * spec.block_size < sink_tokens + # Only wholly-skipped blocks are released. + skipped = manager.get_num_skipped_tokens(30 * chunk_size) + assert last_freed * spec.block_size <= sink_tokens + skipped + + +def test_sink_conversion_is_identity_when_a_frame_is_a_block(): + # Every geometry that paged one frame per block must be untouched. + spec = make_spec(chunk_size=BLOCK, window_chunks=2, sink_chunks=1) + manager = _RecordingManager(spec) + manager.remove_skipped_blocks("req", total_computed_tokens=8 * BLOCK) + assert manager.freed is not None + first_freed, _ = manager.freed + assert first_freed == spec.sink_chunks + + def test_staged_window_reuse_matches_a_full_gather(): """Reusing the staged history must be byte-identical to re-gathering the whole window. diff --git a/tests/diffusion/ar_diffusion/test_paged_attention.py b/tests/diffusion/ar_diffusion/test_paged_attention.py index 7014a60c7..47e9b028e 100644 --- a/tests/diffusion/ar_diffusion/test_paged_attention.py +++ b/tests/diffusion/ar_diffusion/test_paged_attention.py @@ -23,6 +23,7 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import ( ) from vllm_omni.experimental.ar_diffusion.kv_cache import paged_attention as paged_attention_module from vllm_omni.experimental.ar_diffusion.kv_cache.config import KV_GATHER_ENV +from vllm_omni.experimental.ar_diffusion.kv_cache.paged import ChunkWindowManager from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] @@ -36,10 +37,26 @@ NEG = "negative" def make_state( - *, num_layers=1, window_chunks=2, dtype=torch.float32, device=torch.device("cpu"), reuse_history_staging=False + *, + num_layers=1, + window_chunks=2, + chunk_size=BLOCK, + sink_chunks=0, + reset_at_boundary=False, + dtype=torch.float32, + device=torch.device("cpu"), + reuse_history_staging=False, ): + """Build a cache. ``chunk_size`` defaults to the block size, but the two are + independent -- the shipped 832x480 gives 1560 tokens per frame against + 16-token blocks, so a frame is 97.5 blocks.""" cfg = ARDiffusionKVConfig( - enable=True, chunk_size=BLOCK, window_chunks=window_chunks, reuse_history_staging=reuse_history_staging + enable=True, + chunk_size=chunk_size, + window_chunks=window_chunks, + sink_chunks=sink_chunks, + reset_at_boundary=reset_at_boundary, + reuse_history_staging=reuse_history_staging, ) kv = ARDiffusionKVCache( cfg, @@ -109,11 +126,13 @@ def _commit_video_span( n_chunks: int, dtype: torch.dtype, device: torch.device, + chunk_size: int = BLOCK, ) -> tuple[torch.Tensor, torch.Tensor]: - ctx = st.get_kv_caches(kv_branch, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx + span = n_chunks * chunk_size + ctx = st.get_kv_caches(kv_branch, seq_len=span, commit_current=True)[0].forward_ctx ctx.ensure_video_slots(device) - k = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device) - v = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + k = torch.randn(1, span, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + v = torch.randn(1, span, N_HEADS, HEAD_DIM, dtype=dtype, device=device) kv._k_pools[0][ctx.current_video_slot_mapping] = k[0] kv._v_pools[0][ctx.current_video_slot_mapping] = v[0] st.commit_paged_context(kv_branch) @@ -547,6 +566,644 @@ def test_custom_op_compiles_fullgraph_without_recompile_on_value_change(monkeypa torch._dynamo.reset() +# ── a frame that is not a whole number of blocks ──────────────────────────── + +# 24 tokens per chunk against 16-token blocks leaves 8 over, the same remainder +# the shipped 832x480 produces (1560 % 16 == 8). Every fixture above holds +# chunk_size == BLOCK, so nothing there can reach this path. +RAGGED_CHUNK = 24 + + +@pytest.mark.parametrize("commit_current", [False, True]) +@pytest.mark.parametrize("action_len", [0, 3]) +def test_a_ragged_chunk_still_matches_the_dense_reference(commit_current, action_len): + """Attention reads a sequence, not a set of blocks. + + With one committed chunk of 24 tokens the history stops 8 slots into its + last block. The committing path writes straight after it and is fine. The + scratch path used to restart at slot zero of a fresh region, which left + those 8 slots unwritten -- and since the kernel consumes the block table as + one contiguous run, it read them as if they were tokens and shifted every + token after them by 8. Shapes stayed correct throughout, so only the values + showed it. + """ + torch.manual_seed(0) + device = torch.device("cpu") + dtype = torch.float32 + # window_chunks=2 with one chunk of history keeps the whole sequence + # resident, so the window does not also round up here; that case is + # test_a_ragged_window_keeps_the_block_table_a_fixed_shape below. + kv, st = make_state(dtype=dtype, device=device, window_chunks=2, chunk_size=RAGGED_CHUNK) + + history_k, history_v = _commit_video_span( + kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device, chunk_size=RAGGED_CHUNK + ) + + ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=commit_current)[0].forward_ctx + ctx.ensure_video_slots(device) + assert ctx.start_offset == RAGGED_CHUNK % BLOCK == 8 + + current_k = torch.randn(1, RAGGED_CHUNK, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + current_v = torch.randn(1, RAGGED_CHUNK, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0] + kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0] + + action_k = action_v = None + if action_len: + ctx.ensure_action_slots(action_len, device) + action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + kv._k_pools[0][ctx.action_slot_mapping] = action_k[0] + kv._v_pools[0][ctx.action_slot_mapping] = action_v[0] + + query = torch.randn(1, RAGGED_CHUNK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table( + action_len=action_len, query_len=query.shape[1], device=device + ) + paged = ar_diffusion_paged_attention( + query, + kv.key_cache(0), + kv.value_cache(0), + block_table=block_table, + query_start_loc=query_start_loc, + seq_lens=seq_lens, + max_query_len=max_query_len, + max_seq_len=max_seq_len, + softmax_scale=HEAD_DIM**-0.5, + causal=False, + ) + + new_k = torch.cat([history_k, current_k], dim=1) + new_v = torch.cat([history_v, current_v], dim=1) + if action_len: + new_k = torch.cat([new_k, action_k], dim=1) + new_v = torch.cat([new_v, action_v], dim=1) + ref = _dense_attention(query, new_k, new_v) + + torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5) + + +def test_a_ragged_scratch_chunk_is_physically_continuous_with_its_history(): + """The write targets themselves must leave no hole. + + This is the property the values above depend on, asserted directly so a + regression names the cause rather than showing a numeric mismatch. + """ + device = torch.device("cpu") + kv, st = make_state(device=device, window_chunks=2, chunk_size=RAGGED_CHUNK) + _commit_video_span(kv, st, kv_branch=POS, n_chunks=1, dtype=torch.float32, device=device, chunk_size=RAGGED_CHUNK) + + ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=False)[0].forward_ctx + ctx.ensure_video_slots(device) + + history_tail_block = ctx.history_block_ids[-1] + # The chunk starts inside the history's own last block, not at a fresh one. + assert ctx.current_video_block_ids[0] == history_tail_block + first_slot = int(ctx.current_video_slot_mapping[0]) + assert first_slot == history_tail_block * BLOCK + 8 + + # And the action region must not be charged for that managed block. + ctx.ensure_action_slots(3, device) + assert ctx._scratch_blocks_used == len(ctx.current_video_block_ids) - 1 + assert ctx.action_scratch_block_ids[0] not in ctx.current_video_block_ids + + +def test_a_ragged_window_keeps_the_block_table_a_fixed_shape(): + """The table's shape must not track how full the window is. + + With a sink the visible window is 2*24 + 24 = 72 tokens, which is 4.5 + blocks. Deriving the width by flooring that would understate the capacity, + the max() against the live block count would take over, and the shape would + change as the window filled -- recompiling the graph. The same floor would + also let max_seq_len come out under the kv_len actually passed. + """ + device = torch.device("cpu") + kv, st = make_state(device=device, window_chunks=2, sink_chunks=1, chunk_size=RAGGED_CHUNK) + # The window is deliberately not a whole number of blocks. + max_video_tokens = kv.spec.sliding_window + kv.spec.sink_chunks * kv.spec.chunk_size + assert max_video_tokens % BLOCK != 0 + + widths: set[int] = set() + kv_lens: list[int] = [] + for _ in range(6): + ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx + ctx.ensure_video_slots(device) + block_table, _, seq_lens, _, max_seq_len = ctx.build_block_table( + action_len=0, query_len=RAGGED_CHUNK, device=device + ) + widths.add(int(block_table.shape[1])) + # A bound handed to the kernel has to actually bound the sequence. + assert int(seq_lens[0]) <= max_seq_len + kv_lens.append(int(seq_lens[0])) + st.commit_paged_context(POS) + + assert len(widths) == 1, f"block table width varied across ticks: {sorted(widths)}" + # The window has to have actually filled, or none of the above was tested. + assert max(kv_lens) > min(kv_lens) + assert max(kv_lens) <= ctx.max_video_blocks * BLOCK + # Whole blocks are read, so each of the two boundaries -- the sink's end and + # the recent window's start -- can contribute up to BLOCK - 1 tokens the + # window no longer needs. + assert max(kv_lens) <= max_video_tokens + 2 * (BLOCK - 1) + + +def _poison_pools(kv: ARDiffusionKVCache) -> None: + for pool in (*kv._k_pools, *kv._v_pools): + pool.fill_(float("nan")) + + +def _write_prepared_chunk( + kv: ARDiffusionKVCache, ctx, *, dtype: torch.dtype, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor]: + """Write fresh K/V for a prepared chunk, poisoning the rest of its blocks first. + + Blocks are handed out whole and a reused block keeps whatever it held, so a + slot read without having been written would not look unusual. Every slot of + the chunk's blocks that this chunk does not write is set to NaN first -- + except the leading slots of a block the history still occupies -- so reading + one turns the attention output NaN. + """ + mapping = ctx.current_video_slot_mapping + first_slot = int(mapping[0]) + for index, block in enumerate(ctx.current_video_block_ids): + start = first_slot if index == 0 else block * BLOCK + kv._k_pools[0][start : (block + 1) * BLOCK] = float("nan") + kv._v_pools[0][start : (block + 1) * BLOCK] = float("nan") + k = torch.randn(ctx.seq_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + v = torch.randn(ctx.seq_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + kv._k_pools[0][mapping] = k + kv._v_pools[0][mapping] = v + return k, v + + +def _window_positions(end: int, *, sink: int, window: int, written, history: int, resident: set[int]) -> list[int]: + """Token positions attention must read at ``end``. + + Every token of the current chunk, plus every written history token in a + block that is still resident and holds a sink token or one of the window's + most recent tokens. + """ + start = end - window + blocks = {p // BLOCK for p in range(min(sink, end))} | {p // BLOCK for p in range(max(start, 0), end)} + return [p for p in sorted(written) if p >= history or (p // BLOCK in blocks and p // BLOCK in resident)] + + +def _model_block_indices(kv, adapter) -> dict[int, int]: + """Each resident block's index in model positions, by block id. + + Compaction removes the evicted gap after the sink's blocks and shifts the + rest of the table down by it, so a table index past the sink is + ``compacted_tokens`` behind the position it holds. + """ + block_size = kv.block_size + sink_blocks = -(-(kv.spec.sink_chunks * kv.spec.chunk_size) // block_size) + shift = adapter.compacted_tokens // block_size + return { + int(block): index if index < sink_blocks else index + shift + for index, block in enumerate(kv.block_table(adapter)) + if block != kv.null_block_id + } + + +@pytest.mark.parametrize( + ("sink_chunks", "window_chunks", "reset_at_boundary"), + [(1, 2, False), (0, 2, False), (1, 3, False), (2, 2, True), (0, 2, True)], +) +def test_attention_after_eviction_reads_every_kept_token_and_nothing_unwritten( + sink_chunks, window_chunks, reset_at_boundary +): + """Once the window slides, the table must still cover the window exactly. + + The sink and the recent window are separate token ranges, and each can + straddle a block edge on its own: 24-token chunks against 16-token blocks + put a boundary 8 tokens past an edge on alternating chunks. Counting blocks + off the ends of the resident list dropped the block holding the window's + first tokens whenever both boundaries straddled, and the returned length + ran past the last written slot. Neither changes a shape, so this compares + values against dense attention over the tokens the window keeps -- after + eviction, with every unwritten slot poisoned, on both the scratch and the + committing forward of every chunk. + + The reset case keeps only the sink across a boundary, so the recent window + reaches back over blocks that are no longer resident. Those must be skipped + by position rather than read. + + Eight chunks are enough for compaction to drop the evicted gap from the + table, after which storage positions run behind model positions. Removing + a gap that is not a whole number of chunks moves where eviction snaps, and + in the reset case the window then keeps tokens a boundary should have + dropped. + """ + torch.manual_seed(0) + device, dtype = torch.device("cpu"), torch.float32 + kv, st = make_state( + device=device, + window_chunks=window_chunks, + sink_chunks=sink_chunks, + reset_at_boundary=reset_at_boundary, + chunk_size=RAGGED_CHUNK, + ) + _poison_pools(kv) + sink, window = sink_chunks * RAGGED_CHUNK, int(kv.spec.sliding_window) + committed: dict[int, tuple[torch.Tensor, torch.Tensor]] = {} + evicted = compacted = False + + for chunk in range(8): + history = int(st.adapter(POS).absolute_num_computed_tokens) + end = history + RAGGED_CHUNK + for commit_current in (False, True): + ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=commit_current)[0].forward_ctx + ctx.ensure_video_slots(device) + k, v = _write_prepared_chunk(kv, ctx, dtype=dtype, device=device) + written = dict(committed) + written.update({history + i: (k[i], v[i]) for i in range(RAGGED_CHUNK)}) + + query = torch.randn(1, RAGGED_CHUNK, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table( + action_len=0, query_len=RAGGED_CHUNK, device=device + ) + resident = set(_model_block_indices(kv, st.adapter(POS)).values()) + positions = _window_positions( + end, sink=sink, window=window, written=written, history=history, resident=resident + ) + label = f"chunk {chunk}, commit_current={commit_current}" + assert int(seq_lens[0]) == len(positions), f"{label}: kv_len does not match the written slots it covers" + paged = ar_diffusion_paged_attention( + query, + kv.key_cache(0), + kv.value_cache(0), + block_table=block_table, + query_start_loc=query_start_loc, + seq_lens=seq_lens, + max_query_len=max_query_len, + max_seq_len=max_seq_len, + softmax_scale=HEAD_DIM**-0.5, + causal=False, + ) + ref = _dense_attention( + query, + torch.stack([written[p][0] for p in positions]).unsqueeze(0), + torch.stack([written[p][1] for p in positions]).unsqueeze(0), + ) + torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5, msg=lambda m, label=label: f"{label}: {m}") + committed = written + st.commit_paged_context(POS) + compacted = compacted or st.adapter(POS).compacted_tokens > 0 + evicted = evicted or compacted or kv.null_block_id in kv.block_table(st.adapter(POS)) + + assert evicted, "the window never slid, so nothing after eviction was tested" + assert compacted, "the table was never compacted, so storage and model positions never diverged" + + +@pytest.mark.parametrize( + ("sink_chunks", "window_chunks", "reset_at_boundary"), + [(1, 2, False), (0, 2, False), (1, 3, False), (2, 2, True), (0, 2, True)], +) +def test_compaction_does_not_change_which_tokens_stay_resident( + monkeypatch, sink_chunks, window_chunks, reset_at_boundary +): + """Compaction renumbers the block table; it must not change what eviction keeps. + + Storage positions run ``compacted_tokens`` behind model positions, and + eviction snaps to chunk boundaries counted in storage positions. Removing a + gap that is not a whole number of chunks moves every later snap, so eviction + keeps a different set of tokens -- in the reset case, tokens a boundary + should have dropped, which the window then reads. Attention checked against + what is resident cannot see that, so the oracle here is the same rollout + with compaction switched off. + """ + + def resident_blocks_per_tick(compact: bool) -> tuple[list[set[int]], int]: + with monkeypatch.context() as patch: + if not compact: + patch.setattr(ChunkWindowManager, "compact_block_table", lambda self, request_id: 0) + kv, st = make_state( + window_chunks=window_chunks, + sink_chunks=sink_chunks, + reset_at_boundary=reset_at_boundary, + chunk_size=RAGGED_CHUNK, + ) + ticks = [] + for _ in range(12): + ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx + ctx.ensure_video_slots(torch.device("cpu")) + ticks.append(set(_model_block_indices(kv, st.adapter(POS)).values())) + st.commit_paged_context(POS) + return ticks, st.adapter(POS).compacted_tokens + + expected, _ = resident_blocks_per_tick(compact=False) + actual, compacted_tokens = resident_blocks_per_tick(compact=True) + assert compacted_tokens > 0, "the table was never compacted, so nothing was tested" + for tick, (want, got) in enumerate(zip(expected, actual)): + assert got == want, f"tick {tick}: compaction kept blocks {sorted(got - want)} and dropped {sorted(want - got)}" + + +def test_history_staging_holds_a_ragged_window_and_restages_it_whole(monkeypatch): + """The staging buffers must fit the table width this module pads to. + + A sink or window that is not a whole number of blocks spans one more block + than its token count suggests, so sizing staging from tokens left it short + and the first prepared forward refused to run. And because such a chunk + shares its first block with the history's tail, the history a probe sees is + not left untouched by the next probe, so it is restaged rather than reused. + """ + monkeypatch.setenv(KV_GATHER_ENV, "1") + device = torch.device("cpu") + kv, st = make_state( + device=device, window_chunks=2, sink_chunks=1, chunk_size=RAGGED_CHUNK, reuse_history_staging=True + ) + capacity = int(kv.history_staging[0][0].shape[0]) + for chunk in range(6): + for probe in range(2): + ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=False)[0].forward_ctx + ctx.prepare(device=device, action_len=0, query_len=RAGGED_CHUNK) + assert int(ctx.max_seq_len) <= capacity + assert not ctx.reuse_history, f"chunk {chunk}, probe {probe}: a ragged chunk reused the staged history" + ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx + ctx.prepare(device=device, action_len=0, query_len=RAGGED_CHUNK) + st.commit_paged_context(POS) + + +@pytest.mark.parametrize("commit_current", [False, True]) +def test_action_tokens_after_a_partly_written_video_block_are_refused(commit_current): + """Action K/V follows the video blocks in the table, and the kernel reads it as one run. + + A 24-token chunk from an empty history ends 8 slots into its second 16-token + block, so the run would read those 8 unwritten slots as the first action + tokens and never reach the last ones. Every shape stays right, which is why + this is refused rather than computed. + """ + _, st = make_state(window_chunks=2, chunk_size=RAGGED_CHUNK) + ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=commit_current)[0].forward_ctx + with pytest.raises(ValueError, match="partly written video block"): + ctx.build_block_table(action_len=3, query_len=RAGGED_CHUNK + 3, device=torch.device("cpu")) + + +def test_the_shipped_geometry_reads_exactly_the_window_on_every_tick(): + """832x480 with the checkpoint's own sink and window, slot by slot. + + A 9-frame sink and a 9-frame window are each 8 tokens past a 16-token edge + but 28080 tokens together -- a whole number of blocks -- so a width taken + from their sum was one block short on every tick once the window had + filled. Twelve ticks of three frames, checking every slot the kernel reads + against the token that slot is supposed to hold. No attention is computed: + at this size the bookkeeping is the whole question. + """ + from vllm_omni.experimental.ar_diffusion.runner import paging_block_size + + frame, frames_per_tick = (480 // 16) * (832 // 16), 3 + block = paging_block_size(frame) + seq_len = frames_per_tick * frame + cfg = ARDiffusionKVConfig(enable=True, chunk_size=frame, window_chunks=9, sink_chunks=9) + kv = ARDiffusionKVCache( + cfg, + num_layers=1, + num_kv_heads=1, + head_size=8, + dtype=torch.float16, + block_size=block, + max_model_len=1 << 20, + available_bytes=1 << 27, + kv_branches=(ARDiffusionKVBranchSpec(POS, 0),), + session_capacity=1, + frames_per_block=frames_per_tick, + max_scratch_tokens_per_branch=seq_len, + device=torch.device("cpu"), + ) + st = ARDiffusionKVState(kv, "s1", {POS: kv.begin_request("r-pos")}, num_layers=1) + sink, window = 9 * frame, int(kv.spec.sliding_window) + holds: dict[int, int] = {} # slot -> token position last written there + widths: set[int] = set() + filled = compacted = False + + for tick in range(12): + history = int(st.adapter(POS).absolute_num_computed_tokens) + end = history + seq_len + ctx = st.get_kv_caches(POS, seq_len=seq_len, commit_current=True)[0].forward_ctx + ctx.ensure_video_slots(torch.device("cpu")) + for offset, slot in enumerate(ctx.current_video_slot_mapping.tolist()): + holds[slot] = history + offset + table, _, seq_lens, _, _ = ctx.build_block_table(action_len=0, query_len=seq_len, device=torch.device("cpu")) + widths.add(int(table.shape[1])) + kv_len = int(seq_lens[0]) + + position_of_block = {b: i * block for b, i in _model_block_indices(kv, st.adapter(POS)).items()} + read = [] + for b in table[0].tolist()[: -(-kv_len // block)]: + assert b in position_of_block, f"tick {tick}: read block {b}, which holds no token of this session" + read.extend((b * block + o, position_of_block[b] + o) for o in range(block)) + read = read[:kv_len] + wrong = [(slot, want) for slot, want in read if holds.get(slot) != want] + assert not wrong, f"tick {tick}: {len(wrong)} slots read hold a token other than the one their position implies" + + start = end - window + kept = set(range(min(sink, end))) | set(range(max(start, 0), end)) + missed = kept - {want for _, want in read} + assert not missed, f"tick {tick}: {len(missed)} tokens the window keeps were not read" + + filled = filled or end > sink + window + st.commit_paged_context(POS) + compacted = compacted or st.adapter(POS).compacted_tokens > 0 + + assert filled, "the window never filled, so the case that fails was never reached" + assert compacted, "the table was never compacted, so storage and model positions never diverged" + assert len(widths) == 1, f"block table width varied across ticks: {sorted(widths)}" + + +class _CountedBlockIds: + """Wraps a vLLM block manager's ``get_blocks`` and counts the block ids read. + + Every path to a request's block ids goes through ``get_blocks``, so a copy of + the whole table is counted in full even if only a few entries are indexed. + """ + + def __init__(self, manager): + self.reads = 0 + self._get_blocks = manager.get_blocks + + def __call__(self, request_id): + from vllm.v1.core.kv_cache_manager import KVCacheBlocks + + blocks = self._get_blocks(request_id) + return KVCacheBlocks(tuple([_CountedBlock(block, self) for block in group] for group in blocks.blocks)) + + +class _CountedBlock: + def __init__(self, block, counter: _CountedBlockIds): + self._block = block + self._counter = counter + + @property + def block_id(self) -> int: + self._counter.reads += 1 + return self._block.block_id + + +def test_reading_the_window_touches_only_the_blocks_it_can_keep(monkeypatch): + """However long the block table is, reading the window must not cost more. + + Evicted positions stay in the table as null entries until compaction drops + them a whole number of chunks at a time -- up to two frames' worth of blocks + at 832x480 -- so the table can still be far longer than the window. Only the + sink's blocks and the recent window's blocks can hold a kept token, so those + are all the read may touch. Compaction is switched off here so the table + grows to many times the window, and reads are counted where the ids leave + vLLM's block manager. + """ + monkeypatch.setattr(ChunkWindowManager, "compact_block_table", lambda self, request_id: 0) + device = torch.device("cpu") + kv, st = make_state(device=device, window_chunks=2, sink_chunks=1, chunk_size=RAGGED_CHUNK) + reads: list[int] = [] + table_lengths: list[int] = [] + for _ in range(60): + ctx = st.get_kv_caches(POS, seq_len=RAGGED_CHUNK, commit_current=True)[0].forward_ctx + ctx.ensure_video_slots(device) + counted = _CountedBlockIds(kv.manager) + kv.manager.get_blocks = counted + try: + ctx.build_block_table(action_len=0, query_len=RAGGED_CHUNK, device=device) + finally: + del kv.manager.get_blocks + reads.append(counted.reads) + table_lengths.append(len(kv.block_table(st.adapter(POS)))) + st.commit_paged_context(POS) + + assert table_lengths[-1] > 10 * ctx.max_video_blocks, "the table never outgrew the window, so nothing was tested" + assert max(reads) <= ctx.max_video_blocks, ( + f"read {max(reads)} table entries for a {ctx.max_video_blocks}-block window" + ) + + +def test_the_checkpoints_own_default_resolution_can_build_a_cache(): + """832x480 is what LingBot World v2 ships as its default, and it could not run. + + A frame-sized block made the resolution a kernel-compatibility question: + 1560 tokens per frame is not a multiple of 16, so FlashAttention's paged + kernel rejected it and the default resolution had no realtime path at all. + """ + from vllm_omni.experimental.ar_diffusion.runner import paging_block_size + + tokens_per_frame = (480 // 16) * (832 // 16) + assert tokens_per_frame == 1560 + assert tokens_per_frame % 16 == 8, "the whole point is that a frame is not a legal block" + + block_size = paging_block_size(tokens_per_frame) + cfg = ARDiffusionKVConfig(enable=True, chunk_size=tokens_per_frame, window_chunks=2) + kv = ARDiffusionKVCache( + cfg, + num_layers=1, + num_kv_heads=N_HEADS, + head_size=HEAD_DIM, + dtype=torch.float32, + block_size=block_size, + max_model_len=1 << 16, + available_bytes=1 << 28, + kv_branches=(ARDiffusionKVBranchSpec(POS, 0),), + session_capacity=1, + frames_per_block=2, + max_scratch_tokens_per_branch=block_size, + device=torch.device("cpu"), + ) + + assert kv.block_size % 16 == 0 + # The eviction unit is still the frame; only the paging unit changed. + assert kv.spec.chunk_size == tokens_per_frame + assert kv.blocks_per_frame == -(-tokens_per_frame // block_size) == 98 + + +def test_the_shipped_resolution_geometry_matches_dense_attention(): + """The real number, not a scaled-down stand-in: 1560 tokens per frame. + + This is the case the issue reported failing at ~4e-02 on the scratch path. + It runs on CPU because the reference is dense attention, not a kernel. + """ + torch.manual_seed(0) + device = torch.device("cpu") + dtype = torch.float32 + tokens_per_frame = (480 // 16) * (832 // 16) + assert tokens_per_frame == 1560 + + kv, st = make_state(dtype=dtype, device=device, window_chunks=2, chunk_size=tokens_per_frame) + history_k, history_v = _commit_video_span( + kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device, chunk_size=tokens_per_frame + ) + + # The non-committing path is the one that was wrong: four of the five + # forwards per generated block take it. + ctx = st.get_kv_caches(POS, seq_len=tokens_per_frame, commit_current=False)[0].forward_ctx + ctx.ensure_video_slots(device) + assert ctx.start_offset == 8, "1560 % 16 == 8 is the whole reason this case exists" + + current_k = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + current_v = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0] + kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0] + + query = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table( + action_len=0, query_len=tokens_per_frame, device=device + ) + paged = ar_diffusion_paged_attention( + query, + kv.key_cache(0), + kv.value_cache(0), + block_table=block_table, + query_start_loc=query_start_loc, + seq_lens=seq_lens, + max_query_len=max_query_len, + max_seq_len=max_seq_len, + softmax_scale=HEAD_DIM**-0.5, + causal=False, + ) + ref = _dense_attention( + query, + torch.cat([history_k, current_k], dim=1), + torch.cat([history_v, current_v], dim=1), + ) + torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5) + + +@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required") +@pytest.mark.parametrize("commit_current", [False, True]) +def test_the_shipped_resolution_geometry_on_the_real_kernel(commit_current): + """The shipped 832x480 geometry through FlashAttention's paged kernel. + + The CPU test above proves the addressing is right against a dense + reference. This one proves the real kernel agrees, at the real 1560 + tokens per frame, through the production path -- host prep followed by + the fused write+attend op, which is what a forward actually calls. + """ + torch.manual_seed(0) + device = torch.device("cuda") + dtype = torch.float16 + tokens_per_frame = (480 // 16) * (832 // 16) + assert tokens_per_frame == 1560 + + kv, st = make_state(dtype=dtype, device=device, window_chunks=2, chunk_size=tokens_per_frame) + history_k, history_v = _commit_video_span( + kv, st, kv_branch=POS, n_chunks=1, dtype=dtype, device=device, chunk_size=tokens_per_frame + ) + + layer_ctx = st.get_kv_caches(POS, seq_len=tokens_per_frame, commit_current=commit_current)[0] + ctx = layer_ctx.forward_ctx + current_k = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + current_v = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + query = torch.randn(1, tokens_per_frame, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + + ctx.prepare(device=device, action_len=0, query_len=tokens_per_frame) + assert ctx.start_offset == 8, "1560 % 16 == 8 is the whole reason this case exists" + assert kv.block_size == 16 + + inputs = layer_ctx.to_layer_inputs() + paged = paged_write_attn(inputs, query[0], current_k[0], current_v[0], None, None, HEAD_DIM**-0.5).unsqueeze(0) + + new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :] + new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :] + ref = _dense_attention(query, new_k, new_v) + + torch.testing.assert_close(paged, ref, rtol=2e-2, atol=2e-2) + + @pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required") @pytest.mark.parametrize("history_chunks", [0, 1, 3]) @pytest.mark.parametrize("action_len", [0, 3]) diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py index efa5cc2c1..b178981ff 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py @@ -42,6 +42,7 @@ from vllm_omni.experimental.ar_diffusion.kv_cache.paged import ( chunk_slot_mapping, pool_write_chunk, resident_block_ids, + visible_window_blocks, ) _log = init_logger(__name__) @@ -282,7 +283,12 @@ class ARDiffusionKVCache: # forward needs one block per current frame plus space for any # model-declared action/state tokens that coexist with video KV. declared_scratch_blocks = (max_scratch_tokens_per_branch + block_size - 1) // block_size - minimum_scratch_blocks = self.frames_per_block + declared_scratch_blocks + # frames_per_block counts *frames*; convert to blocks rather than + # assuming one block per frame. A frame spans ceil(chunk_size / + # block_size) blocks once the two sizes are allowed to differ. + blocks_per_frame = -(-config.chunk_size // block_size) + self.blocks_per_frame = blocks_per_frame + minimum_scratch_blocks = self.frames_per_block * blocks_per_frame + declared_scratch_blocks override = os.environ.get("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH") override_blocks = int(override) if override is not None else 0 if override_blocks < 0: @@ -313,8 +319,23 @@ class ARDiffusionKVCache: ) def _required_managed_blocks(capacity: int) -> int: - resident_per_session = config.sink_chunks + config.window_chunks - return self.num_local_kv_branches * (capacity * resident_per_session + self.frames_per_block) + 2 + """Managed blocks needed to hold ``capacity`` resident sessions. + + Every term here counts *frames*, so each is converted to blocks. + Before block size and frame size were allowed to differ this + conversion was the identity and the frame counts could be used + directly; they cannot once a frame spans several blocks. + + Converting rather than rewriting keeps the budget the same number + of *tokens* it always was, so a finer block size does not inflate + the requirement -- it only stops the final block of each term from + rounding up a whole frame's worth of memory. + """ + resident_frames = config.sink_chunks + config.window_chunks + per_session_blocks = resident_frames * self.blocks_per_frame + in_flight_blocks = self.frames_per_block * self.blocks_per_frame + headroom_blocks = 2 * self.blocks_per_frame + return self.num_local_kv_branches * (capacity * per_session_blocks + in_flight_blocks) + headroom_blocks # reuse_history_staging keeps one contiguous K and V buffer per layer, # sized to the padded visible window (sink + window plus the @@ -327,9 +348,15 @@ class ARDiffusionKVCache: self.history_staging_tokens = 0 if config.reuse_history_staging: if contiguous_kv_gather_enabled(): - self.history_staging_tokens = ( - self.spec.sliding_window + config.sink_chunks * config.chunk_size + block_size + # The same width build_block_table pads to: a sink or window + # that is not a whole number of blocks spans one more block. + window_blocks = visible_window_blocks( + chunk_size=config.chunk_size, + block_size=block_size, + sink_tokens=config.sink_chunks * config.chunk_size, + window_tokens=self.spec.sliding_window, ) + self.history_staging_tokens = (window_blocks + 1) * block_size self.history_staging_reserved_bytes = int( 2 * num_layers * self.history_staging_tokens * num_kv_heads * head_size * dtype.itemsize ) @@ -649,6 +676,16 @@ class ARDiffusionKVCache: def block_table(self, adapter: ARDiffusionRequestAdapter) -> list[int]: return list(self.manager.get_block_ids(adapter.request_id)[0]) + def block_ids_at(self, adapter: ARDiffusionRequestAdapter, indices: Sequence[int]) -> list[int]: + """Block ids at ``indices`` of the request's block table, without copying it. + + The table keeps a null entry for every evicted position, so it grows with + the session. :meth:`block_table` reads every entry; this reads only the + ones asked for, straight from the manager's own block list. + """ + blocks = self.manager.get_blocks(adapter.request_id).blocks[0] + return [blocks[index].block_id for index in indices] + def chunk_write_slots(self, adapter: ARDiffusionRequestAdapter) -> torch.Tensor: """Slot mapping for the in-flight chunk — the K/V write target.""" return chunk_slot_mapping( diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py index 0ef639c98..669fe7a2f 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py @@ -17,6 +17,7 @@ Cosmos port) reuses unchanged. Three concerns live here: from __future__ import annotations +import math from collections.abc import Sequence from dataclasses import dataclass @@ -196,6 +197,19 @@ def chunk_window_skipped_tokens( return (skipped // chunk_size) * chunk_size +def visible_window_blocks(*, chunk_size: int, block_size: int, sink_tokens: int, window_tokens: int) -> int: + """Most blocks a sink plus a recent window can span; see ``max_video_blocks``. + + Each range is converted on its own: the sink ends at a fixed position, and + the recent range starts on a chunk boundary, so at most + ``block_size - gcd(chunk_size, block_size)`` tokens past a block edge. + """ + worst_start_offset = block_size - math.gcd(chunk_size, block_size) + sink_blocks = -(-sink_tokens // block_size) + recent_blocks = -(-(window_tokens + worst_start_offset) // block_size) + return sink_blocks + recent_blocks + + class ChunkWindowManager(SlidingWindowManager): """``SlidingWindowManager`` that evicts at chunk boundaries. @@ -219,14 +233,26 @@ class ChunkWindowManager(SlidingWindowManager): Call only after commit/eviction, then shift the request's storage position by the returned amount. Physical pages and model positions are unchanged, including any allocated but not yet committed tail. + + The gap starts where :meth:`remove_skipped_blocks` starts freeing: after + the sink's blocks, which are ``sink_chunks`` only while a frame is a + block. Its removal is also rounded down to a whole number of chunks. + Eviction snaps to chunk boundaries and reads its position in storage + coordinates, and paged attention reads each block's offset there too; + shifting storage by whole chunks *and* whole blocks leaves both exactly + as they were in model positions. When a frame is a block, every gap + already is one. """ if self.enable_caching: raise RuntimeError("AR block-table compaction requires prefix caching to be disabled") blocks = self.req_to_blocks.get(request_id, []) - start = self.kv_cache_spec.sink_chunks + spec = self.kv_cache_spec + start = -(-(spec.sink_chunks * spec.chunk_size) // self.block_size) end = start while end < len(blocks) and blocks[end] == self._null_block: end += 1 + blocks_per_step = spec.chunk_size // math.gcd(spec.chunk_size, self.block_size) + end = start + (end - start) // blocks_per_step * blocks_per_step if end == start: return 0 del blocks[start:end] @@ -241,17 +267,32 @@ class ChunkWindowManager(SlidingWindowManager): total_computed_tokens: int, num_prompt_tokens: int | None = None, ) -> None: - """Free the middle gap while preserving the leading attention sink.""" + """Free the middle gap while preserving the leading attention sink. + + ``sink_chunks`` counts frames while the block table is indexed in + blocks, so the sink is converted before it is used as an index. The two + were interchangeable while a frame was a block; this commit is what + separates them. + + Neither boundary is guaranteed to land on a block edge -- 832x480 is + 1560 tokens per frame against a 16-token block -- so a block can hold + the end of one frame and the start of the next. The sink end rounds up + so a straddling block stays protected, and the freed range ends on a + rounded-down boundary so only wholly-skipped blocks are released. Both + are exact when the sizes divide, which is every geometry that paged one + frame per block before. + """ del num_prompt_tokens num_skipped_tokens = self.get_num_skipped_tokens(total_computed_tokens) if num_skipped_tokens <= 0: return - sink_blocks = self.kv_cache_spec.sink_chunks - num_skipped_blocks = num_skipped_tokens // self.block_size + sink_tokens = self.kv_cache_spec.sink_chunks * self.kv_cache_spec.chunk_size + first_live_block = -(-sink_tokens // self.block_size) + last_freed_block = (sink_tokens + num_skipped_tokens) // self.block_size self._remove_blocks_in_range( request_id, - sink_blocks, - sink_blocks + num_skipped_blocks, + first_live_block, + last_freed_block, ) diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py index 2e7f72522..0ea57e39f 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py @@ -6,12 +6,15 @@ from __future__ import annotations import weakref from dataclasses import dataclass, field -from typing import Any, ClassVar, NamedTuple +from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple import torch from vllm_omni.experimental.ar_diffusion.kv_cache.config import contiguous_kv_gather_enabled -from vllm_omni.experimental.ar_diffusion.kv_cache.paged import compute_slot_mapping +from vllm_omni.experimental.ar_diffusion.kv_cache.paged import compute_slot_mapping, visible_window_blocks + +if TYPE_CHECKING: + from vllm.v1.attention.backend import MultipleOf _LAYER_IDX_TENSORS: dict[int, torch.Tensor] = {} @@ -91,6 +94,21 @@ class ARDiffusionPagedForwardContext: action_slot_mapping: torch.Tensor | None = None query_len: int = 0 kv_len: int = 0 + # Tokens already committed when this forward began. Recorded rather than + # derived from the block count: the last block of the history is only + # partly written when a chunk is not a whole number of blocks, so + # blocks * block_size overstates it and would make attention read past + # what was written. + _history_tokens: int = 0 + # Scratch blocks the current video tokens occupy. Not the same as + # len(current_video_block_ids): when the history ends mid-block, the first + # entry there is the history's own tail block, which is managed, not + # scratch. Action K/V is placed after this, so it must not count it. + _scratch_blocks_used: int = 0 + # Slot offset of this chunk's first token inside its first block. The + # chunk's blocks then hold positions ``history - offset + j * block_size``, + # which is how the read path places them among the history's blocks. + _current_offset: int = 0 _allocated_video: bool = False _committed: bool = False _action_len: int = 0 @@ -106,14 +124,77 @@ class ARDiffusionPagedForwardContext: def block_size(self) -> int: return int(self.kv_cache.block_size) + @property + def chunk_size(self) -> int: + """Tokens in one chunk -- one latent frame. The eviction unit.""" + return int(self.kv_cache.spec.chunk_size) + + @property + def start_offset(self) -> int: + """Where in its first block this forward's tokens begin. + + Non-zero exactly when the committed history does not end on a block + boundary, which is the normal case once a frame is not a whole number + of blocks: 1560 tokens per frame against 16-token blocks leaves the + history 8 slots into its last block on every odd-numbered chunk. + """ + return int(self.adapter.num_computed_tokens) % self.block_size + + @property + def sink_tokens(self) -> int: + """Leading tokens the window keeps however far it has slid.""" + return int(self.kv_cache.spec.sink_chunks) * self.chunk_size + + @property + def window_tokens(self) -> int: + """Most recent tokens the window keeps, this forward's own included.""" + return self.max_video_tokens - self.sink_tokens + + @property + def max_video_blocks(self) -> int: + """Most blocks the visible window can span, which fixes the table width. + + The window is two token ranges, the sink and the most recent tokens, and + each is converted to blocks on its own because each can straddle a block + edge on its own. Rounding the two ranges' *sum* up once was a block short + whenever both straddled: at 832x480 a 9-frame sink and a 9-frame window + are each 8 tokens past a 16-token edge but 28080 tokens together, a whole + number of blocks, so it was short on every tick once the window had + filled. + + The sink ends at a fixed position, so its blocks are just its tokens + rounded up. The recent range starts at a whole number of chunks, so its + start sits at most ``block_size - gcd(chunk_size, block_size)`` tokens + past a block edge, and it can need that much more room. A chunk that is + a whole number of blocks never straddles, so every geometry that paged + one frame per block keeps exactly the width it had. + + The width is a constant of the configuration, which is what keeps the + block table a fixed shape as the window fills. + """ + return visible_window_blocks( + chunk_size=self.chunk_size, + block_size=self.block_size, + sink_tokens=self.sink_tokens, + window_tokens=self.window_tokens, + ) + @property def num_current_video_blocks(self) -> int: - if self.seq_len % self.block_size != 0: - raise AssertionError( - "AR-Diffusion paged attention expects frame-aligned seq_len " - f"(multiple of block_size={self.block_size}), got {self.seq_len}" - ) - return self.seq_len // self.block_size + """Blocks this forward's tokens occupy, counted from where they start. + + A chunk need not be a whole number of blocks, so the count depends on + the offset it begins at, not only on its length: thirty tokens + starting at position thirty span three sixteen-token blocks, not two. + Rounding up leaves the tail of the last block unwritten, which is only + safe because nothing reads past ``kv_len`` -- see + :meth:`video_block_table`. + + Both paths count the same way. The committing path writes straight + after the history; the scratch path is made to line up with it by + :meth:`ensure_video_slots`. + """ + return -(-(self.start_offset + self.seq_len) // self.block_size) def ensure_video_slots(self, device: torch.device) -> None: """Allocate/write targets for the current video tokens, once per KV branch.""" @@ -121,8 +202,10 @@ class ARDiffusionPagedForwardContext: return n_blocks = self.num_current_video_blocks + self._history_tokens = int(self.adapter.num_computed_tokens) if self.commit_current: start = int(self.adapter.num_computed_tokens) + self._current_offset = start % self.block_size self.kv_cache.allocate_token_slots(self.adapter, self.seq_len) table = self.kv_cache.block_table(self.adapter) start_block = start // self.block_size @@ -132,8 +215,27 @@ class ARDiffusionPagedForwardContext: compute_slot_mapping(table, positions, self.block_size), device ) else: - self.current_video_block_ids = self.kv_cache.scratch_block_ids(self.kv_branch, 0, n_blocks) - positions = torch.arange(self.seq_len, dtype=torch.long) + # The scratch region cannot simply start at slot zero. The kernel + # reads the history and this chunk as one contiguous run, so if the + # history stopped mid-block, starting here at zero would leave the + # rest of that block unwritten and the kernel would read those dead + # slots as if they were tokens, shifting the whole sequence. + # + # Instead the chunk begins where the history left off, spilling its + # first few tokens into the free tail of the history's own last + # block. That block is already this session's -- blocks are + # allocated whole -- and those slots hold nothing yet; the + # committing forward overwrites the same slots with the clean K/V + # later, so the committed state is unchanged either way. + # A history whose tail block is no longer resident has nothing to + # spill into; the window then starts at this chunk and zero is right. + offset = self.start_offset if self.history_block_ids else 0 + self._current_offset = offset + self._scratch_blocks_used = n_blocks - (1 if offset else 0) + scratch_ids = self.kv_cache.scratch_block_ids(self.kv_branch, 0, self._scratch_blocks_used) + tail_block = [self.history_block_ids[-1]] if offset else [] + self.current_video_block_ids = tail_block + scratch_ids + positions = torch.arange(offset, offset + self.seq_len, dtype=torch.long) self.current_video_slot_mapping = _to_device_async( compute_slot_mapping(self.current_video_block_ids, positions, self.block_size), device ) @@ -152,7 +254,9 @@ class ARDiffusionPagedForwardContext: return action_blocks = (action_len + self.block_size - 1) // self.block_size - scratch_offset = 0 if self.commit_current else len(self.current_video_block_ids) + # Count scratch blocks, not entries: the video list may lead with the + # history's tail block, which lives in the managed pool. + scratch_offset = 0 if self.commit_current else self._scratch_blocks_used self.action_scratch_block_ids = self.kv_cache.scratch_block_ids( self.kv_branch, scratch_offset, @@ -165,23 +269,75 @@ class ARDiffusionPagedForwardContext: self._action_len = action_len def video_block_table(self, device: torch.device) -> tuple[list[int], int]: + """Blocks the attention reads, and how many of their slots it reads. + + The kernel reads the listed blocks as one contiguous run and stops after + the returned length, so two things must hold: every block holding a + token the window keeps is listed, in token order, and only the last + listed block may be partly written. + + Blocks are chosen by the token positions they hold rather than counted + off the ends of the resident list. The window is the sink ``[0, S)`` + plus the most recent tokens, and each of those can straddle a block + edge on its own. Taking a fixed number of blocks off the tail dropped + the block holding the window's first tokens exactly when both + straddled, and a length capped by a block count rather than counted + from what was written then ran past the last written slot. + + Whole blocks are read, so each boundary block contributes up to + ``block_size - 1`` tokens the window has already let go of. They are + real K/V at their own positions; a slightly wider window is a far + smaller error than a truncated one. + """ self.ensure_video_slots(device) - if self.max_video_tokens % self.block_size != 0: - raise AssertionError( - "AR-Diffusion paged attention requires max_video_tokens to be block-aligned, " - f"got max_video_tokens={self.max_video_tokens}, block_size={self.block_size}" - ) - all_video_blocks = self.history_block_ids + self.current_video_block_ids - max_video_blocks = self.max_video_tokens // self.block_size - sink_blocks = int(self.kv_cache.spec.sink_chunks) - if len(all_video_blocks) <= max_video_blocks: - visible_video_blocks = all_video_blocks - else: - tail_blocks = max_video_blocks - sink_blocks - visible_video_blocks = all_video_blocks[:sink_blocks] - if tail_blocks: - visible_video_blocks += all_video_blocks[-tail_blocks:] - video_len = len(visible_video_blocks) * self.block_size + block_size = self.block_size + history = self._history_tokens + end = history + self.seq_len + sink_end = self.sink_tokens + recent_start = end - self.window_tokens + + # Position of each block's first slot. History comes from the block + # table, where an evicted block is a null placeholder, so a block's index + # is its position -- in storage positions, which is what + # num_computed_tokens counts once compaction has dropped an evicted gap. + # Compaction shifts storage by whole chunks and whole blocks, so the sink, + # the window and every block offset sit exactly where they would in model + # positions. Reading the table after ensure_video_slots means a block + # this chunk's own allocation evicted is never listed. The current + # chunk's blocks follow from where its first token was placed; its first + # block is the history's own tail block whenever the history stopped + # part way through one, and keeps the history's position. + first_position: dict[int, int] = {} + # + # Only two index ranges can hold a token the window keeps: the sink's + # blocks, and the blocks from the recent window's start onward. Evicted + # positions stay in the table as null entries until compaction drops + # them a whole number of chunks at a time, so the table can still be + # much longer than the window; reading just these two ranges, without + # copying the table, keeps this at the window's size. + history_blocks = -(-history // block_size) + sink_blocks = min(-(-sink_end // block_size), history_blocks) + recent_first_block = max(max(recent_start, 0) // block_size, sink_blocks) + indices = [*range(sink_blocks), *range(recent_first_block, history_blocks)] + for index, block in zip(indices, self.kv_cache.block_ids_at(self.adapter, indices)): + if block != self.kv_cache.null_block_id: + first_position.setdefault(block, index * block_size) + current_start = history - self._current_offset + for index, block in enumerate(self.current_video_block_ids): + first_position.setdefault(int(block), current_start + index * block_size) + + visible_video_blocks: list[int] = [] + video_len = 0 + for block, first in sorted(first_position.items(), key=lambda item: item[1]): + if first >= sink_end and first + block_size <= recent_start: + continue # wholly between the sink and the recent window + if video_len % block_size: + raise RuntimeError( + "AR-Diffusion paged attention would read past a partly written block: " + f"block {visible_video_blocks[-1]} is followed by block {block}." + ) + visible_video_blocks.append(block) + video_len += min(block_size, end - first) return visible_video_blocks, video_len def build_block_table( @@ -200,6 +356,21 @@ class ARDiffusionPagedForwardContext: ``ceil(seq_lens/block_size)`` entries, so padding is never read. """ video_blocks, video_len = self.video_block_table(device) + # Action K/V is listed after the video blocks and the kernel reads the + # table as one run, so the video run has to end on a block edge. If it + # stops part way through its last block, the run reads that block's + # unwritten slots as the first action tokens and never reaches the last + # ones -- a wrong attention with every shape intact. A frame that is a + # whole number of blocks always ends on an edge, so this refuses only + # geometries that could not page at all before the paging unit was + # separated from the frame. + if action_len > 0 and video_len % self.block_size: + raise ValueError( + "AR-Diffusion paged attention cannot place action tokens after a partly written video block: " + f"the video run ends {video_len % self.block_size} slots into block {video_blocks[-1]}, so the " + "kernel would read its unwritten slots as action tokens. Action tokens need tokens_per_frame " + f"to be a whole number of blocks (block_size={self.block_size})." + ) self.ensure_action_slots(action_len, device) action_blocks = self.action_scratch_block_ids if action_len > 0 else [] block_ids = video_blocks + action_blocks @@ -208,12 +379,19 @@ class ARDiffusionPagedForwardContext: # Fixed capacity: full visible video window + one action-capacity block. action_capacity_blocks = max(1, (action_len + self.block_size - 1) // self.block_size) - width = max(self.max_video_tokens // self.block_size + action_capacity_blocks, len(block_ids)) + # Both of these count in blocks. Deriving them from the token count by + # flooring would understate the capacity whenever the window is not a + # whole number of blocks: the width would fall back on len(block_ids) + # and start tracking how full the window is -- the exact shape churn + # the fixed width exists to prevent -- and max_seq_len could come out + # below the kv_len actually being passed, by up to block_size - 1. + capacity_blocks = self.max_video_blocks + action_capacity_blocks + width = max(capacity_blocks, len(block_ids)) padded = block_ids + [0] * (width - len(block_ids)) self.query_len = int(query_len) self.kv_len = int(video_len + action_len) - max_seq_len = int(self.max_video_tokens + action_capacity_blocks * self.block_size) + max_seq_len = int(capacity_blocks * self.block_size) # Built on the host once per AR block; the copies go through the same # pinned, non-blocking path as the slot mappings above so the CPU does # not wait for the stream to drain three times before the first layer. @@ -281,7 +459,11 @@ class ARDiffusionPagedForwardContext: int(self.seq_len), ) same_session = state.adapter is not None and state.adapter() is self.adapter - reuse = same_session and state.signature == signature + # A chunk that is not a whole number of blocks shares its first block with the history's tail (see + # ensure_video_slots), and the attention op restages the current chunk in whole blocks of its own, + # so such a forward restages the whole window instead. + block_aligned = self.start_offset == 0 and self.seq_len % self.block_size == 0 + reuse = same_session and state.signature == signature and block_aligned state.adapter = weakref.ref(self.adapter) state.signature = signature self.staging_enabled = True @@ -461,6 +643,36 @@ def _resolve_fa_version(head_size: int) -> int: return version +def supported_kernel_block_sizes() -> list[int | MultipleOf]: + """Block sizes the kernel this module dispatches to will accept. + + Same shape as vLLM's ``AttentionBackend.get_supported_kernel_block_sizes`` + -- a plain int is that exact size, ``MultipleOf(b)`` is any positive + multiple of ``b`` -- but answered here rather than read off a backend, + because AR-Diffusion does not go through backend selection. It calls + ``flash_attn_varlen_func`` itself, choosing between vLLM's CUDA build and + ROCm's AITER a few lines below, and only this module knows which. + + It lives next to that choice so there is one place to change. The + constraint is a property of the kernel, not of the card, and the kernels + reachable from here agree on 16 today: vLLM's CUDA FlashAttention, ROCm + AITER, and upstream ``flash_attn`` all advertise ``MultipleOf(16)``. Other + backends in the same tree do not -- ``hpc_attn`` accepts only 64, and + FlashInfer advertises pages of 128 or more solely on Blackwell -- so a + caller must treat this as data to be queried, never as the number 16. + + vLLM's FlashAttention backend does advertise a single 128-token page when FA4 + runs its dedicated head-size-256 kernel on SM100/SM110. That kernel is not + reachable from here: ``get_flash_attn_version`` selects it only when the + caller passes ``supports_fa4_hd256=True``, and :func:`_resolve_fa_version` + does not, so a head size of 256 falls back to FA2, which pages at any + multiple of 16. Passing that flag later has to change this answer as well. + """ + from vllm.v1.attention.backend import MultipleOf + + return [MultipleOf(16)] + + def _rocm_flash_attn_varlen_func(): """Resolve a ROCm varlen kernel, preferring AITER when available. diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/state.py b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py index ec14a9177..670046d0b 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/state.py +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Model-facing session state backed by paged AR-Diffusion KV storage.""" from __future__ import annotations diff --git a/vllm_omni/experimental/ar_diffusion/runner.py b/vllm_omni/experimental/ar_diffusion/runner.py index 858cf843d..4180ae587 100644 --- a/vllm_omni/experimental/ar_diffusion/runner.py +++ b/vllm_omni/experimental/ar_diffusion/runner.py @@ -5,8 +5,9 @@ from __future__ import annotations import dataclasses import time from collections import OrderedDict -from collections.abc import Iterator +from collections.abc import Iterator, Sequence from contextlib import contextmanager +from typing import TYPE_CHECKING import torch from vllm.logger import init_logger @@ -29,9 +30,61 @@ from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVStat from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionTickRequest from vllm_omni.platforms import current_omni_platform +if TYPE_CHECKING: + from vllm.v1.attention.backend import MultipleOf + logger = init_logger(__name__) +def paging_block_size(tokens_per_frame: int, supported: Sequence[int | MultipleOf] | None = None) -> int: + """Choose the paging unit, which is not the eviction unit. + + AR-Diffusion evicts by frame; the attention kernel tiles by a fixed number + of tokens. Using the frame as the block binds every output resolution to + that kernel constraint, so the resolutions that run today do so by + arithmetic accident -- and the LingBot World v2 checkpoint's own default + 832x480 is not one of them: + + (480/16) x (832/16) = 30 x 52 = 1560 tokens per frame, 1560 % 16 = 8 + + A frame that is already a legal block stays one, so every resolution that + works today keeps precisely the paging it has now. Anything else pages at + the finest legal unit, which costs a longer block table and gives back the + rounding waste of a block one whole frame wide -- at 832x480 that block is + over a gigabyte. + + What counts as legal is asked, not assumed. ``supported`` carries vLLM's + own vocabulary: a plain int is that exact size, ``MultipleOf(b)`` is any + positive multiple of ``b``. The kernels AR-Diffusion reaches agree on + multiples of 16, but others in the same tree do not -- ``hpc_attn`` takes + only 64, and FlashInfer offers pages of 128 and up on Blackwell alone. A + constant here would be a silently wrong block size on the first machine + whose kernel differs, and a wrong block size is not an error: it is a + kernel reading tokens that were never written. + """ + if supported is None: + # Imported here, not at module scope: paged_attention registers a + # custom op on import, and the tests that pop it from sys.modules to + # re-import it depend on nothing else holding it first. + from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import supported_kernel_block_sizes + + supported = supported_kernel_block_sizes() + if not supported: + raise ValueError("The attention kernel advertises no legal block size.") + + def legal(size: int) -> bool: + return size > 0 and any( + size == entry if isinstance(entry, int) else size % entry.base == 0 for entry in supported + ) + + if legal(tokens_per_frame): + return tokens_per_frame + finest = min(entry if isinstance(entry, int) else entry.base for entry in supported) + if not legal(finest): + raise ValueError(f"No legal paging unit among {supported!r}.") + return finest + + def resolve_ar_diffusion_kv_config(od_config: OmniDiffusionConfig) -> ARDiffusionKVConfig: """Resolve optional deployment overrides and enable AR-Diffusion KV.""" raw = getattr(od_config, "ar_diffusion_kv_config", None) @@ -139,7 +192,7 @@ class ARDiffusionModelRunner(DiffusionModelRunner): num_kv_heads=spec.num_kv_heads, head_size=spec.head_size, dtype=self.od_config.dtype, - block_size=spec.tokens_per_frame, + block_size=paging_block_size(spec.tokens_per_frame), max_model_len=spec.max_model_len, available_bytes=self._available_memory_bytes() if available_bytes is None else available_bytes, kv_branches=spec.kv_branches, From ae58c2d8b9e7ce0044054a817b695aadb1924342 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Fri, 2 Oct 2026 05:17:06 -0700 Subject: [PATCH 13/22] [Model] YuE2: async single-stage synthesis and CUDA graph acceleration (#8407) Signed-off-by: Sy03 <1370724210@qq.com> Signed-off-by: princepride Signed-off-by: Wang Zhipeng Co-authored-by: princepride Co-authored-by: Claude Opus 5.5 --- examples/offline_inference/yue2/end2end.py | 15 +- recipes/m-a-p/YuE2-3B.md | 61 +- tests/e2e/features/yue2/worker_extension.py | 133 ++ tests/e2e/online_serving/test_yue2.py | 223 ++- tests/worker/test_gpu_ar_model_runner.py | 61 + tests/worker/test_yue2_request_state.py | 215 +++ tests/worker/test_yue2_sampler.py | 168 ++ vllm_omni/deploy/yue2.yaml | 32 +- .../entrypoints/openai/tts_adapters/yue2.py | 13 +- vllm_omni/model_executor/models/yue2/yue2.py | 1509 ++++++++++++++--- vllm_omni/tokenizers/yue2_prompt.py | 5 + vllm_omni/worker/gpu_ar_model_runner.py | 43 +- 12 files changed, 2182 insertions(+), 296 deletions(-) create mode 100644 tests/e2e/features/yue2/worker_extension.py create mode 100644 tests/worker/test_yue2_request_state.py create mode 100644 tests/worker/test_yue2_sampler.py diff --git a/examples/offline_inference/yue2/end2end.py b/examples/offline_inference/yue2/end2end.py index 39fddd371..4d04a5633 100644 --- a/examples/offline_inference/yue2/end2end.py +++ b/examples/offline_inference/yue2/end2end.py @@ -28,7 +28,9 @@ import torch from vllm_omni import Omni from vllm_omni.model_executor.models.yue2.yue2 import ( ABC_SAMPLING, + CONTEXT, KEY_MAX_AUDIO_FRAMES, + KEY_MAX_HOLD_STEPS, KEY_MIN_TOKENS, KEY_PENALTY_WINDOW, KEY_PHASE, @@ -41,10 +43,12 @@ from vllm_omni.model_executor.models.yue2.yue2 import ( KEY_TOP_P, SEMANTIC_SAMPLING, STOP_TOKEN_IDS, + SYNTHESIS_HOLD_STEPS, ) from vllm_omni.tokenizers.yue2_prompt import ( abc_ids_from_generated, abc_prefix_ids, + semantic_frames, semantic_prefix_ids, ) from vllm_omni.tokenizers.yue2_tokenizer import YuE2TextTokenizer @@ -71,7 +75,13 @@ def sampling_params(engine, *, phase, seed, max_frames, prompt_ids): # conditioning prefix from the scheduled tokens alone. KEY_PREFIX_IDS: list(prompt_ids), } - params.max_tokens = preset["max_tokens"] if phase == "abc" else max_frames + 1 + if phase == "abc": + params.max_tokens = preset["max_tokens"] + else: + # The song is synthesized on a side stream while the request is held; + # leave max_tokens headroom for the hold tokens. + params.max_tokens = min(max_frames + 1 + SYNTHESIS_HOLD_STEPS, CONTEXT - len(prompt_ids)) + params.extra_args[KEY_MAX_HOLD_STEPS] = max(0, params.max_tokens - max_frames - 1) params.stop_token_ids = list(STOP_TOKEN_IDS) params.detokenize = False return params @@ -144,6 +154,9 @@ def main() -> None: payload = { "prompt_token_ids": prompt_ids, "generated_token_ids": generated_ids, + # Codec frames only: without the end token and the hold tokens + # the request emits while its song is synthesized. + "semantic_frames": semantic_frames(generated_ids), "seed": args.seed, "cot": args.cot, "max_frames": args.max_frames, diff --git a/recipes/m-a-p/YuE2-3B.md b/recipes/m-a-p/YuE2-3B.md index 8ad3b008d..197156e66 100644 --- a/recipes/m-a-p/YuE2-3B.md +++ b/recipes/m-a-p/YuE2-3B.md @@ -40,9 +40,11 @@ user-downloaded weights and bundles none. | Output | 48 kHz stereo, whole song delivered when the semantic request finishes | Generation is one or two requests on a single AR stage: the abc phase writes -the score, the semantic phase generates codec frames and, on its last step, -solves the ODE and decodes the whole song in-engine (weights load once; the -VAE loads at startup from `$YUE2_VAE` or the hub id). +the score, the semantic phase generates codec frames, then queues the ODE and VAE on a +high-priority CUDA stream. Other requests keep decoding while the finishing +request emits HOLD tokens (up to 4096 steps, within its token/context budget). +The scheduler is unchanged. Weights load once; the VAE loads at startup from +`$YUE2_VAE` or the hub id. ## Running @@ -67,15 +69,50 @@ ending. ## Notes -- **Memory (RTX 4090, 24 GB, `gpu_memory_utilization: 0.70`):** the engine - reserves ~15.8 GiB after startup; with 4 concurrent requests pinned to the - 9000-frame cap (worst case), peak usage during the terminal NAR/VAE - finishing pass reaches ~20.9 GiB, leaving ~3 GiB of headroom. At 0.85 the - same workload OOMs inside the finishing pass, so the deploy yaml pins 0.70. -- **Known limitations:** 4 concurrent full-length requests are the verified - shape (`max_num_seqs: 4`); CUDA graph capture is a - follow-up; the abc phase runs eagerly - after prefill (its tokens are the product, not audio). +The lifecycle e2e regression forces KV-cache preemption with chunked and +unchunked recompute, mixes ABC and semantic sampling, and cancels a request +after NAR work has started. It checks history alignment, released synthesis +buffers, and successful generation after cancellation. Run with one H100/H200 +and cached YuE2-3B and YuE2-Vae weights: + +```bash +export YUE2_MODEL_DIR=/path/to/YuE2-3B +export YUE2_VAE=/path/to/YuE2-Vae +CUDA_VISIBLE_DEVICES=0 python -m pytest tests/e2e/online_serving/test_yue2.py \ + -k preemption_and_synthesis_abort -m 'slow and tts' --run-level full_model -q -s +``` + +The test retains the module's weekly TTS routing; it is not a per-PR CI gate. + +- **Memory:** `gpu_memory_utilization` budgets the vLLM engine; NAR K/V, + acoustic graph buffers and VAE activations also need room during synthesis. + The 0.70 / 4-slot defaults require validation on the target card. The earlier + RTX 4090 measurement (15.8 GiB after startup, 20.9 GiB peak at 9000 frames) + predates the async/compiled implementation; it does not validate this + version or long ABC prefixes on a 4090. Reduce the fraction if the target + workload needs more synthesis memory, especially for long ABC prefixes and + the 9000-frame cap. Releasing completed chunks bounds live NAR buffers, + while the shared graph allocator retains reserved storage for reuse. +- **CUDA graphs:** the default deploy keeps the AR backbone eager. Both ABC + and semantic sampling use prewarmed power-of-two graph buckets. Custom + sampler captures have a bounded cache and fall back to eager sampling when + it fills. The NAR velocity pass uses a chunk graph on the serialized synthesis + stream. Chunk graphs share one allocator pool; completed chunks release + their engines and buffers after their own events. FA3 runs on Hopper; other + CUDA cards use SDPA, which is also warmed at startup. +- **Serving:** run `vllm serve m-a-p/YuE2-3B --omni` with the default + `yue2.yaml`, which uses eager AR with 4 slots. Validate memory capacity + for the target card and workload. + For the tested single-H200 graph settings, reuse this deploy with + `--stage-overrides '{"0":{"max_num_seqs":32,"gpu_memory_utilization":0.5,"enforce_eager":false,"compilation_config":{"cudagraph_mode":"FULL_AND_PIECEWISE"}}}'`. + Decode uses FULL and mixed prefill uses PIECEWISE; VAE and the Python + synthesis queue remain outside these graphs. + Audio is delivered as a whole song; time to first audio equals completion + latency. Aborting a running song stops future work units and waits only for + the at-most-two submitted units before releasing its buffers. +- **Preemption:** before drawing for a resumed row, model-owned sampling + reconciles the codec history and terminal flags with the engine's retained + sequence. Discarded draws cannot advance its penalty window or frame budget. - Audio is not expected to match the upstream torch reference bit-for-bit (fused vs eager kernels). Measured against the upstream torch reference (bd90e4c, same style/lyrics/seed 831001, cot=off, 200 frames): prompt token diff --git a/tests/e2e/features/yue2/worker_extension.py b/tests/e2e/features/yue2/worker_extension.py new file mode 100644 index 000000000..ed3557560 --- /dev/null +++ b/tests/e2e/features/yue2/worker_extension.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Observe YuE2 lifecycle contracts inside a real engine worker.""" + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner + + +class Yue2LifecycleWorkerExtension: + model_runner: "GPUARModelRunner" + + def start_yue2_probe(self): + from unittest.mock import patch + + from vllm.compilation.cuda_graph import CUDAGraphWrapper + from vllm.config import CUDAGraphMode + from vllm.forward_context import get_forward_context, is_forward_context_available + + from vllm_omni.model_executor.models.yue2.yue2 import ABC_END, HOLD_TOKEN, MUSIC_END + + model = self.model_runner.get_model() + self.yue2_probe = { + "preemptions": 0, + "rollbacks": 0, + "checked_rows": 0, + "mixed_steps": 0, + "synthesis_started": False, + "ar_full_replays": 0, + "ar_piecewise_replays": 0, + } + graph_call = CUDAGraphWrapper.__call__ + + def count_ar_replays(wrapper, *args, **kwargs): + if is_forward_context_available(): + context = get_forward_context() + mode = context.cudagraph_runtime_mode + entry = wrapper.concrete_cudagraph_entries.get(context.batch_descriptor) + if mode == wrapper.runtime_mode and entry is not None and entry.cudagraph is not None: + if mode == CUDAGraphMode.FULL: + self.yue2_probe["ar_full_replays"] += 1 + elif mode == CUDAGraphMode.PIECEWISE: + self.yue2_probe["ar_piecewise_replays"] += 1 + return graph_call(wrapper, *args, **kwargs) + + # This worker is created for this test and exits with the engine. + # Count vLLM AR replays, separately from model-owned sampler/NAR graphs. + self.yue2_graph_patch = patch.object(CUDAGraphWrapper, "__call__", count_ar_replays) + self.yue2_graph_patch.start() + self.yue2_cancel_job = None + previous: dict[str, int] = {} + prepare = model.prepare_runner_inputs + reconcile = model._reconcile_history + queue = model._queue_synthesis + sample = model.sample + resolve = model._resolve_pending + emitted: dict[str, int] = {} + + def resolve_pending(): + if model._pending is not None: + event, host, states = model._pending + event.synchronize() + for state, token in zip(states, host[0].tolist()): + end = ABC_END if state.constants.phase == "abc" else MUSIC_END + expected = HOLD_TOKEN if token == end and not state.constants.skip_synthesis else token + assert emitted[state.request_id] == expected, "engine token and model history disagree" + resolve() + + def sample_tokens(logits, sampling_metadata): + result = sample(logits, sampling_metadata) + phases = {model._states[rid].constants.phase for rid, _, _ in model._step_rows if rid in model._states} + if len(phases) == 2: + self.yue2_probe["mixed_steps"] += 1 + for (rid, _, _), row in zip(model._step_rows, result.sampled_token_ids.tolist()): + emitted[rid] = row[0] + return result + + def prepare_inputs(**kwargs): + for rid, computed in zip(kwargs["req_ids"], kwargs["num_computed_tokens"]): + if int(computed) < previous.get(rid, 0): + self.yue2_probe["preemptions"] += 1 + previous[rid] = int(computed) + return prepare(**kwargs) + + def reconcile_history(state, accepted): + if len(state.history) + int(state.end_drawn) > accepted: + self.yue2_probe["rollbacks"] += 1 + reconcile(state, accepted) + assert len(state.history) + int(state.end_drawn) <= accepted + if not state.finish_ready and state.job is None and not state.finished: + assert len(state.history) == accepted + assert not state.truncated and not state.end_drawn + self.yue2_probe["checked_rows"] += 1 + + def queue_synthesis(state): + queue(state) + if "cancel-synthesis" in state.request_id: + job = state.job + assert job.started and job.work is not None and job.keep + assert job.event is None, "test must intercept an active synthesis job" + # Keep an actual NAR job active until the driver cancels it. + # Its first work units ran normally; stop pumping later units + # so cancellation cannot race against a short song finishing. + self.yue2_cancel_advance = job.advance + job.advance = lambda: None + self.yue2_cancel_job = job + self.yue2_probe["synthesis_started"] = True + + model.prepare_runner_inputs = prepare_inputs + model._reconcile_history = reconcile_history + model._queue_synthesis = queue_synthesis + model.sample = sample_tokens + model._resolve_pending = resolve_pending + return True + + def get_yue2_probe(self): + result = dict(self.yue2_probe) + job = self.yue2_cancel_job + if job is not None: + if job.cancelled: + # Exercise the original pump after cleanup: it must not submit + # further units even though this test retains the job object. + self.yue2_cancel_advance() + model = self.model_runner.get_model() + result["cancelled"] = job.cancelled + result["released"] = job._released and job.work is None and not job.keep and job.stream is None + result["removed"] = ( + job.request_id not in model._states + and model._synthesis.active is not job + and job not in model._synthesis.waiting + ) + return result diff --git a/tests/e2e/online_serving/test_yue2.py b/tests/e2e/online_serving/test_yue2.py index 35b25b1b0..17b020517 100644 --- a/tests/e2e/online_serving/test_yue2.py +++ b/tests/e2e/online_serving/test_yue2.py @@ -32,9 +32,9 @@ above the byte floor) and passes on that alone, following the MiniMax Music 3 precedent. Weekly rather than per-PR: each request is a 16 s song and the verification -checkpoints add ~5.5 GB of downloads. Validated on a single RTX 4090 -(24 GB); the hardware mark targets an H100 because the CI fleet has no -consumer cards. +checkpoints add ~5.5 GB of downloads. The async/graph implementation was +validated on H200; the hardware mark targets the CI fleet's H100. The older +RTX 4090 result does not establish this implementation's memory requirements. """ import concurrent.futures @@ -60,8 +60,7 @@ DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 900.0 # 400 frames = 16 s of song. The byte floor sits far below the observed # payload (~3.0 MB of 48 kHz stereo 16-bit) to leave room for an early EOS. _FRAMES = 400 -# Concurrent pair stays short: it exists to exercise row re-indexing, not -# to double the wall time. +# Short concurrent requests exercise row re-indexing as budgets differ. _CONCURRENT_FRAMES = 250 _MIN_BYTES = 1_000_000 @@ -173,7 +172,7 @@ def test_yue2_metal_twinkle_001(omni_server, openai_client, tmp_path) -> None: # CI runs without the verification assets, so the structural checks below # are the only guard against the two real bugs this PR fixed during review: # mono/half-speed output (WAV header + duration bounds) and the int32 - # crash under concurrent row re-indexing (two simultaneous requests). + # crash under concurrent row re-indexing (simultaneous requests). # Duration bounds: the end token is masked for the first 200 steps # (min_tokens), and the 400-frame budget caps the song, so a healthy # request lands in [8, 16.5] s; the half-speed bug produced ~32 s. @@ -190,13 +189,15 @@ def test_yue2_metal_twinkle_001(omni_server, openai_client, tmp_path) -> None: # (see _http_speech); the OpenAI SDK response does not expose headers. _http_speech(omni_server, seed=SEED, frames=_FRAMES) - # Two concurrent requests with different seeds: both must succeed and - # produce different songs (the model-owned sampler re-indexes rows across - # concurrent requests; a dtype mismatch here crashed the engine before). - with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: - futures = [pool.submit(_http_speech, omni_server, seed=SEED + i, frames=_CONCURRENT_FRAMES) for i in range(2)] - pair = [f.result() for f in futures] - assert pair[0] != pair[1], "different seeds produced identical audio" + # Different budgets exercise row re-indexing as shorter requests finish. + # Compare the same-budget pair so length cannot hide an ignored seed. + with concurrent.futures.ThreadPoolExecutor(max_workers=3) as pool: + futures = [ + pool.submit(_http_speech, omni_server, seed=SEED + i, frames=frames) + for i, frames in enumerate((_CONCURRENT_FRAMES, _FRAMES, _CONCURRENT_FRAMES)) + ] + audios = [f.result() for f in futures] + assert audios[0] != audios[2], "different seeds produced identical audio at the same frame budget" # Without the verification assets the structural assertions above # (non-empty payload, byte floor) are the whole test — pass on them. @@ -235,3 +236,199 @@ def test_yue2_metal_twinkle_001(omni_server, openai_client, tmp_path) -> None: assert melody_score > control_score, ( f"calibration: ABC-conditioned {melody_score:.3f} not above control {control_score:.3f}" ) + + +@pytest.mark.slow +@pytest.mark.tts +@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1}) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_yue2_abort_then_next_request_succeeds(omni_server) -> None: + """Serving smoke check after a client timeout, not proof of job cleanup.""" + # The chosen seed and budget time out on the tested cards. Require the + # timeout so this cannot silently become another successful request; + # whether it lands in AR or synthesis depends on the card. + with pytest.raises(httpx.TimeoutException): + httpx.post( + f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech", + json={ + "model": omni_server.model, + "input": LYRICS, + "instructions": CAPTION, + "seed": SEED + 2, + "max_new_tokens": 3000, + "stream": False, + "response_format": "wav", + "extra_params": {"cot": "off"}, + }, + timeout=1.0, + ) + _http_speech(omni_server, seed=SEED + 3, frames=_CONCURRENT_FRAMES) + + +@pytest.mark.slow +@pytest.mark.tts +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.asyncio +@pytest.mark.parametrize("chunked", [True, False], ids=["chunked", "unchunked"]) +@pytest.mark.parametrize("enforce_eager", [True, False], ids=["eager", "ar-graphs"]) +async def test_yue2_preemption_and_synthesis_abort(chunked, enforce_eager, monkeypatch) -> None: + """Force real preemption and abort only after NAR work has been submitted.""" + import asyncio + from pathlib import Path + + import torch + from vllm import SamplingParams + + from tests.helpers.runtime import get_model_prefix + from tests.helpers.stage_config import modify_stage_config + from vllm_omni.entrypoints.async_omni import AsyncOmni + from vllm_omni.model_executor.models.yue2.yue2 import ( + ABC_END, + EOD, + KEY_MAX_AUDIO_FRAMES, + KEY_MAX_HOLD_STEPS, + KEY_MIN_TOKENS, + KEY_PHASE, + KEY_PREFIX_IDS, + KEY_SEED, + KEY_SKIP_SYNTHESIS, + STOP_TOKEN_IDS, + ) + from vllm_omni.tokenizers.yue2_prompt import abc_prefix_ids, semantic_frames, semantic_prefix_ids + from vllm_omni.tokenizers.yue2_tokenizer import YuE2TextTokenizer + from vllm_omni.transformers_utils.repo_utils import hf_api + + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + model_path = Path(os.environ.get("YUE2_MODEL_DIR", get_model_prefix() + MODEL)) + if not model_path.is_dir(): + model_path = Path(hf_api().snapshot_download(str(model_path))) + tokenizer = YuE2TextTokenizer(model_path / "qwen.tiktoken") + prefix = semantic_prefix_ids(tokenizer.encode, CAPTION, LYRICS, "off") + frames = 256 + # Each song fits alone, but four full songs cannot fit in 64 KV blocks. + # Chunked recompute also replays output tokens past the original prompt. + deploy = modify_stage_config( + get_deploy_config_path("yue2.yaml"), + updates={ + "stages": { + 0: { + "enforce_eager": enforce_eager, + "compilation_config": { + "cudagraph_mode": "NONE" if enforce_eager else "FULL_AND_PIECEWISE", + }, + "enable_chunked_prefill": chunked, + "async_scheduling": True, + "max_num_batched_tokens": 256 if chunked else 1024, + "max_model_len": 1000, + "num_gpu_blocks_override": 64, + "block_size": 16, + "max_num_seqs": 4, + "tokenizer": str(model_path), + } + } + }, + ) + engine = AsyncOmni( + model=str(model_path), + deploy_config=deploy, + worker_extension_cls="tests.e2e.features.yue2.worker_extension.Yue2LifecycleWorkerExtension", + stage_init_timeout=600, + init_timeout=900, + ) + + async def probe(): + return (await engine.collective_rpc("get_yue2_probe", stage_ids=[0]))[0][0] + + async def song(request_id, seed): + params = SamplingParams( + max_tokens=1000 - len(prefix), + stop_token_ids=STOP_TOKEN_IDS, + detokenize=False, + extra_args={ + KEY_PREFIX_IDS: prefix, + KEY_SEED: seed, + KEY_MAX_AUDIO_FRAMES: frames, + KEY_MIN_TOKENS: frames, + KEY_MAX_HOLD_STEPS: 1000 - len(prefix) - frames - 1, + }, + ) + audio = [] + tokens = [] + async for output in engine.generate( + {"prompt_token_ids": prefix}, request_id=request_id, sampling_params_list=[params] + ): + tokens = list(output.outputs[0].token_ids) + multimodal = getattr(output.outputs[0], "multimodal_output", None) + if multimodal is not None and "audio" in multimodal: + data = multimodal["audio"] + audio.extend(data if isinstance(data, list) else [data]) + assert audio, "request completed without audio" + waveform = torch.cat([part.detach().cpu().reshape(-1) for part in audio]) + assert torch.isfinite(waveform).all() and waveform.abs().max() > 0 + assert len(semantic_frames(tokens)) == frames + # The VAE convolution trims a small number of boundary samples. + assert waveform.numel() == pytest.approx(frames * 1920 * 2, abs=256) + + async def plan(): + abc_prefix = abc_prefix_ids(tokenizer.encode, CAPTION, LYRICS, "melody") + params = SamplingParams( + max_tokens=64, + stop_token_ids=STOP_TOKEN_IDS, + detokenize=False, + extra_args={ + KEY_PHASE: "abc", + KEY_PREFIX_IDS: abc_prefix, + KEY_SEED: SEED, + KEY_SKIP_SYNTHESIS: True, + }, + ) + tokens = [] + async for output in engine.generate( + {"prompt_token_ids": abc_prefix}, request_id="abc-plan", sampling_params_list=[params] + ): + tokens = list(output.outputs[0].token_ids) + assert tokens and all(t < EOD or t == ABC_END for t in tokens) + + tasks = [] + try: + assert await engine.collective_rpc("start_yue2_probe", stage_ids=[0]) == [[True]] + tasks = [asyncio.create_task(plan())] + tasks.extend(asyncio.create_task(song(f"preempt-{i}", SEED + i)) for i in range(4)) + await asyncio.wait_for(asyncio.gather(*tasks), timeout=600) + observations = await probe() + assert observations["preemptions"] > 0, observations + assert observations["rollbacks"] > 0, observations + assert observations["checked_rows"] >= 4 * frames, observations + assert observations["mixed_steps"] > 0, observations + if enforce_eager: + assert observations["ar_full_replays"] == 0, observations + else: + assert observations["ar_full_replays"] > 0, observations + + cancelled = asyncio.create_task(song("cancel-synthesis", SEED + 7)) + tasks.append(cancelled) + deadline = asyncio.get_running_loop().time() + 120 + while not (await probe())["synthesis_started"]: + assert not cancelled.done(), "song finished before the active-job cancellation probe" + assert asyncio.get_running_loop().time() < deadline, "NAR job never started" + await asyncio.sleep(0.05) + cancelled.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled + deadline = asyncio.get_running_loop().time() + 30 + while True: + observations = await probe() + if observations.get("removed"): + assert observations["cancelled"] and observations["released"], observations + break + assert asyncio.get_running_loop().time() < deadline, observations + await asyncio.sleep(0.05) + await asyncio.wait_for(song("after-cancel", SEED + 8), timeout=120) + observations = await probe() + assert observations["released"] + print("YuE2 lifecycle:", observations) + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + engine.shutdown() diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index 79d8f610d..daa8eabce 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/tests/worker/test_gpu_ar_model_runner.py @@ -16,6 +16,7 @@ import pytest import torch from vllm.sampling_params import SamplingParams from vllm.v1.sample.metadata import SamplingMetadata +from vllm.v1.utils import CpuGpuBuffer from vllm.v1.worker import gpu_input_batch from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch @@ -1885,3 +1886,63 @@ class TestPreferModelSamplerNoneFallback: "least tolerates -- that fallback, then add its directory name to " "`expected` above. If you REMOVED one, drop its name." ) + + +@pytest.mark.parametrize("accepts_extra_args", [False, True]) +def test_prepare_hook_keeps_legacy_signature_and_orders_opt_in_metadata(accepts_extra_args: bool) -> None: + runner = object.__new__(GPUARModelRunner) + runner.model = torch.nn.Module() + runner.model.accepts_runner_sampling_extra_args = accepts_extra_args + runner.requests = { + rid: CachedRequestState( + req_id=rid, + prompt_token_ids=[1], + mm_features=[], + sampling_params=params, + generator=None, + block_ids=([],), + num_computed_tokens=0, + output_token_ids=[], + ) + for rid, params in ( + ("a", SamplingParams(extra_args={"seed": 7})), + ("b", None), + ("unused", SamplingParams(extra_args={"seed": 9})), + ) + } + runner.discard_request_mask = CpuGpuBuffer(3, dtype=torch.bool, device=torch.device("cpu"), pin_memory=False) + runner.discard_request_mask.np[:] = [True, False, True] + ids, positions = torch.tensor([3, 4]), torch.tensor([0, 1]) + step_inputs = dict( + input_ids=ids, + positions=positions, + inputs_embeds=None, + num_computed_tokens=np.array([0, 1]), + num_scheduled_tokens=np.array([1, 1]), + input_ids_buffer=ids, + ) + seen = {} + + def legacy_hook( + *, input_ids, positions, inputs_embeds, req_ids, num_computed_tokens, num_scheduled_tokens, input_ids_buffer + ): + seen["req_ids"] = req_ids + return input_ids, positions + + def opt_in_hook(*, sampling_extra_args, discard_mask, **kwargs): + seen["extra_args"] = sampling_extra_args + seen["discard"] = discard_mask.tolist() + return legacy_hook(**kwargs) + + out = runner._call_prepare_runner_inputs( + opt_in_hook if accepts_extra_args else legacy_hook, + req_ids=["b", "a"], + **step_inputs, + ) + assert out[0] is ids and out[1] is positions + assert seen["req_ids"] == ["b", "a"] + if accepts_extra_args: + assert seen["extra_args"] == [{}, {"seed": 7}] + assert seen["discard"] == [True, False] + else: + assert set(seen) == {"req_ids"} diff --git a/tests/worker/test_yue2_request_state.py b/tests/worker/test_yue2_request_state.py new file mode 100644 index 000000000..75b68c8f3 --- /dev/null +++ b/tests/worker/test_yue2_request_state.py @@ -0,0 +1,215 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""YuE2's worker-side sampler follows retained engine outputs on recompute.""" + +import numpy as np +import pytest +import torch +from vllm.sampling_params import SamplingParams +from vllm.v1.utils import CpuGpuBuffer +from vllm.v1.worker.gpu_input_batch import CachedRequestState + +from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest +from vllm_omni.entrypoints.openai.tts_adapters.yue2 import Yue2Adapter +from vllm_omni.model_executor.models.yue2 import yue2 as y +from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def model(mocker): + """Run real state/sampling control flow with CPU staging and fixed draws.""" + model = object.__new__(y.Yue2ForCausalLM) + torch.nn.Module.__init__(model) + model._states = {} + model._pending = None + model._step_rows = [] + model._step_discard = None + model._decode_t0 = {} + model._last_mm = None + model._audio_queue = [] + model._max_model_len = 4096 + model._stage = torch.empty(1024, dtype=torch.long) + model._pinned = torch.empty((2, 64), dtype=torch.long) + model._synthesis = mocker.Mock(spec=y._SynthesisQueue) + mocker.patch.object(model, "_finish_request_safely") + mocker.patch.object(model, "_sampler_graph", return_value=None) + # CPU copies complete synchronously; this replaces only CUDA completion. + event = mocker.Mock(spec=torch.Event) + mocker.patch.object(y.torch, "Event", return_value=event) + mocker.patch.object(y, "sample_rows", return_value=(torch.tensor([y.CODEC_OFFSET + 7]), torch.tensor([False]))) + constants = y._RowConstants("r", "semantic", 1.0, 0.9, 100, 1.1, 2, 1, 3, 17, False, 8) + model._states["r"] = y._RequestState("r", constants, prompt_len=3, prefix_ids=[1, 2, 3]) + return model + + +def test_pending_draw_is_discarded_after_chunked_recompute(model) -> None: + """A preempted in-flight draw never advances history or the frame limit.""" + logits = torch.zeros((1, y.VOCAB_SIZE)) + model._step_rows = [("r", 0, 3)] + model.sample(logits, None) + state = model._states["r"] + # Other requests run while r waits for KV space. Resolving the host copy + # may resolve a tentative token, but the resumed row must roll it back. + model._step_rows = [] + model.sample(logits[:0], None) + assert state.job is None + # A partial recompute must neither accept the pending draw nor draw again. + model._step_rows = [("r", 0, 2)] + model._step_discard = [True] + model.sample(logits, None) + assert state.job is None + # Completing recompute has retained zero outputs: discard the old draw. + model._step_rows = [("r", 2, 1)] + model._step_discard = [False] + model.sample(logits, None) + assert state.history == [] + # Only this new draw is retained in the next decode input. + model._step_rows = [("r", 3, 1)] + model.sample(logits, None) + assert state.history == [y.CODEC_OFFSET + 7] + assert not state.finish_ready and not state.truncated + + +@pytest.mark.parametrize("token", [y.MUSIC_END, y.CODEC_OFFSET + 7]) +def test_discarded_terminal_draw_does_not_start_synthesis(model, mocker, token: int) -> None: + """Dropped EOS/frame-limit draws cannot prematurely finish a song.""" + state = model._states["r"] + state.history = [y.CODEC_OFFSET + 2] * 2 + mocker.patch.object(y, "sample_rows", return_value=(torch.tensor([token]), torch.tensor([False]))) + logits = torch.zeros((1, y.VOCAB_SIZE)) + model._step_rows = [("r", 4, 1)] + model.sample(logits, None) + model._step_rows = [] + model.sample(logits[:0], None) + assert state.job is None + model._step_rows = [("r", 0, 5)] + model.sample(logits, None) + assert state.history == [y.CODEC_OFFSET + 2] * 2 + assert not state.finish_ready and not state.truncated and not state.end_drawn + assert state.job is None and not state.finished + model._finish_request_safely.assert_not_called() + + +def test_recompute_clears_terminal_state_and_cancels_stale_job(model) -> None: + """Rolling back a drawn end also releases synthesis and revives the row.""" + state = model._states["r"] + state.history = [y.CODEC_OFFSET + 2] * 2 + state.end_drawn = True + state.finished = state.finish_ready = True + state.hold_steps = 3 + state.end_len = 6 + state.delivered = (torch.zeros((2, 0)), False, False) + + def song(keep): + yield + return torch.zeros((2, 0)) + + job = y._SynthesisJob("r", song) + state.job = job + model._reconcile_history(state, 2) + model._synthesis.cancel.assert_called_once_with(job) + assert state.job is None and state.hold_steps == 0 + assert not state.finished and not state.finish_ready and not state.truncated and not state.end_drawn + assert state.end_len is None and state.delivered is None + + +def test_dropped_end_step_redelivers_song_and_end(model, mocker) -> None: + """A preemption that drops the step emitting the end token also drops + that step's audio; the resumed request must deliver both again.""" + state = model._states["r"] + state.history = [y.CODEC_OFFSET + 2] * 3 # at the 3-frame budget + state.truncated = True + state.hold_steps = 2 + audio = torch.ones((2, 5)) + job = mocker.Mock(t_start=1.0, t_done=2.0) + job.done.return_value = True + state.job = job + model._synthesis.complete.return_value = (audio, False, False) + logits = torch.zeros((1, y.VOCAB_SIZE)) + + # Engine holds 3 frames + 2 HOLDs; this step ships the song and its end. + model._step_rows = [("r", 7, 1)] + out = model.sample(logits, None) + assert out.sampled_token_ids.tolist() == [[y.MUSIC_END]] + assert state.finished and state.end_len == 6 + assert [entry[:2] for entry in model._audio_queue] == [("r", audio)] + + # An async lookahead row (end token in flight) delivers nothing again. + model._step_rows = [("r", 8, 1)] + model.sample(logits, None) + assert len(model._audio_queue) == 1 + + # Preempted with the end step in flight: the resumed row lacks it. + model._step_rows = [("r", 0, 8)] + model._step_discard = [False] + out = model.sample(logits, None) + assert out.sampled_token_ids.tolist() == [[y.MUSIC_END]] + assert [entry[:2] for entry in model._audio_queue] == [("r", audio), ("r", audio)] + assert model._audio_queue[1][2:] == (True, False) # truncated, not failed + assert state.end_len == 6 + model._synthesis.complete.assert_called_once() + + +def test_accepted_draw_still_reaches_frame_limit(model) -> None: + """Confirmed codec frames retain the existing delayed finish behavior.""" + logits = torch.zeros((1, y.VOCAB_SIZE)) + for comp, span in ((0, 3), (3, 1), (4, 1)): + model._step_rows = [("r", comp, span)] + model.sample(logits, None) + state = model._states["r"] + assert len(state.history) == 2 + model._resolve_pending() + model._reconcile_history(state, 3) + assert len(state.history) == 3 + assert state.finish_ready and state.truncated + + +def test_real_prepare_hook_admits_reordered_requests(model) -> None: + """Runner opt-in metadata reaches the actual YuE2 admission hook.""" + runner = object.__new__(GPUARModelRunner) + runner.model = model + runner.requests = {} + for rid, seed in (("a", 7), ("b", 9)): + runner.requests[rid] = CachedRequestState( + req_id=rid, + prompt_token_ids=[1, 2, 3], + mm_features=[], + sampling_params=SamplingParams(extra_args={y.KEY_PREFIX_IDS: [1, 2, 3], y.KEY_SEED: seed}), + generator=None, + block_ids=([],), + num_computed_tokens=0, + output_token_ids=[], + ) + runner.discard_request_mask = CpuGpuBuffer(2, dtype=torch.bool, device=torch.device("cpu"), pin_memory=False) + runner.discard_request_mask.np[:] = [True, False] + ids = torch.tensor([1, 2, 1, 2, 3], dtype=torch.int32) + positions = torch.tensor([0, 1, 0, 1, 2]) + runner._call_prepare_runner_inputs( + model.prepare_runner_inputs, + req_ids=["b", "a"], + input_ids=ids, + positions=positions, + inputs_embeds=None, + num_computed_tokens=np.array([0, 0]), + num_scheduled_tokens=np.array([2, 3]), + input_ids_buffer=ids, + ) + assert model._states["b"].constants.seed == 9 + assert model._states["a"].constants.seed == 7 + assert model._states["b"].prompt_len == 3 + assert model._states["b"].prefix_ids == [1, 2, 3] + assert model._step_discard == [True, False] + + +@pytest.mark.parametrize("spare", [0, 1, 4095, 4096, 5000]) +def test_hold_budget_matches_engine_headroom(spare: int) -> None: + """Adapter metadata advertises only the HOLD steps left after clipping.""" + frames = 250 + prompt = {"prompt_token_ids": [1] * (y.CONTEXT - frames - 1 - spare)} + adapter = object.__new__(Yue2Adapter) + request = OpenAICreateSpeechRequest(model="m-a-p/YuE2-3B", input="lyrics", max_new_tokens=frames) + params = adapter.apply_sampling_overrides([SamplingParams()], request, prompt)[0] + assert params.extra_args[y.KEY_MAX_HOLD_STEPS] == min(spare, 4096) + assert params.max_tokens == frames + 1 + min(spare, 4096) diff --git a/tests/worker/test_yue2_sampler.py b/tests/worker/test_yue2_sampler.py new file mode 100644 index 000000000..4a72b0a13 --- /dev/null +++ b/tests/worker/test_yue2_sampler.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""YuE2's batched sampler draws exactly what the per-row reference draws. + +The model samples all rows of a preset as one batch, over the phase's own +columns only, and writes the multinomial out as ``argmax(p / q)`` with +``q ~ Exp(1)`` so the draw stays on the device. The reference below is the +upstream per-row arithmetic over the full vocabulary with +``torch.multinomial``. Same seeds must give the same tokens, and each +request's generator must advance by exactly the same amount per step, or a +request's song changes with the batch it happens to share. +""" + +from __future__ import annotations + +import pytest +import torch + +from vllm_omni.model_executor.models.yue2.yue2 import ( + ABC_END, + ABC_SAMPLING, + CODEC_OFFSET, + CODEC_SIZE, + EOD, + MUSIC_END, + PHASE_VOCAB, + SEMANTIC_SAMPLING, + VOCAB_SIZE, + sample_rows, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _reference_scores(logits: torch.Tensor, preset: dict, history: list[int], phase: str) -> torch.Tensor: + """Upstream per-row scores over the full vocabulary (float32).""" + scores = logits.float() # a new tensor: the logits are bf16 + end = ABC_END if phase == "abc" else MUSIC_END + allowed = torch.full_like(scores, -torch.inf) + if phase == "abc": + allowed[:EOD] = 0 + else: + allowed[CODEC_OFFSET : CODEC_OFFSET + CODEC_SIZE] = 0 + allowed[end] = 0 + scores = scores + allowed + if len(history) < preset["min_tokens"]: + scores[end] = -torch.inf + recent = history[-preset["penalty_window"] :] + if preset["repetition_penalty"] != 1.0 and recent: + freq = torch.zeros_like(scores) + freq.scatter_add_(-1, torch.tensor(recent), torch.ones(len(recent))) + alpha = preset["repetition_penalty"] ** freq + scores = torch.where(scores < 0, scores * alpha, scores / alpha) + if preset["temperature"] == 0: + return scores + scores = scores / preset["temperature"] + threshold = scores.topk(preset["top_k"]).values[-1] + scores = scores.masked_fill(scores < threshold, -torch.inf) + values, indices = scores.sort(descending=True, stable=True) + probabilities = values.softmax(-1) + removed = probabilities.cumsum(-1) - probabilities > preset["top_p"] + removed[:1] = False + values = values.masked_fill(removed, -torch.inf) + return values.scatter(-1, indices, values) + + +def _reference_draw(scores: torch.Tensor, generator: torch.Generator, greedy: bool) -> int: + if greedy: + return int(scores.argmax()) + return int(torch.multinomial(scores.softmax(-1), 1, generator=generator)) + + +def _batched_draw(logits, preset, histories, phase, generators) -> list[int]: + vocab = PHASE_VOCAB[phase] + window = preset["penalty_window"] + recent = [[vocab.col(t) for t in history[-window:]] for history in histories] + width = max(map(len, recent)) + window_cols = torch.tensor([r + [vocab.num_cols] * (width - len(r)) for r in recent]) if width else None + block = [len(history) < preset["min_tokens"] for history in histories] + ids, bad = sample_rows( + logits, + vocab, + temperature=preset["temperature"], + top_p=preset["top_p"], + top_k=preset["top_k"], + repetition_penalty=preset["repetition_penalty"], + window_cols=window_cols, + block_end=torch.tensor(block) if any(block) else None, + generators=generators, + ) + assert not bad.any() + return ids.tolist() + + +@pytest.mark.parametrize( + ("phase", "rows", "min_tokens", "temperature"), + [ + ("semantic", 1, None, None), + ("semantic", 3, None, None), + ("semantic", 3, 0, None), # the end token may be drawn + ("semantic", 2, None, 0.0), # greedy + ("abc", 2, None, None), + ("abc", 2, 0, None), + ], +) +def test_batched_draws_match_per_row_multinomial(phase, rows, min_tokens, temperature) -> None: + preset = dict(ABC_SAMPLING if phase == "abc" else SEMANTIC_SAMPLING) + if min_tokens is not None: + preset["min_tokens"] = min_tokens + if temperature is not None: + preset["temperature"] = temperature + end = ABC_END if phase == "abc" else MUSIC_END + reference = [torch.Generator().manual_seed(1000 + row) for row in range(rows)] + batched = [torch.Generator().manual_seed(1000 + row) for row in range(rows)] + histories: list[list[int]] = [[] for _ in range(rows)] + logits_generator = torch.Generator().manual_seed(7) + for _ in range(40): + logits = (torch.randn((rows, VOCAB_SIZE), generator=logits_generator) * 3).to(torch.bfloat16) + # Peaky over a few ids, like a real LM, so the penalty and top-p matter. + if phase == "semantic": + logits[:, CODEC_OFFSET : CODEC_OFFSET + 300] += 6 + else: + logits[:, :300] += 6 + logits[:, end] += 4 + expected = [ + _reference_draw( + _reference_scores(logits[row], preset, histories[row], phase), reference[row], temperature == 0 + ) + for row in range(rows) + ] + assert _batched_draw(logits, preset, histories, phase, batched) == expected + for row in range(rows): + # A greedy request never reads its stream (its temperature is + # fixed for life), so only sampling rows must advance in step. + if temperature != 0: + assert torch.equal(reference[row].get_state(), batched[row].get_state()) + if expected[row] != end: + histories[row].append(expected[row]) + + +def test_a_row_draws_the_same_alone_or_in_a_batch() -> None: + """A request's song must not depend on which requests share its step.""" + preset = SEMANTIC_SAMPLING + logits = torch.randn((3, VOCAB_SIZE), generator=torch.Generator().manual_seed(3)).to(torch.bfloat16) + histories = [[CODEC_OFFSET + i for i in range(60)], [], [CODEC_OFFSET + 5] * 10] + together = _batched_draw( + logits, preset, histories, "semantic", [torch.Generator().manual_seed(s) for s in (1, 2, 3)] + ) + alone = [ + _batched_draw( + logits[row : row + 1], preset, histories[row : row + 1], "semantic", [torch.Generator().manual_seed(s)] + )[0] + for row, s in enumerate((1, 2, 3)) + ] + assert together == alone + + +@pytest.mark.parametrize("phase", ["abc", "semantic"]) +def test_phase_vocab_columns_round_trip(phase) -> None: + vocab = PHASE_VOCAB[phase] + end = ABC_END if phase == "abc" else MUSIC_END + ids = [end, 0, EOD - 1] if phase == "abc" else [end, CODEC_OFFSET, CODEC_OFFSET + CODEC_SIZE - 1] + cols = torch.tensor([vocab.col(i) for i in ids]) + assert vocab.col(end) == vocab.end_col + assert cols.max() < vocab.num_cols + assert vocab.ids(cols).tolist() == ids + logits = torch.randn((2, VOCAB_SIZE)) + assert torch.equal(vocab.take(logits)[:, cols], logits[:, ids]) diff --git a/vllm_omni/deploy/yue2.yaml b/vllm_omni/deploy/yue2.yaml index 502d950c5..109212824 100644 --- a/vllm_omni/deploy/yue2.yaml +++ b/vllm_omni/deploy/yue2.yaml @@ -1,5 +1,5 @@ # YuE2-3B deploy: single-stage native-AR text-to-music (~3.6B bf16 + VAE). -# Target 1x 24 GB card; verified on a single RTX 4090. +# Four-slot CUDA configuration; validate memory on the target card/workload. # # The architectures pin is required: the checkpoint's config.json names # "YuE2ForCausalLM" with an auto_map that only makes sense under @@ -13,16 +13,17 @@ dtype: bfloat16 stages: - stage_id: 0 + # A finishing song keeps its slot during side-stream synthesis, + # emitting HOLD tokens until its audio is ready. max_num_seqs: 4 - # Measured on a 24 GB RTX 4090 with 4 concurrent requests pinned to the - # 9000-frame cap (worst case): the engine reserves ~15.8 GiB at 0.70 and - # peak usage during the terminal NAR/VAE finishing pass reaches 20.9 GiB, - # leaving ~3 GiB of headroom. 0.85 reserves ~19.4 GiB and the same - # workload OOMs inside the finishing pass. + # Reserve synthesis headroom: NAR K/V, acoustic graph buffers and VAE + # activations are allocated beyond the engine budget. The earlier 4090 + # memory measurement predates async/compiled synthesis; 0.70 is a starting + # point, not a guarantee for 9000 frames or long ABC prefixes. Lower it if + # the target workload needs more synthesis memory (see the model recipe). gpu_memory_utilization: 0.70 - # Phase 1 keeps the backbone eager: the model-owned sampler runs after the - # forward, and the terminal NAR/VAE pass runs inside the finishing step. - # CUDA-graph capture of the standard Qwen3 backbone is a perf follow-up. + # Keep AR capture off by default. Sampler/NAR graphs still need + # separate synthesis headroom. enforce_eager: true # Prefix caching stays off: the omni prefix-cache save snapshots the # step's mm payload before the model-owned sampler runs, but YuE2 @@ -42,14 +43,17 @@ stages: hf_overrides: architectures: ["Yue2ForCausalLM"] default_sampling_params: - # Placeholder values: the model's own sampler overrides everything from - # extra_args. Their presence turns on has_sampling_extra_args, which is - # what threads the per-request yue2_* keys into the model's forward. + # Defaults for callers that omit model-owned sampling options. The + # opt-in prepare_runner_inputs hook reads each request's extra_args; + # it does not depend on the forward's has_sampling_extra_args flag. extra_args: yue2_phase: "semantic" yue2_seed: 831001 yue2_max_audio_frames: 200 + yue2_max_hold_steps: 4096 detokenize: false - # Engine-side backstop; the model stops on its own end token first. - max_tokens: 9000 + stop_token_ids: [151848, 151852] # ABC_END, MUSIC_END + # 200 frames + 4096 synthesis holds + the final end token. Callers + # changing the frame budget must also reserve this token headroom. + max_tokens: 4297 temperature: 1.0 diff --git a/vllm_omni/entrypoints/openai/tts_adapters/yue2.py b/vllm_omni/entrypoints/openai/tts_adapters/yue2.py index bcacc34e7..1d343afb8 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/yue2.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/yue2.py @@ -36,6 +36,7 @@ from vllm_omni.model_executor.models.yue2.yue2 import ( CONTEXT, FRAMES_PER_SECOND, KEY_MAX_AUDIO_FRAMES, + KEY_MAX_HOLD_STEPS, KEY_MIN_TOKENS, KEY_PENALTY_WINDOW, KEY_PHASE, @@ -48,6 +49,7 @@ from vllm_omni.model_executor.models.yue2.yue2 import ( KEY_TOP_P, SEMANTIC_SAMPLING, STOP_TOKEN_IDS, + SYNTHESIS_HOLD_STEPS, ) from vllm_omni.tokenizers.yue2_prompt import semantic_prefix_ids @@ -161,7 +163,9 @@ class Yue2Adapter(ARTTSAdapter): lyrics = request.input or "" prefix = semantic_prefix_ids(encode, instructions, lyrics, cot, abc_ids=abc_ids) max_frames = int(request.max_new_tokens or _DEFAULT_MAX_FRAMES) - if len(prefix) + max_frames > CONTEXT: + # +1: the song is delivered on the step after its last frame (the + # model acts on its draws one step late), which also needs a slot. + if len(prefix) + max_frames + 1 > CONTEXT: return_err = ( f"YuE2 prompt is {len(prefix)} tokens and the {max_frames}-frame budget " f"exceeds the {CONTEXT}-token context. Shorten the lyrics/style or " @@ -209,7 +213,12 @@ class Yue2Adapter(ARTTSAdapter): } # The model needs one decode step beyond the last frame to deliver the # terminal NAR/VAE pass, so do not cap the engine at exactly max_frames. - params.max_tokens = max_frames + 1 + # Synthesize on a side stream while other songs keep decoding; the + # request may emit up to SYNTHESIS_HOLD_STEPS hold tokens meanwhile + # (see the model), so max_tokens leaves that much headroom. + prompt_len = len(prompt["prompt_token_ids"]) + params.max_tokens = min(max_frames + 1 + SYNTHESIS_HOLD_STEPS, CONTEXT - prompt_len) + params.extra_args[KEY_MAX_HOLD_STEPS] = max(0, params.max_tokens - max_frames - 1) params.stop_token_ids = list(STOP_TOKEN_IDS) params.detokenize = False return sampling_params_list diff --git a/vllm_omni/model_executor/models/yue2/yue2.py b/vllm_omni/model_executor/models/yue2/yue2.py index 02549cf84..bd1dd8b94 100644 --- a/vllm_omni/model_executor/models/yue2/yue2.py +++ b/vllm_omni/model_executor/models/yue2/yue2.py @@ -16,7 +16,7 @@ import json import math import os import time -from collections.abc import Callable, Iterable, Sequence +from collections.abc import Callable, Generator, Iterable, Iterator, Sequence from dataclasses import dataclass, field from numbers import Integral from pathlib import Path @@ -45,7 +45,8 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput Values mirror the upstream ``yue2/protocol.py`` (bd90e4c) exactly; changing any of them desynchronizes generation from the reference implementation. The model is single-codebook: one codec token per latent frame, 25 frames per second, -so ``max_tokens`` in engine terms equals frames/25 seconds of music. +so the frame budget divided by 25 is the maximum duration in seconds. +The engine token budget also includes synthesis HOLD and terminal tokens. """ EOD = 151643 @@ -59,6 +60,10 @@ VOCAB_SIZE, CONTEXT = 184704, 24576 # only ever emits the end token of the request's current phase, so a request # never stops on the other phase's end by accident. STOP_TOKEN_IDS = [ABC_END, MUSIC_END] +# Emitted in place of a freshly drawn MUSIC_END: the sampler reads its draws +# back one step late, so the request stays alive for one more step, in which +# the song is synthesized and MUSIC_END is emitted. +HOLD_TOKEN = LATENT_PAD SAMPLE_RATE = 48000 AUDIO_CHANNELS = 2 @@ -107,6 +112,15 @@ KEY_SKIP_SYNTHESIS = "yue2_skip_synthesis" # abc phase: tokens only, no audio # whole prefix, but under a KV prefix-cache hit the engine schedules only the # uncached tail, so the scheduled input_ids slice cannot rebuild it. KEY_PREFIX_IDS = "yue2_prefix_ids" +# How many steps a finishing song may emit HOLD_TOKEN while its NAR/VAE pass +# runs on a side stream (0: synthesize synchronously in the finishing step). +# The caller must leave that many tokens of max_tokens headroom. +KEY_MAX_HOLD_STEPS = "yue2_max_hold_steps" +# What the serving adapter and the offline example grant. A song holds while +# it waits for, and runs on, the side stream; under load the wait behind +# earlier songs dominates (about a thousand steps with 32-48 songs in flight +# on one H200). A request out of budget waits for its song in that step. +SYNTHESIS_HOLD_STEPS = 4096 # -------------------- checkpoint weight routing -------------------- @@ -156,80 +170,187 @@ request-local arithmetic exactly: phase masking (abc = text vocabulary + its end; semantic = codec span + its end), windowed repetition penalty over the request's own history, then temperature/top-k/top-p with a seeded multinomial. Default CFG is off (guidance 1.0 for full/melody, 1.01 for off); the upstream -torch backend treats those the same, so Phase 1 samples a single row. +torch backend treats those the same, so a request samples a single row. + +Rows sharing a preset are sampled as one batch over their phase's columns +only, as a CUDA graph per row count; the multinomial is drawn the way +``torch.multinomial`` draws one sample, so each request's random stream and +every drawn token match the per-row reference arithmetic. """ -def window_penalty(logits: torch.Tensor, recent_ids: list[int], penalty: float) -> torch.Tensor: - """Upweight/downweight ids seen in the window, upstream arithmetic.""" - if penalty == 1.0 or not recent_ids: - return logits - recent = torch.as_tensor(recent_ids, dtype=torch.long, device=logits.device) - if logits.dim() > 1: - recent = recent.reshape(1, -1) - freq = torch.zeros_like(logits) - freq.scatter_add_(-1, recent, torch.ones_like(recent, dtype=logits.dtype)) - alpha = penalty**freq - return torch.where(logits < 0, logits * alpha, logits / alpha) +# The semantic phase may only emit MUSIC_END or a codec token, and MUSIC_END +# sits right below the codec span, so its whole vocabulary is one slice. +assert MUSIC_END + 1 == CODEC_OFFSET -def distribution( +class PhaseVocab: + """Columns one phase may sample, in increasing token-id order. + + Phase masking sets every other column to -inf, so dropping them changes + no finite score, and keeping the id order keeps sort/argmax tie breaking + the same as on the full vocabulary. + """ + + def __init__(self, phase: str): + self.phase = phase + if phase == "abc": + self.end_col = EOD # columns [0, EOD) are ids; column EOD is ABC_END + self.num_cols = EOD + 1 + else: + self.end_col = 0 # column c is id MUSIC_END + c + self.num_cols = CODEC_OFFSET + CODEC_SIZE - MUSIC_END + self._index: dict[torch.device, torch.Tensor] = {} + + def take(self, rows: torch.Tensor) -> torch.Tensor: + """Allowed columns of ``rows`` [R, vocab] as a new float32 tensor.""" + if self.phase == "abc": + index = self._index.get(rows.device) + if index is None: + index = torch.cat((torch.arange(EOD), torch.tensor([ABC_END]))).to(rows.device) + self._index[rows.device] = index + return rows.index_select(-1, index).float() + return rows[..., MUSIC_END : CODEC_OFFSET + CODEC_SIZE].to(torch.float32, copy=True) + + def col(self, token_id: int) -> int: + if self.phase == "abc": + return self.end_col if token_id == ABC_END else token_id + return token_id - MUSIC_END + + def ids(self, cols: torch.Tensor) -> torch.Tensor: + if self.phase == "abc": + return torch.where(cols == self.end_col, ABC_END, cols) + return cols + MUSIC_END + + +PHASE_VOCAB = {"abc": PhaseVocab("abc"), "semantic": PhaseVocab("semantic")} + + +def _sample_core( logits: torch.Tensor, - *, + noise: torch.Tensor, + vocab: PhaseVocab, temperature: float, top_p: float, top_k: int, repetition_penalty: float, - penalty_window: int, - history: list[int], - step: int, - min_tokens: int, - phase: str, -) -> torch.Tensor: - """Masked/penalized/shaped scores for one row (float32 path). - - Accepts either one row ``[vocab]`` (what the model's sampler sees) or a - batch ``[rows, vocab]`` and returns the same shape it was given. - """ - single = logits.dim() == 1 - logits = logits.unsqueeze(0) if single else logits - scores = logits.float().clone() - end = ABC_END if phase == "abc" else MUSIC_END - allowed = torch.full_like(scores, float("-inf")) - if phase == "abc": - allowed[..., :EOD] = 0 - else: - allowed[..., CODEC_OFFSET : CODEC_OFFSET + CODEC_SIZE] = 0 - allowed[..., end] = 0 - scores = scores + allowed - if step < min_tokens: - scores[..., end] = -torch.inf - scores = window_penalty(scores, history[-penalty_window:], repetition_penalty) + window_cols: torch.Tensor | None, + block_end: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Masked, penalized, shaped scores and the multinomial draw, given the + Exp(1) noise [R, vocab].""" + rows = logits.shape[0] + scores = vocab.take(logits) + if block_end is not None: + scores[:, vocab.end_col].masked_fill_(block_end, -torch.inf) + if repetition_penalty != 1.0 and window_cols is not None: + # The padding column is counted, then dropped. + freq = torch.zeros((rows, vocab.num_cols + 1), dtype=scores.dtype, device=scores.device) + freq.scatter_add_(-1, window_cols, torch.ones_like(window_cols, dtype=scores.dtype)) + alpha = repetition_penalty ** freq[:, : vocab.num_cols] + scores = torch.where(scores < 0, scores * alpha, scores / alpha) + bad = ~torch.isfinite(scores).any(-1) if temperature == 0: - return scores.squeeze(0) if single else scores + return vocab.ids(scores.argmax(-1)), bad if temperature != 1: scores = scores / temperature - threshold = scores.topk(min(top_k, scores.shape[-1])).values[..., -1, None] + threshold = scores.topk(min(top_k, vocab.num_cols)).values[..., -1, None] scores = scores.masked_fill(scores < threshold, -torch.inf) if top_p < 1: - values, indices = scores.sort(descending=True) + values, indices = scores.sort(descending=True, stable=True) probabilities = values.softmax(-1) removed = probabilities.cumsum(-1) - probabilities > top_p removed[..., :1] = False values = values.masked_fill(removed, -torch.inf) scores = values.scatter(-1, indices, values) - return scores.squeeze(0) if single else scores - - -def sample_row(scores: torch.Tensor, generator: torch.Generator, *, greedy: bool = False) -> int: - """Draw one token id from prepared scores with the request's generator.""" - if not torch.isfinite(scores).any(): - raise RuntimeError("sampling scores are all -inf; check phase masking") - if greedy: - return int(scores.argmax().item()) probabilities = scores.softmax(-1) - next_id = torch.multinomial(probabilities, 1, generator=generator) - return int(next_id.item()) + return vocab.ids((probabilities / vocab.take(noise)).argmax(-1)), bad + + +def draw_noise(noise: torch.Tensor, generators: Sequence[torch.Generator]) -> torch.Tensor: + """Row r = the Exp(1) draw ``torch.multinomial`` makes from generator r.""" + for row, generator in enumerate(generators): + noise[row].exponential_(1, generator=generator) + return noise + + +def sample_rows( + logits: torch.Tensor, + vocab: PhaseVocab, + *, + temperature: float, + top_p: float, + top_k: int, + repetition_penalty: float, + window_cols: torch.Tensor | None, + block_end: torch.Tensor | None, + generators: list[torch.Generator], +) -> tuple[torch.Tensor, torch.Tensor]: + """Sample one token per row for rows sharing one preset. + + Works on the phase's allowed columns only and never synchronizes: every + input is already on the device (``window_cols`` [R, W] are the recent + ids as vocab columns, padded with ``vocab.num_cols``; ``block_end`` [R] + masks the end token) and it returns device tensors (token ids [R], + all-masked flags [R]). The multinomial draw is written out the way + ``torch.multinomial`` draws one sample (``argmax(p / q)`` with + ``q ~ Exp(1)`` over the full vocabulary), so each request's generator + advances exactly as in the per-row path. + """ + noise = torch.empty(logits.shape, dtype=torch.float32, device=logits.device) + draw_noise(noise, generators) + return _sample_core(logits, noise, vocab, temperature, top_p, top_k, repetition_penalty, window_cols, block_end) + + +class SamplerGraph: + """``sample_rows`` for a row-count bucket and preset, replayed as one graph. + + The sampler is ~30 small kernels per step, each a host-side launch; in a + batch-1 decode step that host time exceeds the backbone's GPU time. The + Exp(1) noise is still drawn per row from each request's own generator + outside the graph, and an all-False ``block_end`` / padded ``window_cols`` + leave every score as the eager path computes it, so draws are unchanged. + """ + + def __init__(self, vocab, rows, vocab_size, logits_dtype, preset, window, device, pool): + self.vocab = vocab + self.preset = preset # (temperature, top_p, top_k, repetition_penalty) + self.logits = torch.zeros((rows, vocab_size), dtype=logits_dtype, device=device) + self.noise = torch.ones((rows, vocab_size), dtype=torch.float32, device=device) + self.block_end = torch.zeros(rows, dtype=torch.bool, device=device) + self.window_cols = ( + torch.full((rows, window), vocab.num_cols, dtype=torch.long, device=device) if window else None + ) + side = torch.cuda.Stream(device) + side.wait_stream(torch.cuda.current_stream(device)) + with torch.cuda.stream(side): + for _ in range(2): + self._run() + torch.cuda.current_stream(device).wait_stream(side) + self.graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(self.graph, pool=pool, capture_error_mode="thread_local"): + self.ids, self.bad = self._run() + + def _run(self): + return _sample_core(self.logits, self.noise, self.vocab, *self.preset, self.window_cols, self.block_end) + + def __call__(self, logits, window_cols, block_end, generators): + rows = logits.shape[0] + if rows > self.logits.shape[0] or len(generators) != rows: + raise ValueError("Sampler inputs exceed the captured bucket or lack a request generator") + # A replay reads its inputs from the addresses it was captured with. + # Padding rows have independent arithmetic and never advance a real + # request's generator. Only the active output rows are returned. + self.logits[:rows].copy_(logits) + draw_noise(self.noise, generators) + if self.window_cols is not None: + self.window_cols[:rows].copy_(window_cols) + if block_end is None: + self.block_end.zero_() + else: + self.block_end[:rows].copy_(block_end) + self.graph.replay() + return self.ids[:rows], self.bad[:rows] # -------------------- NAR flow-matching pass -------------------- @@ -238,11 +359,13 @@ def sample_row(scores: torch.Tensor, generator: torch.Generator, *, greedy: bool """Acoustic flow matching for YuE2-3B, ported from upstream ``nar.py``. One AR prefill per original chunk, then 32 midpoint ODE steps that walk only -the NAR path. The AR prefill re-uses the vLLM backbone's own projections -(fused ``qkv_proj`` split by head counts, ``q_norm``/``k_norm``, the layer's -``o_proj``/``mlp``), with RoPE and SDPA computed here exactly like the -reference: NAR positions attend the full AR prefix plus all NAR positions -bidirectionally, which a per-chunk prefill + concatenated K/V reproduces. +the NAR path: NAR positions attend the full AR prefix plus all NAR positions +bidirectionally. The prefill re-uses the vLLM backbone's own projections and +writes the prefix K/V once into a per-layer buffer; every ODE evaluation +writes its NAR K/V right after it and attends one contiguous sequence. An +evaluation is FlashAttention-3 (SDPA without it) plus the per-layer halves in +``_nar_pre_attention``/``_nar_post_attention`` (compiled), and after the +first call a CUDA graph replay: all 64 evaluations of a chunk share shapes. """ @@ -282,47 +405,158 @@ def _linear(module, x): return out[0] if isinstance(out, tuple) else out -def _attention(q, k, v, *, causal: bool, query_chunk_size: int | None = None): - """SDPA over [tokens, heads, dim]; query tiling only bounds temporaries.""" - if q.ndim != 3 or k.ndim != 3 or v.shape != k.shape or q.shape[-1] != k.shape[-1]: - raise ValueError("Expected Q/K/V [tokens, heads, dim] with matching K/V") - if min(q.shape) < 1 or min(k.shape) < 1 or q.shape[1] % k.shape[1]: - raise ValueError("Invalid attention lengths or grouped-query head count") - if causal and len(q) != len(k): - raise ValueError("Causal prefill requires matching Q/K sequence lengths") - block = query_chunk_size or (len(q) if q.device.type == "cuda" else 256) - query = q.transpose(0, 1).unsqueeze(0) - key = k.transpose(0, 1).unsqueeze(0) - value = v.transpose(0, 1).unsqueeze(0) - grouped = query.shape[1] != key.shape[1] - outputs = [] - for start in range(0, len(q), block): - end = min(start + block, len(q)) - used_key = key[..., :end, :] if causal else key - used_value = value[..., :end, :] if causal else value - mask = None - if causal and start: - mask = torch.arange(end, device=q.device)[None, :] <= torch.arange(start, end, device=q.device)[:, None] - outputs.append( - F.scaled_dot_product_attention( - query[..., start:end, :], - used_key, - used_value, - attn_mask=mask, - is_causal=causal and start == 0, - enable_gqa=grouped, +def _attention(q, k, v, *, causal: bool) -> torch.Tensor: + """SDPA over [tokens, heads, dim] Q/K/V (grouped K/V heads allowed).""" + out = F.scaled_dot_product_attention( + q.transpose(0, 1).unsqueeze(0), + k.transpose(0, 1).unsqueeze(0), + v.transpose(0, 1).unsqueeze(0), + is_causal=causal, + enable_gqa=q.shape[1] != k.shape[1], + ) + return out[0].transpose(0, 1) + + +def _drain(work: Iterator[None]) -> torch.Tensor: + """Run a work generator to completion and return its result.""" + while True: + try: + next(work) + except StopIteration as done: + return done.value + + +def _rms(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + # Same arithmetic as _RMSNorm.forward. + return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + eps).to(x.dtype) * weight + + +def _rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + """_apply_rotary for [T, heads, dim] with [T, dim/2] cos/sin.""" + return _apply_rotary(x, cos[:, None, :], sin[:, None, :]) + + +def _nar_pre_attention(x, w_norm, w_qkv, w_q_norm, w_k_norm, cos, sin, eps: float, heads: int, kv_heads: int): + """input norm -> fused q/k/v projection -> q/k norm -> RoPE, for [T, hidden].""" + tokens = x.shape[0] + head_dim = w_q_norm.shape[0] + qkv = F.linear(_rms(x, w_norm, eps), w_qkv) + q, k, v = qkv.split([heads * head_dim, kv_heads * head_dim, kv_heads * head_dim], dim=-1) + q = _rope(_rms(q.reshape(tokens, heads, head_dim), w_q_norm, eps), cos, sin) + k = _rope(_rms(k.reshape(tokens, kv_heads, head_dim), w_k_norm, eps), cos, sin) + return q, k, v.reshape(tokens, kv_heads, head_dim) + + +def _nar_post_attention(x, attn, w_o, w_norm, w_gate_up, w_down, eps: float): + """o_proj + residual -> pre-MLP norm -> fused gate/up SiLU MLP + residual.""" + x = x + F.linear(attn.flatten(1), w_o) + gate, up = F.linear(_rms(x, w_norm, eps), w_gate_up).chunk(2, dim=-1) + return x + F.linear(F.silu(gate) * up, w_down) + + +_COMPILED: dict[str, Callable] = {} + + +def _nar_kernels() -> tuple[Callable, Callable]: + """The per-layer NAR halves, compiled once for any token count. + + Compilation fuses the RMSNorm/RoPE/SiLU elementwise chains that dominate + the eager NAR pass after its GEMMs and attention. + """ + if "pre" not in _COMPILED: + pre, post = _nar_pre_attention, _nar_post_attention + if torch.cuda.is_available(): + pre = torch.compile(pre, dynamic=True, fullgraph=True, options=_COMPILE_OPTIONS) + post = torch.compile(post, dynamic=True, fullgraph=True, options=_COMPILE_OPTIONS) + _COMPILED.update(pre=pre, post=post) + return _COMPILED["pre"], _COMPILED["post"] + + +def _flash_attention_fn() -> Callable | None: + """FlashAttention-3 varlen from vLLM on Hopper, else None (use SDPA).""" + if "fa3" not in _COMPILED: + fn = None + if torch.cuda.is_available(): + try: + from vllm.vllm_flash_attn import flash_attn_interface, flash_attn_varlen_func + except ImportError: + pass + else: + if flash_attn_interface.is_fa_version_supported(3): + fn = flash_attn_varlen_func + _COMPILED["fa3"] = fn + return _COMPILED["fa3"] + + +def _sequence_bounds(length: int, device: torch.device) -> torch.Tensor: + return torch.tensor([0, length], dtype=torch.int32).to(device, non_blocking=True) + + +def _flash_attention(fa3, q, k, v, *, causal: bool, bounds_q=None, bounds_k=None) -> torch.Tensor: + """One sequence of [tokens, heads, dim] Q/K/V (grouped K/V heads allowed). + + ``fa3`` is FlashAttention-3 varlen, or None for PyTorch SDPA. + """ + if fa3 is None: + return _attention(q, k, v, causal=causal) + bounds_q = _sequence_bounds(q.shape[0], q.device) if bounds_q is None else bounds_q + bounds_k = _sequence_bounds(k.shape[0], q.device) if bounds_k is None else bounds_k + return fa3( + q, + k, + v, + max_seqlen_q=q.shape[0], + cu_seqlens_q=bounds_q, + max_seqlen_k=k.shape[0], + cu_seqlens_k=bounds_k, + causal=causal, + fa_version=3, + ) + + +def fuse_nar_projections(layers: nn.ModuleList) -> list[tuple[torch.Tensor, ...]]: + """Per NAR layer: norm/qkv/q-norm/k-norm/o/norm/gate-up/down weights. + + q/k/v and gate/up are concatenated so each runs as one GEMM; the original + ``nn.Linear`` weights become views into the fused tensors, so the module + tree stays intact at no extra memory. + """ + fused = [] + for layer in layers: + attn, mlp = layer.self_attn, layer.mlp + projections = (attn.q_proj, attn.k_proj, attn.v_proj) + qkv = torch.cat([p.weight.data for p in projections]) + offset = 0 + for p in projections: + rows = p.weight.shape[0] + p.weight.data = qkv[offset : offset + rows] + offset += rows + gate_up = torch.cat([mlp.gate_proj.weight.data, mlp.up_proj.weight.data]) + rows = mlp.gate_proj.weight.shape[0] + mlp.gate_proj.weight.data = gate_up[:rows] + mlp.up_proj.weight.data = gate_up[rows:] + fused.append( + ( + layer.input_layernorm.weight, + qkv, + attn.q_norm.weight, + attn.k_norm.weight, + attn.o_proj.weight, + layer.pre_mlp_layernorm.weight, + gate_up, + mlp.down_proj.weight, ) ) - return torch.cat(outputs, dim=-2)[0].transpose(0, 1) + return fused class CachedNAR: """One acoustic chunk: AR prefix K/V prefilled once, ODE walks the NAR path.""" - def __init__(self, model, chunk: Chunk, query_chunk_size: int | None = None): + def __init__(self, model, chunk: Chunk, pool: tuple[int, int] | None = None): self.model = model self.chunk = chunk - self.query_chunk_size = query_chunk_size + self.pool = pool weight = next(model.vae2llm.parameters()) self.device, self.dtype = weight.device, weight.dtype if chunk.noise.ndim != 2 or chunk.noise.shape[1] != 64 or len(chunk.noise) < 1: @@ -333,65 +567,182 @@ class CachedNAR: if self.ar_length + self.nar_length > model.max_position_embeddings: raise ValueError("Acoustic chunk exceeds the model context") positions = torch.arange(self.ar_length, self.ar_length + self.nar_length, device=self.device)[None] - self.cos, self.sin = model.rotary(positions) + cos, sin = model.rotary(positions) + self.cos, self.sin = cos[0], sin[0] # [NAR, dim/2] local = torch.arange(self.nar_length, device=self.device).clamp(max=model.max_latent_frames - 1) self.pos_emb = model.latent_pos_embed(local)[None] - self.cache: list[tuple[torch.Tensor, torch.Tensor]] = [] + self.fa3 = _flash_attention_fn() + self.bounds_q = _sequence_bounds(self.nar_length, self.device) + self.bounds_k = _sequence_bounds(self.ar_length + self.nar_length, self.device) + self.graph: torch.cuda.CUDAGraph | None = None + self.last_use: torch.cuda.Event | None = None self._prefill() def _prefill(self): - backbone = self.model.model - ids = torch.tensor([self.chunk.ar_tokens], dtype=torch.long, device=self.device) - positions = torch.arange(self.ar_length, device=self.device)[None] - cos, sin = self.model.rotary(positions) - x = backbone.embed_tokens(ids)[0] # [T, H] - for layer in backbone.layers: - # Upstream: layer.self_attn.project_qkv(layer.input_layernorm(x), ...) + model = self.model + backbone = model.model + config = model.config + total = self.ar_length + self.nar_length + shape = (len(backbone.layers), total, int(config.num_key_value_heads), int(config.head_dim)) + self.keys = torch.empty(shape, dtype=self.dtype, device=self.device) + self.values = torch.empty(shape, dtype=self.dtype, device=self.device) + ids = torch.tensor(self.chunk.ar_tokens, dtype=torch.long).to(self.device, non_blocking=True) + cos, sin = model.rotary(torch.arange(self.ar_length, device=self.device)[None]) + x = backbone.embed_tokens(ids) # [T, H] + last = len(backbone.layers) - 1 + for index, layer in enumerate(backbone.layers): # The layernorm is load-bearing: without it the cached AR K/V # conditioning is garbage and the ODE latents blow up (clipped audio). - q, k, v = self.model.ar_project(layer, layer.input_layernorm(x), cos, sin) - self.cache.append((k, v)) - h = _attention(q, k, v, causal=True, query_chunk_size=self.query_chunk_size) + q, k, v = model.ar_project(layer, layer.input_layernorm(x), cos, sin) + self.keys[index, : self.ar_length] = k + self.values[index, : self.ar_length] = v + if index == last: + break # the last layer's output feeds nothing + h = _flash_attention(self.fa3, q, k, v, causal=True) x = x + _linear(layer.self_attn.o_proj, h.flatten(1)) x = x + _linear(layer.mlp, layer.post_attention_layernorm(x)) @torch.inference_mode() def velocity(self, state: torch.Tensor, raw_t: float) -> torch.Tensor: + """One ODE evaluation; from the second call on, a CUDA graph replay. + + After one eager call (compilation, cuBLAS/FA3 setup) the evaluation is + captured once and replayed from static state/time inputs, replacing + ~30 host-side launches per layer with one launch per evaluation. + """ + if self.device.type != "cuda": + return self._velocity(state, torch.tensor(raw_t, dtype=self.dtype, device=self.device)) + if self.graph is None: + if not hasattr(self, "state_in"): + self.state_in = torch.empty_like(state) + self.time_in = torch.empty((), dtype=self.dtype, device=self.device) + self.state_in.copy_(state) + self.time_in.fill_(raw_t) + return self._velocity(self.state_in, self.time_in) + self.graph = torch.cuda.CUDAGraph() + if torch.cuda.current_stream(self.device) == torch.cuda.default_stream(self.device): + with torch.cuda.graph(self.graph, pool=self.pool, capture_error_mode="thread_local"): + self.velocity_out = self._velocity(self.state_in, self.time_in) + else: + # Already on a side stream (the async synthesis path): capture + # in place, without torch.cuda.graph's device-wide sync, which + # would stall the decode loop behind every queued kernel. + self.graph.capture_begin(pool=self.pool, capture_error_mode="thread_local") + try: + self.velocity_out = self._velocity(self.state_in, self.time_in) + finally: + self.graph.capture_end() + # The graph reads its inputs from the addresses it was captured with. + self.state_in.copy_(state) + self.time_in.fill_(raw_t) + self.graph.replay() + return self.velocity_out + + def _velocity(self, state: torch.Tensor, raw_t: torch.Tensor) -> torch.Tensor: model = self.model + pre, post = _nar_kernels() + config = model.config + eps = float(config.rms_norm_eps) + heads, kv_heads = int(config.num_attention_heads), int(config.num_key_value_heads) x_nar = F.pad(state, (0, 0, 1, 1)) - shifted = model.shift_t(raw_t, self.device, self.dtype) + # The timestep shift, from a device scalar so a graph replay can update it. + t_sig = torch.sigmoid(raw_t) + shift = model._timestep_shift + shifted = shift * t_sig / (1 + (shift - 1) * t_sig) x = model.vae2llm(x_nar[None]) x = x + model.time_embedder(shifted.expand(self.nar_length))[None] - x = x + self.pos_emb - for layer, (ar_k, ar_v) in zip(model.nar_layers, self.cache): - q, k, v = layer.self_attn.project_qkv(layer.input_layernorm(x), self.cos, self.sin) - k, v = torch.cat((ar_k, k[0])), torch.cat((ar_v, v[0])) - h = _attention(q[0], k, v, causal=False, query_chunk_size=self.query_chunk_size) - x = x + _linear(layer.self_attn.o_proj, h.flatten(1)[None]) - x = x + layer.mlp(layer.pre_mlp_layernorm(x)) - return model.llm2vae(model.model.norm(x))[0, 1:-1] + x = (x + self.pos_emb)[0] + for index, weights in enumerate(model.nar_weights()): + w_norm, w_qkv, w_q_norm, w_k_norm, w_o, w_mlp_norm, w_gate_up, w_down = weights + q, k, v = pre(x, w_norm, w_qkv, w_q_norm, w_k_norm, self.cos, self.sin, eps, heads, kv_heads) + self.keys[index, self.ar_length :] = k + self.values[index, self.ar_length :] = v + h = _flash_attention( + self.fa3, + q, + self.keys[index], + self.values[index], + causal=False, + bounds_q=self.bounds_q, + bounds_k=self.bounds_k, + ) + x = post(x, h, w_o, w_mlp_norm, w_gate_up, w_down, eps) + return model.llm2vae(model.model.norm(x[None]))[0, 1:-1] - @torch.inference_mode() - def solve(self, steps: int = 32) -> torch.Tensor: + def solve_iter(self, steps: int = 32) -> Iterator[None]: + """The ODE solve as a work generator: yields after enqueuing each + evaluation and returns the (still in flight) device latents.""" if isinstance(steps, bool) or not isinstance(steps, Integral) or steps < 1: raise ValueError("steps must be a positive integer") - state = self.chunk.noise.to(device=self.device, dtype=self.dtype) + # Non-blocking: a blocking copy would wait for the queued prefill. + state = self.chunk.noise.to(device=self.device, dtype=self.dtype, non_blocking=True) dt = 1.0 / steps for step in range(steps): t = 1.0 - step * dt raw = torch.logit(torch.tensor(t, dtype=torch.float64, device="cpu")).clamp(-20, 20).item() first = self.velocity(state, raw) + yield mid = state - first * (dt / 2) raw_mid = torch.logit(torch.tensor(t - dt / 2, dtype=torch.float64, device="cpu")).clamp(-20, 20).item() state = state - self.velocity(mid, raw_mid) * dt - result = state.float().cpu() + yield + return state.float() + + @torch.inference_mode() + def solve(self, steps: int = 32) -> torch.Tensor: + result = _drain(self.solve_iter(steps)).cpu() if not torch.isfinite(result).all(): raise FloatingPointError("Acoustic flow matching produced non-finite latents") return result def close(self): - self.cache.clear() + self.keys = self.values = self.graph = None self.cos = self.sin = self.pos_emb = None + self.state_in = self.time_in = self.velocity_out = None + self.bounds_q = self.bounds_k = None + self.chunk = None + + +class _NARGraphPool: + """Allocator pool shared by graphs on the serialized synthesis stream.""" + + def __init__(self) -> None: + self.pool: tuple[int, int] | None = None + self._anchor: torch.cuda.CUDAGraph | None = None + self._anchor_buffer: torch.Tensor | None = None + + def handle(self, device: torch.device) -> tuple[int, int]: + if self.pool is None: + self.pool = torch.cuda.graph_pool_handle() + # The handle alone does not keep the allocator pool alive. A tiny + # graph owns it after a chunk is released, including across aborts. + self._anchor = torch.cuda.CUDAGraph() + self._anchor.capture_begin(pool=self.pool, capture_error_mode="thread_local") + try: + self._anchor_buffer = torch.zeros((), dtype=torch.int32, device=device) + finally: + self._anchor.capture_end() + return self.pool + + +def synthesize_iter( + model, prefix: Sequence[int], codec: Sequence[int], seed: int, keep: list, graph_pool: _NARGraphPool | None = None +) -> Iterator[None]: + """``synthesize`` as a work generator (one unit per prefill or ODE + evaluation); returns the in-flight device latents, unchecked. The solver + engines go to ``keep``: their graphs and buffers must outlive the work.""" + output = [] + for chunk in song_chunks(prefix, codec, seed): + # Graphs share allocator storage only on this serialized stream. + # Their state and K/V are retired after each chunk, with no engine cache. + pool = None if graph_pool is None else graph_pool.handle(next(model.vae2llm.parameters()).device) + engine = CachedNAR(model, chunk, pool) + keep.append(engine) + yield + output.append((yield from engine.solve_iter(ODE_STEPS))) + engine.last_use = torch.cuda.Event() + engine.last_use.record() + return torch.cat(output, dim=0) @torch.inference_mode() @@ -440,6 +791,26 @@ def snake_beta(x, alpha, beta): return x + (1.0 / (beta + 0.000000001)) * torch.pow(torch.sin(x * alpha), 2) +_SNAKE: list[Callable] = [] +# Inductor otherwise benchmarks a few launch configs per reduction kernel and +# keeps the fastest; timing noise then changes the reduction order, and with it +# the song's last bits, from one process to the next. Deterministic mode keeps +# "same seed, same song" across server restarts. +_COMPILE_OPTIONS = {"deterministic": True} + + +def _snake_kernel() -> Callable: + """snake_beta fused into one kernel: it runs on every [C, samples] FP32 + activation of the decoder, where each unfused elementwise op is a full + memory pass.""" + if not _SNAKE: + fused = torch.cuda.is_available() + _SNAKE.append( + torch.compile(snake_beta, dynamic=True, fullgraph=True, options=_COMPILE_OPTIONS) if fused else snake_beta + ) + return _SNAKE[0] + + class SnakeBeta(nn.Module): def __init__( self, @@ -467,6 +838,8 @@ class SnakeBeta(nn.Module): if self.alpha_logscale: alpha = torch.exp(alpha) beta = torch.exp(beta) + if x.is_cuda and not self.training: + return _snake_kernel()(x, alpha, beta) return snake_beta(x, alpha, beta) @@ -981,11 +1354,13 @@ class YuE2VAE(PreTrainedModel): def decoder_device(self): return next(self.decoder.parameters()).device - def _latent(self, latent): + def _latent(self, latent, check_finite=True): latent = torch.as_tensor(latent) if latent.ndim != 3 or latent.shape[1] != self.config.latent_dim or latent.shape[0] < 1 or latent.shape[-1] < 1: raise ValueError(f"Expected nonempty [B,{self.config.latent_dim},T] latents") - if not torch.isfinite(latent).all(): + # The check reads the values back, a device sync; an async caller + # checks the decoded audio on the host instead. + if check_finite and not torch.isfinite(latent).all(): raise ValueError("VAE latents contain non-finite values") if next(self.decoder.parameters()).dtype != torch.float32: raise ValueError("VAE decoder weights must remain FP32") @@ -1024,9 +1399,9 @@ class YuE2VAE(PreTrainedModel): return latent @torch.inference_mode() - def decode(self, latent): + def decode(self, latent, check_finite=True): """Full waveform, FP32 [B,2,1920*T-64], without clipping.""" - latent = self._latent(latent) + latent = self._latent(latent, check_finite) with torch.autocast(device_type=self.decoder_device.type, enabled=False): return self.decoder(latent.to(device=self.decoder_device, dtype=torch.float32)) @@ -1049,6 +1424,20 @@ class YuE2VAE(PreTrainedModel): halo_frames=None, output_device="cpu", on_progress: Callable[[int, int], None] | None = None, + check_finite=True, + ): + return _drain( + self.decode_tiled_iter(latent, core_frames, halo_frames, output_device, on_progress, check_finite) + ) + + def decode_tiled_iter( + self, + latent, + core_frames=None, + halo_frames=None, + output_device="cpu", + on_progress: Callable[[int, int], None] | None = None, + check_finite=True, ): """Decode bounded tiles, retaining exact cores with natural end length. @@ -1059,7 +1448,7 @@ class YuE2VAE(PreTrainedModel): no extra synchronization is added. With a CUDA output device, queued work may still be executing. Callback exceptions propagate. """ - latent = self._latent(latent) + latent = self._latent(latent, check_finite) core_frames = self.config.decode_core_frames if core_frames is None else core_frames halo_frames = self.config.decode_halo_frames if halo_frames is None else halo_frames if not isinstance(core_frames, int) or core_frames < 1: @@ -1078,7 +1467,7 @@ class YuE2VAE(PreTrainedModel): end = min(frames, start + core_frames) left = max(0, start - halo_frames) right = min(frames, end + halo_frames) - tile = self.decode(latent[..., left:right]) + tile = self.decode(latent[..., left:right], check_finite) out_start, out_end = start * ratio, min(end * ratio, total) crop_start = (start - left) * ratio crop = tile[..., crop_start : crop_start + out_end - out_start] @@ -1088,6 +1477,7 @@ class YuE2VAE(PreTrainedModel): del tile, crop if on_progress is not None: on_progress(tile_index + 1, tiles) + yield return audio def decode_audio(self, latent, chunked=True, **kwargs): @@ -1180,15 +1570,6 @@ class _NARAttention(nn.Module): self.q_norm = _RMSNorm(self.head_dim, config.rms_norm_eps) self.k_norm = _RMSNorm(self.head_dim, config.rms_norm_eps) - def project_qkv(self, x, cos, sin): - B, T, _ = x.shape - q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim) - k = self.k_proj(x).view(B, T, self.num_kv_heads, self.head_dim) - v = self.v_proj(x).view(B, T, self.num_kv_heads, self.head_dim) - q, k = self.q_norm(q), self.k_norm(k) - rc, rs = cos.unsqueeze(2), sin.unsqueeze(2) - return _apply_rotary(q, rc, rs), _apply_rotary(k, rc, rs), v - class _NARMLP(nn.Module): def __init__(self, config): @@ -1265,6 +1646,7 @@ class _RowConstants: max_audio_frames: int seed: int skip_synthesis: bool + max_hold_steps: int = 0 @dataclass @@ -1273,13 +1655,263 @@ class _RequestState: request_id: str constants: _RowConstants - prompt_tokens: int = 0 prompt_len: int = 0 # full prompt length; a completing prefill row reaches it prefix_ids: list[int] = field(default_factory=list) history: list[int] = field(default_factory=list) generator: torch.Generator | None = None finished: bool = False truncated: bool = False + # Set when the host learns (one step after the draw) that the request + # drew its end token or filled its frame budget; the next scheduled step + # runs the finishing pass and emits the end token. + finish_ready: bool = False + end_drawn: bool = False # the engine counts a drawn end (or its HOLD) too + job: _SynthesisJob | None = None # async synthesis in flight + hold_steps: int = 0 # HOLD_TOKEN steps emitted while it runs + # Engine output count once the emitted end token is accepted, and the + # shipped (audio, truncated, error). A preemption can drop the step that + # delivered both; the resumed request then delivers them again. + end_len: int | None = None + delivered: tuple[torch.Tensor, bool, bool] | None = None + + +SongWork = Callable[[list], Generator[None, None, torch.Tensor]] + + +class _SynthesisJob: + """One song's NAR + VAE pass on a side stream, fed a little per step. + + The pass is a generator of work units (the prefill, each ODE evaluation, + each VAE tile). ``advance`` enqueues units only while fewer than + ``MAX_INFLIGHT`` are queued ahead of the GPU and never waits: enqueuing + the whole song at once fills the stream's launch queue, after which every + launch blocks the host until the GPU drains, and the decode loop with it. + The decode steps call ``advance`` every step, so the side stream stays + busy while the other requests keep decoding; kernels of the two streams + run concurrently within the process. The finishing request emits + HOLD_TOKEN until ``done()``. + """ + + MAX_INFLIGHT = 2 + + def __init__(self, request_id: str, song: SongWork): + """``song(keep)`` yields once per enqueued work unit and returns the + device audio [2, samples]; what it appends to ``keep`` stays alive + until the job is released.""" + self.request_id = request_id + self.song = song + self.stream: torch.cuda.Stream | None = None + self.work: Iterator[None] | None = None + self.inflight: list[torch.cuda.Event] = [] + self.event: torch.cuda.Event | None = None # recorded after the last unit + self.audio: torch.Tensor | None = None # pinned host copy + self.finite: torch.Tensor | None = None # pinned pre-clamp validity flag + self.error: BaseException | None = None + self.cancelled = False + self._released = False + # Solver engines (captured graphs, K/V buffers), latents and audio: + # they must outlive the queued kernels that use them. + self.keep: list[CachedNAR | torch.Tensor] = [] + self.t_start: float | None = None + self.t_done: float | None = None + + @property + def started(self) -> bool: + return self.t_start is not None + + def _run(self) -> Iterator[None]: + audio = yield from self.song(self.keep) + # Check before clipping: clipping maps infinities to finite +/-1. + # Both copies target pinned allocations and complete before self.event. + self.finite = torch.empty((), dtype=torch.bool, pin_memory=True) + self.finite.copy_(torch.isfinite(audio).all(), non_blocking=True) + audio = audio.clamp(-1, 1) + # Pinned, so the device-to-host copy stays async on the side stream. + self.audio = torch.empty(audio.shape, dtype=audio.dtype, pin_memory=True) + self.audio.copy_(audio, non_blocking=True) + self.keep.append(audio) + + def start(self, stream: torch.cuda.Stream) -> None: + self.stream = stream + self.t_start = time.perf_counter() + self.work = self._run() + self.advance() + + def advance(self) -> None: + """Enqueue work units up to MAX_INFLIGHT ahead of the GPU; never waits.""" + if self.work is None or self.event is not None or self.error is not None: + return + with torch.cuda.stream(self.stream), torch.inference_mode(): + while True: + # Completed chunks need not keep all their K/V and private + # graph pools alive until the whole song finishes. + for item in self.keep[:]: + if isinstance(item, CachedNAR) and item.last_use is not None and item.last_use.query(): + item.close() + self.keep.remove(item) + while self.inflight and self.inflight[0].query(): + self.inflight.pop(0) + if len(self.inflight) >= self.MAX_INFLIGHT: + return + try: + next(self.work) + except StopIteration: + self.event = torch.cuda.Event() + self.event.record(self.stream) + self.inflight.clear() + return + except (RuntimeError, ValueError) as exc: + # CUDA errors and OOM are RuntimeErrors; song_chunks + # rejects a bad song with ValueError. Either fails only + # this request (see finish); anything else is a bug. + self.error = exc + return + unit = torch.cuda.Event() + unit.record(self.stream) + self.inflight.append(unit) + + def done(self) -> bool: + if self.cancelled or self.error is not None: + return True + self.advance() + finished = self.event is not None and self.event.query() + if finished and self.t_done is None: + self.t_done = time.perf_counter() + return finished + + def drain(self) -> None: + """Enqueue and wait for everything (a request out of hold budget).""" + while self.work is not None and self.event is None and self.error is None: + if self.inflight: + self.inflight[0].synchronize() + self.advance() + + def release(self) -> None: + """Release once, after this job's last submitted kernel has finished. + + The queue may already be running a later job on the same stream when + this job is delivered. Waiting on the stream would wait for that job. + """ + if self._released: + return + if self.stream is not None and self.event is None: + # An error or cancellation can precede the normal terminal event. + self.event = torch.cuda.Event() + self.event.record(self.stream) + if self.event is not None: + self.event.synchronize() + if self.work is not None: + self.work.close() + self.work = None + for item in self.keep: + if isinstance(item, CachedNAR): + item.close() + self.keep.clear() + self.stream = None + self._released = True + + def cancel(self) -> None: + """Stop submitting work and release the already submitted units. + + Waiting is bounded by MAX_INFLIGHT, not the rest of the song. Cleanup + happens here even if the engine becomes idle after the cancellation. + """ + self.cancelled = True + self.release() + + def finish(self) -> tuple[torch.Tensor, bool]: + """Wait if needed, release, and return (audio [2, T] on the CPU, failed).""" + self.drain() + self.release() + if self.t_done is None: + self.t_done = time.perf_counter() + if self.cancelled: + return torch.zeros((2, 0)), True + if self.error is not None: + logger.error("YuE2 synthesis failed for req=%s: %r", self.request_id, self.error) + return torch.zeros((2, 0)), True + if self.finite is None or not bool(self.finite): + logger.error("YuE2 synthesis produced non-finite audio for req=%s", self.request_id) + return torch.zeros((2, 0)), True + return self.audio, False + + +class _SynthesisQueue: + """Songs waiting for, or running on, the synthesis side stream. + + One song runs at a time: the NAR pass is compute-bound, so a second one + would only double the memory. ``pump`` runs every decode step; it feeds + the running song and starts the next once it is done. + """ + + def __init__(self) -> None: + self._stream: torch.cuda.Stream | None = None + self.graph_pool = _NARGraphPool() + self.active: _SynthesisJob | None = None + self.waiting: list[_SynthesisJob] = [] + + def stream(self) -> torch.cuda.Stream: + """The side stream, at the highest priority the device offers. + + Once several songs are in flight the synthesis pass, not the decode, + bounds throughput, and its kernels leave no room to co-schedule the + decode's anyway: they fill every SM. So synthesis goes first and the + decode fills the gaps. + """ + if self._stream is None: + self._stream = torch.cuda.Stream(priority=torch.cuda.Stream.priority_range()[1]) + return self._stream + + def submit(self, job: _SynthesisJob) -> None: + self.waiting.append(job) + self.pump() + + def pump(self) -> None: + """Feed the running song; start the next queued one once it is done.""" + if self.active is not None: + if not self.active.done(): # done() also feeds it + return + self.active.release() # frees its buffers; the audio stays + self.active = None + while self.waiting: + job = self.waiting.pop(0) + if job.error is not None: + continue # fails at completion, without touching the GPU + job.start(self.stream()) + self.active = job + return + + def complete(self, job: _SynthesisJob) -> tuple[torch.Tensor, bool, bool]: + """Finish ``job`` now, ahead of the queue if it has not started. + + Returns (audio [2, T] on the CPU, failed, whether it had to wait). + """ + if job.cancelled: + return torch.zeros((2, 0)), True, False + if job in self.waiting: + self.waiting.remove(job) + if not job.started and job.error is None: + if self.active is not None: + # Finish the running song first (its own request ships it). + self.active.drain() + self.active.release() + self.active = None + job.start(self.stream()) + self.active = job + waited = not job.done() + audio, failed = job.finish() + if self.active is job: + self.active = None + self.pump() + return audio, failed, waited + + def cancel(self, job: _SynthesisJob) -> None: + """Drop an aborted song and release its bounded in-flight work.""" + if job in self.waiting: + self.waiting.remove(job) + job.cancel() + if self.active is job: + self.active = None class Yue2ForCausalLM(nn.Module): @@ -1288,6 +1920,9 @@ class Yue2ForCausalLM(nn.Module): have_multimodal_outputs = True prefer_model_sampler = True has_postprocess = False + # Request state is created in prepare_runner_inputs, which runs on every + # step; forward() is skipped on a full CUDA graph replay. + accepts_runner_sampling_extra_args = True # The song is decoded in-model and shipped via make_omni_output; per-step # hidden states must NOT ride the pooler payload, or the output pipeline # remaps the accumulated "hidden" rows to the audio modality key and the @@ -1328,7 +1963,6 @@ class Yue2ForCausalLM(nn.Module): prefix=maybe_prefix(prefix, "lm_head"), ) self.logits_processor = LogitsProcessor(int(config.vocab_size)) - self.nar_layers = nn.ModuleList([_NARLayer(config) for _ in range(int(config.num_hidden_layers))]) self.vae2llm = nn.Linear(int(getattr(config, "latent_dim", LATENT_DIM)), hidden) self.llm2vae = nn.Linear(hidden, int(getattr(config, "latent_dim", LATENT_DIM))) @@ -1343,27 +1977,31 @@ class Yue2ForCausalLM(nn.Module): self.max_position_embeddings = int(config.max_position_embeddings) self.max_latent_frames = int(getattr(config, "max_latent_frames", self.max_position_embeddings)) self._timestep_shift = float(getattr(config, "timestep_shift", 1.0)) - - self._states: dict[str, _RequestState] = {} - self._row_constants: dict[str, _RowConstants] = {} - self._audio_queue: list[tuple[str, torch.Tensor, bool, bool]] = [] - self._deferred_cleanup_ids: set[str] = set() - self._step_rows: list[tuple[str, int, int]] = [] # (req_id, computed, scheduled) - self._decode_t0: dict[str, float] = {} # first decode-step wall time, for AR tok/s - self._last_mm: dict[str, Any] | None = None # current step's make_omni_output payload + # Fused NAR projection weights (fuse_nar_projections); see nar_weights. + self.nar_fused: list[tuple[torch.Tensor, ...]] | None = None self._vae: YuE2VAE | None = None self._vae_device: torch.device | None = None + self._states: dict[str, _RequestState] = {} + self._audio_queue: list[tuple[str, torch.Tensor, bool, bool]] = [] + self._step_rows: list[tuple[str, int, int]] = [] # (req_id, computed, scheduled) + self._step_discard: list[bool] | None = None + self._decode_t0: dict[str, float] = {} # first decode-step wall time, for AR tok/s + self._last_mm: dict[str, Any] | None = None # current step's make_omni_output payload + # Tokens drawn by the last sample() call, still on their way to the + # host: (copy-done event, pinned [2, rows] ids/flags, states). + self._pending: tuple[torch.Event, torch.Tensor, list[_RequestState]] | None = None + self._pinned: torch.Tensor | None = None + self._stage: torch.Tensor | None = None + self._logits: torch.Tensor | None = None # semantic-only logits buffer + self._sampler_graphs: dict[tuple, SamplerGraph] = {} + self._sampler_graph_pool: tuple[int, int] | None = None + self._sampler_capacity = int(vllm_config.scheduler_config.max_num_seqs) + self._max_model_len = int(vllm_config.model_config.max_model_len) + self._synthesis = _SynthesisQueue() + # ------------------------------------------------------------ sampling - def shift_t(self, raw_t: float, device: torch.device, dtype: torch.dtype) -> torch.Tensor: - t_sig = torch.sigmoid(torch.tensor(raw_t, dtype=dtype, device=device)) - shift = self._timestep_shift - return shift * t_sig / (1 + (shift - 1) * t_sig) - - def rotary(self, positions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - return self.rotary_emb(positions) - def ar_project(self, layer, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): """AR-path Q/K/V from the vLLM backbone's fused projection.""" attn = layer.self_attn @@ -1387,6 +2025,108 @@ class Yue2ForCausalLM(nn.Module): k = _apply_rotary(k.unsqueeze(0), rc, rs).squeeze(0) return q, k, v.view(T, num_kv_heads, head_dim) + def nar_weights(self) -> list[tuple[torch.Tensor, ...]]: + """The fused NAR layer weights, fused on first use. + + load_weights fuses them, but a dummy-weight load (load_format=dummy, + as the CI smoke tests run) never calls it. + """ + if self.nar_fused is None: + self.nar_fused = fuse_nar_projections(self.nar_layers) + return self.nar_fused + + def rotary(self, positions: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return self.rotary_emb(positions) + + def synthesize_song( + self, prefix_ids: Sequence[int], codec: Sequence[int], seed: int + ) -> tuple[torch.Tensor, float, float]: + """Channels-first stereo audio for one song, plus NAR and VAE seconds.""" + t0 = time.perf_counter() + latents = synthesize(self, prefix_ids, codec, seed, steps=ODE_STEPS) + logger.debug( + "YUE2_LATENTS frames=%d absmax=%.4f mean=%.5f std=%.4f", + latents.shape[0], + float(latents.abs().max()), + float(latents.mean()), + float(latents.std()), + ) + t1 = time.perf_counter() + audio = self._decode_latents(latents) + return audio, t1 - t0, time.perf_counter() - t1 + + @torch.inference_mode() + def _warm_up_synthesis(self) -> None: + """Compile the NAR kernels and touch the VAE once at startup. + + Otherwise the first finishing song pays the compile (seconds) inside + a decode step. One ODE step over a few frames covers every shape: + the NAR kernels are compiled for dynamic token counts. + """ + if not torch.cuda.is_available() or self.nar_fused is None: + return + t0 = time.perf_counter() + training, default_dtype = self.training, torch.get_default_dtype() + # load_weights runs before the loader switches to eval, and under the + # loader's bf16 default dtype, which the compiled kernels guard on: + # warm up in the serving-time state or the first song recompiles. + self.eval() + torch.set_default_dtype(torch.float32) + try: + latents = synthesize(self, [EOD, ABC_START, ABC_END, MUSIC_START], list(range(16)), DEFAULT_SEED, steps=1) + self._decode_latents(latents) + finally: + torch.set_default_dtype(default_dtype) + self.train(training) + logger.info("YuE2 synthesis warm-up took %.1f s", time.perf_counter() - t0) + + def _vae_model(self, device: torch.device) -> YuE2VAE: + if self._vae is None: + vae_path = os.environ.get("YUE2_VAE", DEFAULT_VAE_ID) + logger.info("Loading YuE2 VAE decoder from %s", vae_path) + # Keep the VAE out of the module tree (__dict__, not an nn.Module + # attribute): the checkpoint has no VAE keys, so a registered + # submodule would trip DefaultModelLoader.track_weights_loading, + # and a model-wide dtype/device pass would change VAE numerics. + self.__dict__["_vae"] = YuE2VAE.from_pretrained(vae_path, decoder_only=True, device="cpu") + if self._vae_device != device: + self._vae.to(device) + self._vae_device = device + return self._vae + + def _decode_latents_iter(self, latents: torch.Tensor) -> Iterator[None]: + """Device latents -> device stereo [2, samples] as a work generator + (one unit per VAE tile); nothing is read back to the host.""" + model = self._vae_model(latents.device) + audio = yield from model.decode_tiled_iter( + latents.T.unsqueeze(0), + core_frames=VAE_CORE_FRAMES, + halo_frames=VAE_HALO_FRAMES, + output_device=latents.device, + check_finite=False, + ) + # _SynthesisJob checks finite values before clipping and transferring. + return audio[0].contiguous() + + def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor: + """[frames, 64] FP32 CPU latents -> channels-first stereo [2, samples].""" + device = f"cuda:{torch.accelerator.current_device_index()}" if torch.cuda.is_available() else "cpu" + model = self._vae_model(torch.device(device)) + z = latents.T.unsqueeze(0) # [1, 64, T] + with torch.inference_mode(): + audio = model.decode_tiled( + z, + core_frames=VAE_CORE_FRAMES, + halo_frames=VAE_HALO_FRAMES, + output_device="cpu", + ) + if not torch.isfinite(audio).all(): + raise RuntimeError("VAE produced non-finite audio") + logger.debug("YUE2_AUDIO_PRECLAMP absmax=%.4f", float(audio.abs().max())) + # Channels-first [2, T], the same layout MiniMax Music 3 ships in + # model_outputs; serving's create_audio transposes it for soundfile. + return audio[0].float().clamp(-1, 1).contiguous() + # ------------------------------------------------------------ weights def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -1398,12 +2138,15 @@ class Yue2ForCausalLM(nn.Module): loader = AutoWeightsLoader(self) loaded = loader.load_weights(iter(ar_pairs)) + self.nar_fused = fuse_nar_projections(self.nar_layers) # Load the VAE now, not lazily at the first finishing request: a bad # path or a failed download must surface at startup, and the decoder # weights must sit on the GPU before vLLM's memory profiling sizes # the KV cache. device = f"cuda:{torch.accelerator.current_device_index()}" if torch.cuda.is_available() else "cpu" self._vae_model(torch.device(device)) + self._warm_up_synthesis() + self._warm_up_sampler() # The NAR/projection modules were populated by hand above; # DefaultModelLoader.track_weights_loading diffs every named # parameter against the returned set, so the side keys must be @@ -1423,19 +2166,28 @@ class Yue2ForCausalLM(nn.Module): num_scheduled_tokens: Any, input_ids: torch.Tensor, positions: torch.Tensor, + sampling_extra_args: list[dict[str, Any]] | None = None, + discard_mask: Sequence[bool] | None = None, **_: Any, ) -> tuple[torch.Tensor, torch.Tensor]: - """Bind rows to requests for the step; prompt tokens are captured in - ``forward``, where the state is created (see ``_capture_constants``). + """Bind rows to requests and create new request state for the step. + + This runs before the backbone on every step, including CUDA graph + replays, so request admission never depends on forward() running. Decode rows feed tokens the model itself sampled, so they never contribute to the prompt prefix.""" computed = [int(v) for v in num_computed_tokens] scheduled = [int(v) for v in num_scheduled_tokens] self._step_rows = [(str(rid), computed[i], scheduled[i]) for i, rid in enumerate(req_ids)] + # The runner's own verdict on which rows' samples it keeps: a partial + # prefill, including a recompute after preemption (prompt + output + # tokens replayed), must not draw from the request's generator. + self._step_discard = None if discard_mask is None else [bool(v) for v in discard_mask] # make_omni_output runs every step; without this reset the first step's # payload would stay live forever and the _audio_queue fallback in # _ship_audio could never trigger. self._last_mm = None + self._capture_constants(sampling_extra_args, input_ids) return input_ids, positions def forward( @@ -1446,8 +2198,7 @@ class Yue2ForCausalLM(nn.Module): inputs_embeds: torch.Tensor | None = None, **kwargs: Any, ) -> torch.Tensor: - del intermediate_tensors - self._capture_constants(kwargs, input_ids) + del intermediate_tensors, kwargs hidden = self.model( input_ids=input_ids if inputs_embeds is None else None, positions=positions, @@ -1455,10 +2206,9 @@ class Yue2ForCausalLM(nn.Module): ) if isinstance(hidden, tuple): hidden = hidden[0] - self._flush_deferred_cleanup() return hidden - def _capture_constants(self, kwargs: dict[str, Any], input_ids: torch.Tensor | None) -> None: + def _capture_constants(self, extra_args: list[dict[str, Any]] | None, input_ids: torch.Tensor | None) -> None: """Create per-request state on the request's FIRST scheduled step. With KV prefix caching the first step may arrive with ``comp > 0`` @@ -1469,7 +2219,6 @@ class Yue2ForCausalLM(nn.Module): truth for the NAR conditioning prefix and for the prompt length the sampler uses to tell a completing prefill row from a mid-chunk one. """ - extra_args = kwargs.get("sampling_extra_args") if extra_args is None or input_ids is None: return offset = 0 @@ -1514,15 +2263,14 @@ class Yue2ForCausalLM(nn.Module): max_audio_frames=int(args.get(KEY_MAX_AUDIO_FRAMES, preset["max_tokens"])), seed=int(seed), skip_synthesis=bool(args.get(KEY_SKIP_SYNTHESIS, phase == "abc")), + max_hold_steps=int(args.get(KEY_MAX_HOLD_STEPS, 0) or 0), ) device = input_ids.device generator = torch.Generator(device=device if device.type != "mps" else "cpu") generator.manual_seed(constants.seed) - self._row_constants[req_id] = constants self._states[req_id] = _RequestState( request_id=req_id, constants=constants, - prompt_tokens=prompt_len, prompt_len=prompt_len, prefix_ids=prefix_ids, generator=generator, @@ -1530,65 +2278,349 @@ class Yue2ForCausalLM(nn.Module): offset += span def compute_logits(self, hidden_states: torch.Tensor, sampling_metadata: Any = None) -> torch.Tensor: - return self.logits_processor(self.lm_head, hidden_states) + if not self._semantic_only_step(): + return self.logits_processor(self.lm_head, hidden_states) + # Every row samples the semantic phase, whose sampler reads only the + # MUSIC_END + codec columns (18% of the vocabulary): project onto + # those alone and leave the rest of the row at -inf. + rows = hidden_states.shape[0] + if self._logits is None or self._logits.shape[0] < rows or self._logits.device != hidden_states.device: + self._logits = torch.full( + (max(rows, 8), self.vocab_size), -torch.inf, dtype=hidden_states.dtype, device=hidden_states.device + ) + logits = self._logits[:rows] + lo, hi = MUSIC_END, CODEC_OFFSET + CODEC_SIZE + torch.matmul(hidden_states, self.lm_head.weight[lo:hi].T, out=logits[:, lo:hi]) + return logits + + def _semantic_only_step(self) -> bool: + """True when every request in the step is a semantic-phase request.""" + if not self._step_rows or not hasattr(self.lm_head, "weight"): + return False + for req_id, _, _ in self._step_rows: + state = self._states.get(req_id) + if state is None or state.constants.phase != "semantic": + return False + return True # ------------------------------------------------------------ sampler def sample(self, logits: torch.Tensor, sampling_metadata: Any) -> SamplerOutput | None: - """Draw one token per decoding row with the request's own arithmetic.""" + """Draw one token per decoding row with the request's own arithmetic. + + The draw stays on the device: its ids reach the host through an + async copy that the NEXT call resolves, so the host never waits for + the GPU inside a decode step and async scheduling can overlap the + next step's preparation with this step's kernels. Every draw still + sees the full history (the previous draw is resolved first); only + acting on an end token or a full frame budget moves one step later. + """ del sampling_metadata rows = int(logits.shape[0]) # The runner's input_ids buffer is int32; long ids crash its scatter # when concurrent requests re-index rows. - token_ids = torch.zeros((rows, 1), dtype=torch.int32, device=logits.device) + # Build fixed HOLD/end outputs on the host and transfer the whole batch + # once. Assigning a Python scalar to a CUDA row performs a blocking H2D + # copy. The pinned allocator tracks this copy's lifetime, so the buffer + # can be retired without waiting for it or rewriting an in-flight copy. + host_token_ids = torch.zeros((rows, 1), dtype=torch.int32, pin_memory=logits.is_cuda) + self._resolve_pending() + self._synthesis.pump() if not self._step_rows: - return SamplerOutput(sampled_token_ids=token_ids, logprobs_tensors=None) + return SamplerOutput( + sampled_token_ids=host_token_ids.to(logits.device, non_blocking=True), logprobs_tensors=None + ) + # Rows sharing a phase and preset are sampled as one batch. + groups: dict[tuple, list[tuple[int, _RequestState]]] = {} for row, (req_id, comp, span) in enumerate(self._step_rows): if row >= rows: break state = self._states.get(req_id) - if state is None or state.finished: + if state is None: continue - if comp + span < state.prompt_len: - # Mid-chunk prefill rows: their sampled token is discarded. - # (A completing prefill row reaches the full prompt length - # even when a prefix-cache hit left comp > 0.) + if self._step_discard is not None and row < len(self._step_discard): + discarded = self._step_discard[row] + else: + # A completing prefill row reaches the full prompt length even + # when a prefix-cache hit left comp > 0. + discarded = comp + span < state.prompt_len + if discarded: + continue + accepted = max(0, comp + span - state.prompt_len) + self._reconcile_history(state, accepted) + c = state.constants + end = ABC_END if c.phase == "abc" else MUSIC_END + if state.finished: + # Normally an async lookahead row the engine discards. If the + # scheduler dropped the step that emitted the end token (the + # request was preempted with it in flight), that step's audio + # was dropped too: deliver both again. + if state.end_len is not None and accepted < state.end_len: + if state.delivered is not None: + audio, truncated, error = state.delivered + self._ship_audio(state.request_id, audio, truncated, error=error) + state.end_len = accepted + 1 + host_token_ids[row, 0] = end + continue + if state.finish_ready: + state.finish_ready = False + if self._hold_allowed(state): + self._queue_synthesis(state) + host_token_ids[row, 0] = HOLD_TOKEN + continue + state.finished = True + state.end_len = accepted + 1 + self._finish_request_safely(state, hit_end=not state.truncated) + host_token_ids[row, 0] = end + continue + if state.job is not None: + if state.job.done() or not self._hold_allowed(state): + self._complete_synthesis(state) + state.end_len = accepted + 1 + host_token_ids[row, 0] = end + else: + state.hold_steps += 1 + host_token_ids[row, 0] = HOLD_TOKEN continue - constants = state.constants if state.request_id not in self._decode_t0: self._decode_t0[state.request_id] = time.perf_counter() - scores = distribution( - logits[row], - temperature=constants.temperature, - top_p=constants.top_p, - top_k=constants.top_k, - repetition_penalty=constants.repetition_penalty, - penalty_window=constants.penalty_window, - history=state.history, - step=len(state.history), - min_tokens=constants.min_tokens, - phase=constants.phase, + # With no penalty the window has no effect. Canonicalize it so + # equivalent groups cannot replay and overwrite the same graph's + # output buffers before their draws are collected below. + window = 0 if c.repetition_penalty == 1.0 else c.penalty_window + key = (c.phase, c.temperature, c.top_p, c.top_k, c.repetition_penalty, window) + groups.setdefault(key, []).append((row, state)) + token_ids = host_token_ids.to(logits.device, non_blocking=True) + if not groups: + return SamplerOutput(sampled_token_ids=token_ids, logprobs_tensors=None) + + # Stage every small host input of this step in one pinned buffer and + # move it with a single async copy: a pageable copy or a list index + # would make the host wait for the GPU. The previous step's copies + # have landed (_resolve_pending waited on its event), so the staging + # buffer can be rewritten. + plans = [] + flat: list[int] = [] + for (phase, temperature, top_p, top_k, penalty, window), members in groups.items(): + vocab = PHASE_VOCAB[phase] + # With no penalty there is no history input: [-0:] would walk the + # entire history every step before discarding it at width=0. + windows = ( + [] if penalty == 1.0 else [[vocab.col(t) for t in state.history[-window:]] for _, state in members] ) - token = sample_row(scores, state.generator, greedy=constants.temperature == 0) - end = ABC_END if constants.phase == "abc" else MUSIC_END - if token == end: - state.finished = True - state.truncated = False - self._finish_request_safely(state, hit_end=True) - token_ids[row, 0] = end + # A fixed width (the preset window) keeps the sampler graph static. + width = 0 if penalty == 1.0 else window if window > 0 else max(map(len, windows)) + block = [len(state.history) < state.constants.min_tokens for _, state in members] + # A request that still needs its song must not stop on its end + # token when drawn: the host only sees the draw next step. + end = ABC_END if phase == "abc" else MUSIC_END + held = [-1 if state.constants.skip_synthesis else end for _, state in members] + base = len(flat) + flat += [row for row, _ in members] + block + held + for cols in windows if width else (): + flat += cols + [vocab.num_cols] * (width - len(cols)) + plans.append((vocab, (temperature, top_p, top_k, penalty), members, base, width, any(block))) + if self._stage is None or self._stage.numel() < len(flat): + self._stage = torch.empty(max(len(flat), 1024), dtype=torch.long, pin_memory=True) + self._stage[: len(flat)].copy_(torch.as_tensor(flat)) + staged = self._stage[: len(flat)].to(logits.device, non_blocking=True) + + # Graphs share a pool and may replay in a different order from capture. + # Preserve each group's outputs outside that pool before replaying the + # next graph: its intermediates may alias the previous graph's outputs. + count = sum(len(members) for members in groups.values()) + drawn = torch.empty((2, count), dtype=torch.long, device=logits.device) + emitted: list[torch.Tensor] = [] + drawn_rows: list[int] = [] + states: list[_RequestState] = [] + for vocab, (temperature, top_p, top_k, penalty), members, base, width, blocks in plans: + n = len(members) + index = staged[base : base + n] + rows_logits = logits if n == rows else logits.index_select(0, index) + window_cols = staged[base + 3 * n : base + 3 * n + n * width].view(n, width) if width else None + block_end = staged[base + n : base + 2 * n].bool() if blocks else None + generators = [state.generator for _, state in members] + graph = self._sampler_graph(vocab, n, rows_logits, (temperature, top_p, top_k, penalty), width) + if graph is not None: + ids, bad = graph(rows_logits, window_cols, block_end, generators) else: - state.history.append(token) - budget = constants.max_audio_frames - if constants.phase == "semantic" and len(state.history) >= budget: - state.finished = True - state.truncated = True - self._finish_request_safely(state, hit_end=False) - token_ids[row, 0] = end - else: - token_ids[row, 0] = token + ids, bad = sample_rows( + rows_logits, + vocab, + temperature=temperature, + top_p=top_p, + top_k=top_k, + repetition_penalty=penalty, + window_cols=window_cols, + block_end=block_end, + generators=generators, + ) + offset = len(drawn_rows) + drawn[0, offset : offset + n].copy_(ids) + drawn[1, offset : offset + n].copy_(bad) + emitted.append(torch.where(ids == staged[base + 2 * n : base + 3 * n], HOLD_TOKEN, ids)) + drawn_rows += [row for row, _ in members] + states += [state for _, state in members] + out = torch.cat(emitted).to(torch.int32) + if drawn_rows == list(range(rows)): + token_ids[:, 0] = out + else: + index = torch.cat([staged[base : base + len(m)] for _, _, m, base, _, _ in plans]) + token_ids.view(-1).index_copy_(0, index, out) + + if self._pinned is None or self._pinned.shape[1] < count: + self._pinned = torch.empty((2, max(count, 64)), dtype=torch.long, pin_memory=True) + host = self._pinned[:, :count] + # Async into pinned memory; the next sample() call resolves it. + host.copy_(drawn, non_blocking=True) + event = torch.Event() + event.record() + self._pending = (event, host, states) return SamplerOutput(sampled_token_ids=token_ids, logprobs_tensors=None) + def _sampler_graph(self, vocab, rows, logits, preset, width) -> SamplerGraph | None: + """Use a power-of-two bucket; bound captures for custom sampling presets. + + The serving presets are captured before vLLM profiles memory. Other + presets run eagerly once the small custom-preset allowance is full. + """ + if logits.device.type != "cuda": + return None + bucket = 1 << (rows - 1).bit_length() + key = (vocab.phase, bucket, logits.dtype, preset, width) + graph = self._sampler_graphs.get(key) + if graph is None: + if len(self._sampler_graphs) >= 2 * self._sampler_capacity.bit_length() + 8: + return None + if self._sampler_graph_pool is None: + self._sampler_graph_pool = torch.cuda.graph_pool_handle() + graph = SamplerGraph( + vocab, bucket, logits.shape[1], logits.dtype, preset, width, logits.device, self._sampler_graph_pool + ) + self._sampler_graphs[key] = graph + return graph + + def _warm_up_sampler(self) -> None: + """Capture default phase/preset buckets before serving and KV profiling. + + Lazy captures use torch.cuda.graph, which synchronizes the device. + Admission and batch-size changes should replay pre-existing graphs. + """ + weight = self.lm_head.weight + if weight.device.type != "cuda": + return + bucket = 1 + while bucket < self._sampler_capacity * 2: + for phase, preset in (("semantic", SEMANTIC_SAMPLING), ("abc", ABC_SAMPLING)): + shape = (preset["temperature"], preset["top_p"], preset["top_k"], preset["repetition_penalty"]) + logits = torch.empty((bucket, self.vocab_size), dtype=weight.dtype, device=weight.device) + self._sampler_graph(PHASE_VOCAB[phase], bucket, logits, shape, preset["penalty_window"]) + if bucket >= self._sampler_capacity: + break + bucket *= 2 + + # ------------------------------------------------------------ async synthesis + + def _hold_allowed(self, state: _RequestState) -> bool: + """May this request emit one more HOLD_TOKEN step? + + Needs the caller's hold budget (and its max_tokens headroom) and room + in the context for the HOLD and the final end token. + """ + c = state.constants + if c.skip_synthesis or state.hold_steps >= c.max_hold_steps or not torch.cuda.is_available(): + return False + # prompt + frames + the drawn-end HOLD + holds so far + this HOLD + end + return state.prompt_len + len(state.history) + state.hold_steps + 3 <= self._max_model_len + + def _song_iter( + self, prefix_ids: list[int], codec: list[int], seed: int, keep: list + ) -> Generator[None, None, torch.Tensor]: + """NAR + VAE for one song as work units; returns the device audio.""" + latents = yield from synthesize_iter(self, prefix_ids, codec, seed, keep, self._synthesis.graph_pool) + keep.append(latents) + return (yield from self._decode_latents_iter(latents)) + + def _queue_synthesis(self, state: _RequestState) -> None: + codec = [t - CODEC_OFFSET for t in state.history] + prefix_ids, seed = state.prefix_ids, state.constants.seed + job = _SynthesisJob(state.request_id, lambda keep: self._song_iter(prefix_ids, codec, seed, keep)) + if not codec or min(codec) < 0 or max(codec) >= CODEC_SIZE: + job.error = RuntimeError("semantic history is empty or contains non-codec tokens") + state.job = job + state.hold_steps = 1 + self._synthesis.submit(job) + + def _complete_synthesis(self, state: _RequestState) -> None: + """Ship the song on this step (waiting for it if the hold ran out).""" + job = state.job + assert job is not None + audio, failed, waited = self._synthesis.complete(job) + state.job = None + state.finished = True + t0 = self._decode_t0.pop(state.request_id, None) + t_ar = (job.t_start - t0) if t0 is not None and job.t_start is not None else float("nan") + logger.debug( + "YUE2_TIMING req=%s phase=%s ar_tokens=%d ar_s=%.2f ar_tps=%.1f synth_s=%.2f hold_steps=%d waited=%s", + state.request_id, + state.constants.phase, + len(state.history), + t_ar, + len(state.history) / t_ar if t_ar else float("nan"), + (job.t_done - job.t_start) if job.t_start is not None and job.t_done is not None else 0.0, + state.hold_steps, + waited, + ) + self._ship_audio(state.request_id, audio, state.truncated, error=failed) + + def _reconcile_history(self, state: _RequestState, accepted: int) -> None: + """Drop in-flight draws the scheduler discarded on preemption. + + Only a non-discarded row has replayed the complete accepted sequence; + partial recompute chunks cannot establish this boundary. A drawn EOS + occupies one engine token even though it is not in the codec history. + """ + if accepted >= len(state.history) + int(state.end_drawn): + return + if state.job is not None: + self._synthesis.cancel(state.job) + state.job = None + del state.history[accepted:] + state.end_drawn = False + state.finished = False + state.end_len = None + state.delivered = None + state.hold_steps = 0 + state.truncated = state.constants.phase == "semantic" and len(state.history) >= state.constants.max_audio_frames + state.finish_ready = state.truncated + + def _resolve_pending(self) -> None: + """Apply the previous step's draws to the host-side request state.""" + if self._pending is None: + return + event, host, states = self._pending + self._pending = None + event.synchronize() + ids, bad = host.tolist() + if any(bad): + raise RuntimeError("sampling scores are all -inf; check phase masking") + for state, token in zip(states, ids): + constants = state.constants + if token == (ABC_END if constants.phase == "abc" else MUSIC_END): + state.end_drawn = True + state.truncated = False + if constants.skip_synthesis: + # Emitted as drawn: the engine has already stopped it. + state.finished = True + else: + state.finish_ready = True + continue + state.history.append(token) + if constants.phase == "semantic" and len(state.history) >= constants.max_audio_frames: + state.truncated = True + state.finish_ready = True + # ------------------------------------------------------------ audio def _finish_request_safely(self, state: _RequestState, *, hit_end: bool) -> None: @@ -1602,7 +2634,9 @@ class Yue2ForCausalLM(nn.Module): """ try: self._finish_request(state, hit_end=hit_end) - except Exception: + except (RuntimeError, ValueError, FloatingPointError): + # CPU synthesis also reports non-finite latents as FloatingPointError. + # Programming errors propagate on both CUDA completion paths. logger.exception("YuE2 synthesis failed for req=%s; failing only this request", state.request_id) self._ship_audio(state.request_id, torch.zeros((2, 0)), state.truncated, error=True) @@ -1611,25 +2645,22 @@ class Yue2ForCausalLM(nn.Module): constants = state.constants if constants.skip_synthesis or not state.history: return + if torch.cuda.is_available(): + # No hold headroom means wait now, but still use the serialized + # synthesis stream and its bounded in-flight work. Independent sync + # graphs would overlap a running job and accumulate private pools. + self._queue_synthesis(state) + state.hold_steps = 0 + self._complete_synthesis(state) + return codec = [t - CODEC_OFFSET for t in state.history] if min(codec) < 0 or max(codec) >= CODEC_SIZE: raise RuntimeError("semantic history contains non-codec tokens") t0 = self._decode_t0.pop(state.request_id, None) t_nar0 = time.perf_counter() - latents = synthesize(self, state.prefix_ids, codec, constants.seed, steps=ODE_STEPS) - logger.info( - "YUE2_LATENTS frames=%d absmax=%.4f mean=%.5f std=%.4f", - latents.shape[0], - float(latents.abs().max()), - float(latents.mean()), - float(latents.std()), - ) - t_nar = time.perf_counter() - t_nar0 - t_vae0 = time.perf_counter() - audio = self._decode_latents(latents) - t_vae = time.perf_counter() - t_vae0 + audio, t_nar, t_vae = self.synthesize_song(state.prefix_ids, codec, constants.seed) t_ar = (t_nar0 - t0) if t0 is not None else float("nan") - logger.info( + logger.debug( "YUE2_TIMING req=%s phase=%s ar_tokens=%d ar_s=%.2f ar_tps=%.1f nar_s=%.2f vae_s=%.2f", state.request_id, constants.phase, @@ -1651,7 +2682,12 @@ class Yue2ForCausalLM(nn.Module): sparse routing on the request's last step in the output batch. Leaving the payload for a later step would drop it: the request is gone by then (gepard_talker flushes at its last step for the same reason). + The request keeps a reference until the engine finishes it, in case + the scheduler drops this step (see ``_RequestState.end_len``). """ + state = self._states.get(req_id) + if state is not None: + state.delivered = (audio, truncated, error) mm = self._last_mm if mm is None: # No forward built a payload this step; fall back to the queue. @@ -1666,39 +2702,6 @@ class Yue2ForCausalLM(nn.Module): meta["truncated"].append(int(truncated)) meta["error"].append(int(error)) - def _vae_model(self, device: torch.device) -> YuE2VAE: - if self._vae is None: - vae_path = os.environ.get("YUE2_VAE", DEFAULT_VAE_ID) - logger.info("Loading YuE2 VAE decoder from %s", vae_path) - # Keep the VAE out of the module tree (__dict__, not an nn.Module - # attribute): the checkpoint has no VAE keys, so a registered - # submodule would trip DefaultModelLoader.track_weights_loading, - # and a model-wide dtype/device pass would change VAE numerics. - self.__dict__["_vae"] = YuE2VAE.from_pretrained(vae_path, decoder_only=True, device="cpu") - if self._vae_device != device: - self._vae.to(device) - self._vae_device = device - return self._vae - - def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor: - """[frames, 64] FP32 CPU latents -> channels-first stereo [2, samples].""" - device = f"cuda:{torch.accelerator.current_device_index()}" if torch.cuda.is_available() else "cpu" - model = self._vae_model(torch.device(device)) - z = latents.T.unsqueeze(0) # [1, 64, T] - with torch.inference_mode(): - audio = model.decode_tiled( - z, - core_frames=VAE_CORE_FRAMES, - halo_frames=VAE_HALO_FRAMES, - output_device="cpu", - ) - if not torch.isfinite(audio).all(): - raise RuntimeError("VAE produced non-finite audio") - logger.info("YUE2_AUDIO_PRECLAMP absmax=%.4f", float(audio.abs().max())) - # Channels-first [2, T], the same layout MiniMax Music 3 ships in - # model_outputs; serving's create_audio transposes it for soundfile. - return audio[0].float().clamp(-1, 1).contiguous() - def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: if isinstance(model_outputs, OmniOutput): return model_outputs @@ -1733,18 +2736,20 @@ class Yue2ForCausalLM(nn.Module): return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs=mm) def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None: - # Fires before forward; defer the free so the in-flight step can read. - for rid in finished_req_ids: - self._deferred_cleanup_ids.add(_request_key(rid)) - - def _flush_deferred_cleanup(self) -> None: - for req_id in self._deferred_cleanup_ids: + # Runs at the start of execute_model, after the previous step's + # sample(), so the state can be freed right away. + finished = {_request_key(rid) for rid in finished_req_ids} + for req_id in finished: state = self._states.pop(req_id, None) + if state is not None and state.job is not None: + # Aborted mid-synthesis. + self._synthesis.cancel(state.job) + state.job = None if state is not None and not state.finished: logger.warning( "YuE2 request %s ended without an end token or its frame " "budget (aborted or preempted); no audio was produced.", req_id, ) - self._row_constants.pop(req_id, None) - self._deferred_cleanup_ids.clear() + self._decode_t0.pop(req_id, None) + self._step_rows = [row for row in self._step_rows if row[0] not in finished] diff --git a/vllm_omni/tokenizers/yue2_prompt.py b/vllm_omni/tokenizers/yue2_prompt.py index 97d7b7dea..f1baa4122 100644 --- a/vllm_omni/tokenizers/yue2_prompt.py +++ b/vllm_omni/tokenizers/yue2_prompt.py @@ -16,6 +16,7 @@ from vllm_omni.model_executor.models.yue2.yue2 import ( ABC_START, CODEC_OFFSET, EOD, + HOLD_TOKEN, MUSIC_END, MUSIC_START, ) @@ -84,6 +85,10 @@ def semantic_frames(generated: list[int]) -> list[int]: ids = list(generated) if ids and ids[-1] == MUSIC_END: ids.pop() + # The model emits HOLD_TOKEN on the step it draws MUSIC_END (it acts on + # its draws one step late) and while the song is synthesized. + while ids and ids[-1] == HOLD_TOKEN: + ids.pop() return [t - CODEC_OFFSET for t in ids] diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 330355c64..63dde07c2 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -893,6 +893,44 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex merged[req_id] = data return merged + def _call_prepare_runner_inputs( + self, + prepare_runner_inputs: Callable[..., tuple[torch.Tensor, torch.Tensor]], + *, + req_ids: list[str], + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None, + num_computed_tokens: np.ndarray, + num_scheduled_tokens: np.ndarray, + input_ids_buffer: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Keep opt-in request metadata outside captured model forwards.""" + if getattr(self.model, "accepts_runner_sampling_extra_args", False): + return prepare_runner_inputs( + input_ids=input_ids, + positions=positions, + inputs_embeds=inputs_embeds, + req_ids=req_ids, + num_computed_tokens=num_computed_tokens, + num_scheduled_tokens=num_scheduled_tokens, + input_ids_buffer=input_ids_buffer, + sampling_extra_args=[ + params.extra_args if params and params.extra_args else {} + for params in (self.requests[rid].sampling_params for rid in req_ids) + ], + discard_mask=self.discard_request_mask.np[: len(req_ids)], + ) + return prepare_runner_inputs( + input_ids=input_ids, + positions=positions, + inputs_embeds=inputs_embeds, + req_ids=req_ids, + num_computed_tokens=num_computed_tokens, + num_scheduled_tokens=num_scheduled_tokens, + input_ids_buffer=input_ids_buffer, + ) + @torch.inference_mode() def execute_model( self, @@ -1150,11 +1188,12 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex # for multimodal position detection, fix decode position offsets). prepare_runner_inputs = getattr(self.model, "prepare_runner_inputs", None) if callable(prepare_runner_inputs): - input_ids, positions = prepare_runner_inputs( + input_ids, positions = self._call_prepare_runner_inputs( + prepare_runner_inputs, + req_ids=req_ids[:num_reqs], input_ids=input_ids, positions=positions, inputs_embeds=inputs_embeds, - req_ids=req_ids[:num_reqs], num_computed_tokens=self.input_batch.num_computed_tokens_cpu[:num_reqs], num_scheduled_tokens=num_scheduled_tokens_np[:num_reqs], input_ids_buffer=self.input_ids.gpu[:num_tokens_padded], From e372aa3a4137bca513827b37e6080b908cd977c7 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Fri, 2 Oct 2026 07:11:04 -0700 Subject: [PATCH 14/22] [Bugfix] Restore Fish Speech and CosyVoice3 audio quality at high concurrency (#8422) Signed-off-by: Sy03 <1370724210@qq.com> --- .buildkite/cuda/test-ready.yml | 10 ++ .../test_fish_speech_async_chunk.py | 79 ++++++++++++++- .../worker_v2/test_capture_query_bound_gpu.py | 1 + .../worker_v2/test_decode_query_bound_gpu.py | 99 +++++++++++++++++++ tests/worker_v2/test_omni_model_state.py | 33 +++++++ .../stage_input_processors/fish_speech.py | 85 ++++++++-------- .../model_states/omni_model_state.py | 10 +- 7 files changed, 270 insertions(+), 47 deletions(-) create mode 100644 tests/worker_v2/test_decode_query_bound_gpu.py diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index ee93db4f5..8f3ca6bdc 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.buildkite/cuda/test-ready.yml @@ -25,6 +25,16 @@ steps: - "pytest -sv tests/model_executor/models/cosyvoice3/test_cosyvoice3_full_batch.py -k packed -m 'core_model and cuda' --run-level=core_model" mirror_hardwares: h100_1 + - label: "Simple · MRV2 FA3 Graph Bounds Test" + timeout_in_minutes: 15 + source_file_dependencies: + - tests/worker_v2/test_capture_query_bound_gpu.py + - tests/worker_v2/test_decode_query_bound_gpu.py + - vllm_omni/worker_v2/model_states/omni_model_state.py + commands: + - "pytest -sv tests/worker_v2/test_capture_query_bound_gpu.py tests/worker_v2/test_decode_query_bound_gpu.py -m 'core_model and cuda and H100' --run-level=core_model" + mirror_hardwares: h100_1 + - label: "Simple · Diffusion Test" timeout_in_minutes: 30 commands: diff --git a/tests/model_executor/stage_input_processors/test_fish_speech_async_chunk.py b/tests/model_executor/stage_input_processors/test_fish_speech_async_chunk.py index 132615b14..c6a68c898 100644 --- a/tests/model_executor/stage_input_processors/test_fish_speech_async_chunk.py +++ b/tests/model_executor/stage_input_processors/test_fish_speech_async_chunk.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from collections import defaultdict from types import SimpleNamespace @@ -18,7 +21,7 @@ def _make_transfer_manager( left_context: int = 1, initial_chunk_frames: int = 0, single_initial_chunk: bool = False, - backlog_chunk_frames: int = 0, + backlog_chunk_frames: int | str = 0, backlog_load_threshold: float = 0.75, max_num_seqs: int = 1, ): @@ -227,3 +230,77 @@ def test_invalid_backlog_config_fails_loudly(): {"audio_codes": torch.tensor([[1, 2, 3]])}, request, ) + + +@pytest.mark.parametrize("single_initial_chunk", [False, True]) +@pytest.mark.parametrize("initial_chunk_frames", [0, 4, 10]) +@pytest.mark.parametrize("tensor_payload", [False, True]) +@pytest.mark.parametrize("pattern", ["steady", "drop_at_finish", "rise_midstream", "oscillate"]) +def test_backlog_changes_preserve_every_frame_once(single_initial_chunk, initial_chunk_frames, tensor_payload, pattern): + tm = _make_transfer_manager( + tensor_payload=tensor_payload, + chunk_frames=25, + left_context=25, + initial_chunk_frames=initial_chunk_frames, + single_initial_chunk=single_initial_chunk, + backlog_chunk_frames=50, + max_num_seqs=64, + ) + request = _make_request() + emitted = [] + for frame in range(1, 141): + high = ( + pattern == "steady" + or (pattern == "drop_at_finish" and frame < 140) + or (pattern == "rise_midstream" and frame >= 37) + or (pattern == "oscillate" and (frame // 13) % 2 == 0) + ) + others = 63 if high else 1 + tm.code_prompt_token_ids = defaultdict(list, {"req": tm.code_prompt_token_ids["req"]}) + for idx in range(others): + tm.code_prompt_token_ids[f"other-{idx}"] = [None] * 50 + out = slow_ar_to_dac_decoder_async_chunk( + tm, + {"audio_codes": torch.tensor([[frame, frame + 1000]]), "audio_code_valid": torch.tensor([True])}, + request, + ) + if out is not None: + codes = out.codes.audio.reshape(2, -1) + emitted.extend(codes[0, out.meta.left_context_size :].tolist()) + tm.put_req_chunk["req"] += 1 + # EOS may arrive in a separate callback after the last valid frame. + out = slow_ar_to_dac_decoder_async_chunk(tm, {}, request, is_finished=True) + if out is not None and out.codes.audio.numel(): + emitted.extend(out.codes.audio.reshape(2, -1)[0, out.meta.left_context_size :].tolist()) + assert emitted == list(range(1, 141)) + assert slow_ar_to_dac_decoder_async_chunk(tm, {}, request, is_finished=True) is None + + +def test_final_backlog_drop_flushes_entire_pending_tail(): + tm = _make_transfer_manager( + tensor_payload=True, + chunk_frames=25, + left_context=25, + initial_chunk_frames=4, + single_initial_chunk=True, + backlog_chunk_frames=50, + max_num_seqs=64, + ) + request = _make_request() + emitted = [] + for frame in range(1, 41): + if frame < 40: + for idx in range(63): + tm.code_prompt_token_ids[f"other-{idx}"] = [None] * 30 + else: + for idx in range(32, 63): + tm.code_prompt_token_ids.pop(f"other-{idx}") + out = slow_ar_to_dac_decoder_async_chunk( + tm, + {"audio_codes": torch.tensor([[frame]])}, + request, + is_finished=frame == 40, + ) + if out is not None: + emitted.extend(out.codes.audio[0, out.meta.left_context_size :].tolist()) + assert emitted == list(range(1, 41)) diff --git a/tests/worker_v2/test_capture_query_bound_gpu.py b/tests/worker_v2/test_capture_query_bound_gpu.py index 93c53f313..c341a575b 100644 --- a/tests/worker_v2/test_capture_query_bound_gpu.py +++ b/tests/worker_v2/test_capture_query_bound_gpu.py @@ -44,6 +44,7 @@ def _check_prefill(): spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=128, dtype=torch.bfloat16) kvconfig = KVCacheConfig(num_blocks=1024, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["attn"], spec)]) state = OmniModelState.__new__(OmniModelState) + state.vllm_config = c state.max_model_len = 256 state.supports_mm_inputs = False with set_current_vllm_config(c): diff --git a/tests/worker_v2/test_decode_query_bound_gpu.py b/tests/worker_v2/test_decode_query_bound_gpu.py new file mode 100644 index 000000000..114e9f343 --- /dev/null +++ b/tests/worker_v2/test_decode_query_bound_gpu.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Real FA3 graph replay must preserve the decode GQA layout.""" + +from dataclasses import dataclass, replace + +import numpy as np +import pytest +import torch +from transformers import Qwen2Config +from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config +from vllm.config.compilation import CUDAGraphMode +from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend, FlashAttentionImpl +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers +from vllm.v1.worker.utils import AttentionGroup + +from tests.helpers.mark import hardware_test +from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState + +pytestmark = [pytest.mark.core_model] + + +@dataclass +class _AttentionScales: + _q_scale: torch.Tensor + _k_scale: torch.Tensor + _v_scale: torch.Tensor + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize("graph_mode", [CUDAGraphMode.FULL_AND_PIECEWISE, CUDAGraphMode.FULL_DECODE_ONLY]) +def test_full_decode_graph_matches_fa3_aot_schedule(graph_mode, tmp_path): + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9: + pytest.skip("requires FA3 on Hopper or newer") + torch.manual_seed(17) + device = torch.device("cuda:0") + c = VllmConfig() + Qwen2Config( + architectures=["Qwen2ForCausalLM"], + hidden_size=1024, + intermediate_size=2048, + num_hidden_layers=1, + num_attention_heads=16, + num_key_value_heads=2, + max_position_embeddings=2048, + ).save_pretrained(tmp_path) + c.model_config = ModelConfig(model=str(tmp_path), skip_tokenizer_init=True, dtype="bfloat16") + c.scheduler_config.max_num_seqs = 64 + c.compilation_config.cudagraph_mode = graph_mode + c.compilation_config.max_cudagraph_capture_size = 128 + c.attention_config.flash_attn_version = 3 + spec = FullAttentionSpec(block_size=16, num_kv_heads=2, head_size=64, dtype=torch.bfloat16) + kvconfig = KVCacheConfig(num_blocks=1024, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["attn"], spec)]) + state = OmniModelState.__new__(OmniModelState) + state.vllm_config = c + state.max_model_len = 2048 + state.supports_mm_inputs = False + with set_current_vllm_config(c): + group = AttentionGroup(FlashAttentionBackend, ["attn"], spec, 0) + group.create_metadata_builders(c, device) + impl = FlashAttentionImpl(16, 64, 64**-0.5, 2, None, None, "auto") + assert impl.vllm_flash_attn_version == 3 + buffers = InputBuffers(64, 128, device) + batch = InputBatch.make_dummy(8, 8, buffers, max_query_len=None) + tables = torch.arange(1024, dtype=torch.int32, device=device).reshape(8, 128) + slots = torch.zeros((1, 8), dtype=torch.int64, device=device) + cache = torch.randn(1024, 2, 16, 128, device=device, dtype=torch.bfloat16) + q = torch.randn(8, 16, 64, device=device, dtype=torch.bfloat16) + k = v = torch.empty(8, 2, 64, device=device, dtype=torch.bfloat16) + layer = _AttentionScales( + _q_scale=torch.ones((), device=device), + _k_scale=torch.ones((), device=device), + _v_scale=torch.ones((), device=device), + ) + md = state.prepare_attn(batch, CUDAGraphMode.FULL, (tables,), slots, [[group]], kvconfig, for_capture=True)[ + "attn" + ] + out = torch.full_like(q, float("nan")) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + impl.forward(layer, q, k, v, cache, md, out) + stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + impl.forward(layer, q, k, v, cache, md, out) + lengths = np.arange(37, 45, dtype=np.int32) + batch.seq_lens.copy_(torch.tensor(lengths, device=device)) + runtime = replace(batch, max_query_len=1, seq_lens_cpu_upper_bound=torch.tensor(lengths)) + eager_md = state.prepare_attn(runtime, CUDAGraphMode.FULL, (tables,), slots, [[group]], kvconfig)["attn"] + expected = torch.empty_like(q) + impl.forward(layer, q, k, v, cache, eager_md, expected) + state.prepare_attn(runtime, CUDAGraphMode.FULL, (tables,), slots, [[group]], kvconfig) + graph.replay() + torch.accelerator.synchronize() + assert md.max_query_len == 1 + torch.testing.assert_close(out, expected, atol=0.002, rtol=0.002) diff --git a/tests/worker_v2/test_omni_model_state.py b/tests/worker_v2/test_omni_model_state.py index b36dbb139..1cbef67e8 100644 --- a/tests/worker_v2/test_omni_model_state.py +++ b/tests/worker_v2/test_omni_model_state.py @@ -4,13 +4,17 @@ graph/eager paths with per-request seed independence, async snapshot ownership.""" from contextlib import nullcontext +from dataclasses import replace from types import SimpleNamespace from unittest.mock import MagicMock, patch import numpy as np import pytest import torch +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers +from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner from vllm.v1.worker.gpu.model_states.default import DefaultModelState +from vllm.v1.worker.gpu.states import RequestState from vllm_omni.model_executor.models.output_templates import OmniOutput, OwnedBatchTensor from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState, _make_safe_get_rope @@ -661,3 +665,32 @@ def test_identity_preprocess_skips_only_decode_rows(prefilling): state.run_preprocess(batch, {"input_ids": torch.tensor([1]), "inputs_embeds": embeds}) assert seen == (["r1"] if prefilling else []) assert torch.equal(embeds, torch.ones(1, 4)) + + +def test_mm_embeddings_exclude_zero_length_graph_padding_rows(): + state = object.__new__(OmniModelState) + state.supports_mm_inputs = True + state.mm_pruner = None + state.prompt_embeds_state = None + state.encoder_runner = MagicMock(spec=EncoderRunner, inputs_embeds=torch.zeros(8, 2)) + state.execute_mm_encoder = lambda _: None + state.gather_mm_embeddings = lambda _: ([torch.ones(1, 2)], torch.ones(5, dtype=torch.bool)) + + def embed(input_ids, *, query_start_loc, multimodal_embeddings, is_multimodal): + assert query_start_loc == [0, 2, 5] + return input_ids[:, None].expand(-1, 2).float() + + state.model = MagicMock(supports_embed_input_ids_query_start_loc=True, embed_input_ids=embed) + buffers = InputBuffers(4, 8, torch.device("cpu")) + batch = replace( + InputBatch.make_dummy(2, 5, buffers), + num_reqs=2, + num_reqs_after_padding=4, + num_tokens=5, + num_tokens_after_padding=8, + input_ids=torch.arange(8), + query_start_loc_np=np.array([0, 2, 5, 5, 5], dtype=np.int32), + ) + result = state.prepare_inputs_embeds({}, batch, MagicMock(spec=RequestState)) + torch.testing.assert_close(result[:5], torch.arange(5).float()[:, None].expand(-1, 2)) + assert result.shape == (8, 2) diff --git a/vllm_omni/model_executor/stage_input_processors/fish_speech.py b/vllm_omni/model_executor/stage_input_processors/fish_speech.py index f428c34a4..cfcb277ea 100644 --- a/vllm_omni/model_executor/stage_input_processors/fish_speech.py +++ b/vllm_omni/model_executor/stage_input_processors/fish_speech.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Stage input processor for Fish Speech S2 Pro: Slow AR → DAC Decoder.""" from collections.abc import Mapping @@ -122,9 +125,9 @@ def slow_ar_to_dac_decoder_async_chunk( ) -> OmniPayloadStruct | None: """Async streaming processor: emit code chunks as they are produced. - Accumulates per-step codes and emits fixed-size chunks with left context - overlap for smooth audio transitions, analogous to - ``talker2code2wav_async_chunk`` in Qwen3 TTS. + Accumulates per-step codes and emits contiguous chunks with left context + overlap. The per-request output frontier remains absolute when the backlog + policy changes the chunk size; a final callback flushes every pending frame. """ request_id = request.external_req_id finished = bool(is_finished or request.is_finished()) @@ -174,49 +177,41 @@ def slow_ar_to_dac_decoder_async_chunk( ) return None - single_initial_chunk = _cfg_bool( - cfg, - "fish_speech_single_initial_chunk", - False, + # Keep the absolute output frontier in the request payload so that changing + # the backlog chunk size cannot reinterpret frames already emitted. Both + # connector runtimes reclaim request_payload on completion/cancellation. + request_payload = getattr(transfer_manager, "request_payload", None) + if request_payload is None: + request_payload = {} + transfer_manager.request_payload = request_payload + state = request_payload.setdefault(request_id, {}).setdefault( + "_fish_speech_async_state", {"emitted_frames": 0, "terminal_sent": False} ) - use_first_chunk = initial_chunk_size > 0 and initial_chunk_size < steady_chunk_size + if state["terminal_sent"]: + return None + emitted_frames = state["emitted_frames"] + pending = length - emitted_frames + if pending <= 0: + if not finished: + return None + state["terminal_sent"] = True + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), + meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)), + ) - if single_initial_chunk and use_first_chunk: - if length <= initial_chunk_size: - if not finished and length < initial_chunk_size: - return None - context_length = length if finished and length < initial_chunk_size else initial_chunk_size - else: - adjusted = length - initial_chunk_size - if adjusted <= 0: - return None - if not finished and adjusted % steady_chunk_size != 0: - return None - chunk_length = adjusted % steady_chunk_size - context_length = chunk_length if chunk_length != 0 else steady_chunk_size - end_index = min(length, left_context_size_config + context_length) - left_context_size = max(0, int(end_index - context_length)) - window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:] - elif initial_chunk_size > 0 and length <= chunk_size: - already_sent = transfer_manager.put_req_chunk[request_id] * initial_chunk_size - pending = length - already_sent - if pending <= 0: - return None - if pending < initial_chunk_size and not finished: - return None - context_length = min(pending, initial_chunk_size) - left_context_size = max(0, length - context_length) - window_frames = transfer_manager.code_prompt_token_ids[request_id][:length] - else: - initial_coverage = (chunk_size // initial_chunk_size) * initial_chunk_size if initial_chunk_size > 0 else 0 - adjusted = length - initial_coverage - chunk_length = adjusted % steady_chunk_size - if chunk_length != 0 and not finished: - return None - context_length = chunk_length if chunk_length != 0 else steady_chunk_size - end_index = min(length, left_context_size_config + context_length) - left_context_size = max(0, int(end_index - context_length)) - window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:] + single_initial_chunk = _cfg_bool(cfg, "fish_speech_single_initial_chunk", False) + initial_coverage = ( + (initial_chunk_size if single_initial_chunk else (chunk_size // initial_chunk_size) * initial_chunk_size) + if initial_chunk_size > 0 + else 0 + ) + target_chunk_size = initial_chunk_size if emitted_frames < initial_coverage else steady_chunk_size + if not finished and pending < target_chunk_size: + return None + + left_context_size = min(emitted_frames, left_context_size_config) + window_frames = transfer_manager.code_prompt_token_ids[request_id][emitted_frames - left_context_size : length] # Pack into codebook-major codes. The tensor path avoids expanding codec # indices into Python ints across the connector boundary; Stage1 schedules @@ -228,6 +223,8 @@ def slow_ar_to_dac_decoder_async_chunk( else: code_predictor_codes = codes_qf.reshape(-1) + state["emitted_frames"] = length + state["terminal_sent"] = finished return OmniPayloadStruct( codes=CodesStruct(audio=code_predictor_codes), meta=MetaStruct( diff --git a/vllm_omni/worker_v2/model_states/omni_model_state.py b/vllm_omni/worker_v2/model_states/omni_model_state.py index f10149f46..34c3c1cdf 100644 --- a/vllm_omni/worker_v2/model_states/omni_model_state.py +++ b/vllm_omni/worker_v2/model_states/omni_model_state.py @@ -448,12 +448,18 @@ class OmniModelState(DefaultModelState): for_capture: bool = False, ubatch_idx: int = 0, ) -> dict[str, Any]: - if for_capture and input_batch.max_query_len is None: + if ( + for_capture + and input_batch.max_query_len is None + and not self.vllm_config.compilation_config.cudagraph_mode.separate_routine() + ): # vLLM 0.30 distributes dummy tokens evenly across requests. For # an unconstrained FULL graph that split is not a query-length # bound: replay may put the entire token bucket in one request. # Attention launch parameters are fixed at capture, so use the # bucket's worst-case query length, not the dummy per-row length. + # Separate decode graphs must retain their uniform query bound: + # FA3's captured GQA layout must match the runtime AOT schedule. # Keep explicit bounds (e.g. varlen decode) and runtime metadata. input_batch = replace(input_batch, max_query_len=input_batch.num_tokens) return super().prepare_attn( @@ -485,7 +491,7 @@ class OmniModelState(DefaultModelState): mm_embeds, is_mm_embed = self.gather_mm_embeddings(input_batch) kwargs: dict[str, Any] = {"multimodal_embeddings": mm_embeds, "is_multimodal": is_mm_embed} if mm_embeds: - kwargs["query_start_loc"] = input_batch.query_start_loc_np.tolist() + kwargs["query_start_loc"] = input_batch.query_start_loc_np[: input_batch.num_reqs + 1].tolist() embeds = self.model.embed_input_ids(input_batch.input_ids[: input_batch.num_tokens], **kwargs) inputs_embeds = self.encoder_runner.inputs_embeds inputs_embeds[: embeds.shape[0]] = embeds From eb1699bfca2ca6f1bb074afdad7236b6809c3514 Mon Sep 17 00:00:00 2001 From: Zheng Wengang Date: Fri, 2 Oct 2026 15:18:51 +0100 Subject: [PATCH 15/22] [FEAT][Fronted] Support incremental prefill for realtime video streaming (#5894) Signed-off-by: ZhengWG Signed-off-by: Zheng Wengang Co-authored-by: amy-why-3459 --- docs/serving/video_stream_api.md | 47 +- .../openai_api/test_serving_video_stream.py | 554 ++++++++++++- .../openai_api/test_video_audio_delta.py | 4 +- vllm_omni/entrypoints/async_omni.py | 4 + .../openai/serving_video_stream.py | 36 +- .../entrypoints/openai/video_stream_base.py | 776 ++++++++++++++---- 6 files changed, 1187 insertions(+), 234 deletions(-) diff --git a/docs/serving/video_stream_api.md b/docs/serving/video_stream_api.md index 22362aac6..8d2194842 100644 --- a/docs/serving/video_stream_api.md +++ b/docs/serving/video_stream_api.md @@ -16,6 +16,10 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \ --trust-remote-code ``` +To enable [incremental prefill](#incremental-prefill), set +`enable_prefix_caching: true` on stage 0 (the thinker) in the deploy config. Otherwise +the server falls back to rebuilding the prompt per query. + ### Run the Example Client ```bash @@ -54,7 +58,7 @@ WebSocket /v1/video/chat/stream When buffered frames include client-supplied `frame_id` values, the server emits `video.frames.consumed` after the first engine output. Its `frame_ids`, `frames`, and `latest_pts_ms` describe the images included in that query's prompt. -Queries stride-sample the buffered frames, keeping the last frame, then exclude frames already known to have failed decoding. Excluded frames are not replaced with other buffered frames. Frames whose background decoding has not finished remain eligible through the image URL path. +Queries stride-sample the buffered frames, keeping the last frame, then exclude frames already known to have failed decoding. With incremental prefill active, every buffered frame is included instead of stride-sampling, with the same decode-failure exclusion. Excluded frames are not replaced with other buffered frames. The incremental path waits for the whole window to finish decoding. The legacy path waits only for its sampled positions, so an unselected frame still decoding does not delay the query. **Bugfix compatibility note:** `video.frames.consumed` now excludes known decode failures that were already excluded from the prompt. Older versions could report those frames and their timestamps as consumed. The event name and fields are unchanged; an empty selection reports empty lists and `latest_pts_ms: null`. @@ -64,10 +68,10 @@ Queries stride-sample the buffered frames, keeping the last frame, then exclude | ------- | ------ | --------- | ------------- | | `model` | string or null | null | Optional model name. Usually omitted because the server hosts one model. | | `modalities` | list[string] | `["text", "audio"]` | Output modalities. Use `["text"]`, `["audio"]`, or both. | -| `num_frames` | integer, 1-128 | `4` | Number of buffered frames sampled for each query. | -| `max_frames` | integer, 1-256 | `50` | Maximum retained frame buffer size. Oldest frames are evicted first. | +| `num_frames` | integer, 1-128 | `4` | Number of buffered frames sampled for each query. Legacy path only: with incremental prefill active, every buffered frame is submitted and frame density is the client's responsibility (clients push discrete frames at their own rate). | +| `max_frames` | integer, 1-256 | `50` | Maximum retained frame buffer size. Below the cap both paths keep every frame. Legacy evicts the oldest frame once full. Incremental prefill drops the oldest half at once, then appends until full again. | | `system_prompt` | string or null | null | Optional custom system prompt. | -| `use_audio_in_video` | bool | `true` | Include streamed audio chunks in multimodal video understanding when audio is present. | +| `use_audio_in_video` | bool | `true` | Incremental prefill: forwarded on warmup and every query when this is true, so multimodal hashes stay aligned. Legacy: forwarded only when this query has input audio. | | `sampling_params_list` | list or null | null | Optional per-stage parameter dictionaries. Each provided entry replaces that stage's deployment sampling settings. | | `enable_frame_filter` | bool | `true` | Enable EVS near-duplicate frame filtering. | | `frame_filter_threshold` | float, 0.0-1.0 | `0.95` | EVS similarity threshold. Higher keeps more frames; lower drops more near-duplicates. | @@ -107,9 +111,42 @@ The server accepts these legacy field names and rewrites them before validation. EVS compares downsampled frames and drops near-duplicate frames before they enter the session frame buffer. `frame_filter_threshold` controls retention: higher values are more permissive and keep more frames; lower values are more aggressive and drop more similar frames. +## Incremental Prefill + +Server-side only (stage-0 `enable_prefix_caching: true` and `modalities: ["text"]`). +There is no `session.config` flag. Audio-output sessions stay on the legacy path. + +Input is unchanged: `video.frame` and `audio.chunk` are separate buffers. Warmup +prefills `history + frames` as they arrive (`max_tokens=1`). When +`use_audio_in_video` is true, warmup and query both pass that processor kwarg. +`video.query` is the same prompt plus optional input audio and the question text, +so the warmed vision prefix cache is reused and the query pays for the audio/text +suffix. All buffered frames are +submitted in arrival order (`num_frames` subsample is legacy-only). Once the buffer +hits `max_frames`, the oldest half is dropped in one step and a warning is logged; +until then the buffer matches the legacy path. After a turn, +history compresses to the last two text messages and the next warmup follows the +new prefix. + +### Legacy path (incremental prefill inactive) + +Stage-0 prefix caching is off, or `modalities` includes `"audio"`: no warmup, +frames are re-sampled up to `num_frames` each query, and the query waits only +for those sampled frames to decode. Prompt shape is the same +(`image_pil` frames, optional trailing `input_audio`, text). Both paths keep the +frame buffer across queries (`max_frames`) and compress history the same way. + ## Known Limitations -- Session KV reuse and incremental prefill are not implemented in this PR. Each `video.query` rebuilds the model prompt from the retained frame and audio buffers. +- Past `max_frames`, incremental prefill keeps a suffix of the newest frames, not a resample of the whole stream. Rebuilding that longer window is left to a follow-up. +- Incremental prefill applies only to text-only sessions with stage-0 prefix caching enabled + on `WS /v1/video/chat/stream`. It does not accelerate `/v1/realtime` camera input, which + uses duplex conversation items rather than `video.frame` / `video.query`. Audio-output + sessions on this endpoint rebuild the prompt each query. +- Stage-0 prefix caching is safe for audio-output requests: on a hit, the omni prefix cache + restores the thinker hidden states the talker needs. A preempted request that resumes on a + hit does not get them back, so size the deployment to avoid preemption; see + [Prefix Caching](../design/feature/prefix_caching.md). - Back-to-back short replies can still expose an engine-layer scheduler race. The PR notes an observed workaround of at least 200 ms idle between turns when clients repeatedly see idle timeouts. - If the audio buffer exceeds the server limit, the server emits `Audio buffer overflow` and clears the currently buffered audio for the session. - The API is intended for Qwen3-Omni streaming video understanding; other models may not support the same multimodal processor arguments. diff --git a/tests/entrypoints/openai_api/test_serving_video_stream.py b/tests/entrypoints/openai_api/test_serving_video_stream.py index c9c4f6a4b..65450267f 100644 --- a/tests/entrypoints/openai_api/test_serving_video_stream.py +++ b/tests/entrypoints/openai_api/test_serving_video_stream.py @@ -7,10 +7,12 @@ from __future__ import annotations import asyncio import base64 import gc +import hashlib import io import json import threading import weakref +from types import SimpleNamespace from typing import Any from unittest.mock import AsyncMock, MagicMock @@ -26,7 +28,9 @@ from vllm_omni.entrypoints.openai.serving_video_stream import ( QwenOmniStreamingVideoHandler, StreamingVideoSessionConfig, ) -from vllm_omni.entrypoints.openai.video_stream_base import OmniStreamingVideoHandler +from vllm_omni.entrypoints.openai.video_stream_base import ( + OmniStreamingVideoHandler, +) from vllm_omni.outputs import OmniRequestOutput pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -506,17 +510,9 @@ async def test_video_frames_consumed_excludes_bad_frame_in_query_snapshot( pytest.param(2, ["f0", "f1", "f2"], True, [0, 2], id="repeated-image"), ], ) -async def test_video_frame_selection_preserves_uncached_frames( - monkeypatch, num_frames, frame_ids, repeat_image, expected_indices -): - """Prewarm misses still use image URLs, with metadata paired by buffer position.""" +async def test_video_frame_selection_pairs_metadata_by_position(num_frames, frame_ids, repeat_image, expected_indices): + """Queries wait for prewarmed PIL frames, with metadata paired by buffer position.""" engine_prompts = [] - pending_decode = asyncio.Event() - - async def blocked_to_thread(*args, **kwargs): - await pending_decode.wait() - - monkeypatch.setattr(asyncio, "to_thread", blocked_to_thread) class OneOutputEngine: async def generate(self, **kwargs): @@ -550,13 +546,13 @@ async def test_video_frame_selection_preserves_uncached_frames( assert not [message for message in ws.sent if message["type"] == "error"] assert len(engine_prompts) == 1 content = engine_prompts[0]["messages"][-1]["content"] - assert content == [ - *[ - {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{frames[index]}"}} - for index in expected_indices - ], - {"type": "text", "text": "describe"}, - ] + assert content[-1] == {"type": "text", "text": "describe"} + image_parts = content[:-1] + assert [part["type"] for part in image_parts] == ["image_pil"] * len(expected_indices) + for part, index in zip(image_parts, expected_indices): + raw = base64.b64decode(frames[index]) + assert part["image_pil"].tobytes() == Image.open(io.BytesIO(raw)).convert("RGB").tobytes() + assert part["uuid"] == hashlib.md5(raw, usedforsecurity=False).hexdigest() consumed_events = [message for message in ws.sent if message["type"] == "video.frames.consumed"] if not any(frame_ids): assert consumed_events == [] @@ -909,6 +905,9 @@ async def test_interrupted_queries_wait_for_abort_without_fixed_delay(monkeypatc class BlockingEngine: async def generate(self, *, request_id, **kwargs): + submitted = kwargs.get("submitted") + if submitted is not None and not submitted.done(): + submitted.set_result(None) if request_ids: assert aborted[-1] == request_ids[-1] request_ids.append(request_id) @@ -1004,7 +1003,7 @@ async def test_video_done_waits_for_in_flight_query(): @pytest.mark.asyncio -async def test_frame_prewarm_does_not_block_following_query(monkeypatch): +async def test_query_waits_for_frame_prewarm(monkeypatch): decode_started = threading.Event() release_decode = threading.Event() query_started = asyncio.Event() @@ -1035,14 +1034,80 @@ async def test_frame_prewarm_does_not_block_following_query(monkeypatch): assert decode_started.is_set() ws.put({"type": "video.query", "text": "describe"}) - await asyncio.wait_for(query_started.wait(), timeout=2.0) + await asyncio.sleep(0.05) + assert not query_started.is_set() release_decode.set() + await asyncio.wait_for(query_started.wait(), timeout=2.0) ws.put({"type": "video.done"}) await asyncio.wait_for(task, timeout=2.0) assert "session.done" in ws.sent_types() +@pytest.mark.asyncio +async def test_legacy_query_does_not_wait_for_unselected_frame(monkeypatch): + """APC off and num_frames=1: a blocked older frame must not delay the query.""" + frame_a = _make_jpeg(1, 0, 0) + frame_b = _make_jpeg(0, 1, 0) + decode_a_started = threading.Event() + release_a = threading.Event() + decode_b_done = threading.Event() + query_started = asyncio.Event() + + def gated_decode(raw_bytes: bytes): + if raw_bytes == frame_a: + decode_a_started.set() + release_a.wait(timeout=5.0) + image = Image.open(io.BytesIO(raw_bytes)).convert("RGB") + if raw_bytes == frame_b: + decode_b_done.set() + return image + + class RecordingHandler(QwenOmniStreamingVideoHandler): + async def _process_query(self, *args, **kwargs): + query_started.set() + + monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode) + ws = TimedWebSocket() + handler = RecordingHandler(chat_service=object(), engine_client=object(), idle_timeout=5.0) + task = asyncio.create_task(handler.handle_session(ws)) + try: + ws.put( + { + "type": "session.config", + "model": "test", + "modalities": ["text"], + "num_frames": 1, + "enable_frame_filter": False, + } + ) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": _b64(frame_a), "frame_id": "A"}) + for _ in range(100): + if decode_a_started.is_set(): + break + await asyncio.sleep(0.01) + assert decode_a_started.is_set() + ws.put({"type": "video.frame", "data": _b64(frame_b), "frame_id": "B"}) + for _ in range(200): + if decode_b_done.is_set(): + break + await asyncio.sleep(0.01) + assert decode_b_done.is_set() + + ws.put({"type": "video.query", "text": "describe"}) + await asyncio.wait_for(query_started.wait(), timeout=2.0) + assert not release_a.is_set() + release_a.set() + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + finally: + release_a.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + @pytest.mark.asyncio @pytest.mark.parametrize("max_frames", [1, 2], ids=["evicted", "retained"]) async def test_frame_prewarm_only_keeps_retained_images(monkeypatch, max_frames): @@ -1272,3 +1337,452 @@ def test_build_messages_keeps_recent_history_text_only(): assert messages[1] == {"role": "assistant", "content": "recent answer"} assert messages[2] == user_message assert user_message["content"][-1] == {"type": "text", "text": "current question"} + + +# --------------------------------------------------------------------------- +# Incremental prefill (warmup) lifecycle +# --------------------------------------------------------------------------- + + +class RecordingReuseEngine: + """Engine stub with stage-0 prefix caching on; records generate() and abort().""" + + def __init__(self, warmup_delay: float = 0.0, query_delay: float = 0.0): + self.requests: list[dict[str, Any]] = [] + self.aborted: list[str] = [] + self._warmup_delay = warmup_delay + self._query_delay = query_delay + self.engine = SimpleNamespace( + stage_vllm_configs=[SimpleNamespace(cache_config=SimpleNamespace(enable_prefix_caching=True))] + ) + + def warmup_ids(self) -> list[str]: + return [r["request_id"] for r in self.requests if r["request_id"].startswith("video-warmup-")] + + def query_ids(self) -> list[str]: + return [r["request_id"] for r in self.requests if not r["request_id"].startswith("video-warmup-")] + + def generate(self, **kwargs): + self.requests.append(kwargs) + submitted = kwargs.get("submitted") + if submitted is not None and not submitted.done(): + submitted.set_result(None) + is_warmup = kwargs["request_id"].startswith("video-warmup-") + delay = self._warmup_delay if is_warmup else self._query_delay + + async def _gen(): + if delay: + await asyncio.sleep(delay) + yield _text_result("ok") + + return _gen() + + async def abort(self, request_id: str) -> None: + self.aborted.append(request_id) + + +class PromptRecordingHandler(QwenOmniStreamingVideoHandler): + """Records each query turn's user-message content and processor kwargs.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.query_contents: list[list[dict[str, Any]]] = [] + self.warmup_contents: list[list[dict[str, Any]]] = [] + self.mm_processor_kwargs: list[Any] = [] + + async def _preprocess_to_engine_prompt(self, request): + self.mm_processor_kwargs.append(getattr(request, "mm_processor_kwargs", None)) + return {"prompt": "p"} + + def build_engine_prompt( + self, config, frame_buffer, audio_buffer, message_history, query_text, prewarmed_frames, *, frame_indices=None + ): + messages, user_message = super().build_engine_prompt( + config, + frame_buffer, + audio_buffer, + message_history, + query_text, + prewarmed_frames, + frame_indices=frame_indices, + ) + self.query_contents.append(list(user_message["content"])) + return messages, user_message + + def build_engine_prompt_prefix(self, config, frame_buffer, message_history, prewarmed_frames): + messages = super().build_engine_prompt_prefix(config, frame_buffer, message_history, prewarmed_frames) + if messages: + self.warmup_contents.append(list(messages[-1]["content"])) + return messages + + +async def _poll(predicate, timeout: float = 3.0) -> bool: + for _ in range(int(timeout / 0.01)): + if predicate(): + return True + await asyncio.sleep(0.01) + return bool(predicate()) + + +@pytest.mark.asyncio +async def test_warmup_fires_on_frame_and_query_cancels_it(): + engine = RecordingReuseEngine(warmup_delay=10.0, query_delay=0.3) + handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0) + + ws = TimedWebSocket() + task = asyncio.create_task(handler.handle_session(ws)) + ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False}) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": _b64(_make_jpeg())}) + + assert await _poll(lambda: engine.warmup_ids()) + warmup_id = engine.warmup_ids()[0] + assert engine.requests[0]["sampling_params"].max_tokens == 1 + + ws.put({"type": "video.query", "text": "describe"}) + assert await _poll(lambda: engine.query_ids()) + # in-flight warmup was cancelled + aborted, and not restarted mid-query + assert warmup_id in engine.aborted + assert len(engine.warmup_ids()) == 1 + assert handler.mm_processor_kwargs[:2] == [ + {"use_audio_in_video": True}, + {"use_audio_in_video": True}, + ] + + assert await _poll(lambda: "response.text.done" in ws.sent_types()) + # the committed turn changed the context -> a fresh warmup fires + assert await _poll(lambda: len(engine.warmup_ids()) == 2) + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + + +@pytest.mark.asyncio +async def test_cancelled_warmup_reaches_stage0_before_abort(): + """Cancel during submit must not abort a warmup the engine never received.""" + + class _HoldSubmitEngine(RecordingReuseEngine): + def __init__(self): + super().__init__() + self.started = asyncio.Event() + self.release_submit = asyncio.Event() + + def generate(self, **kwargs): + submitted = kwargs.get("submitted") + + async def _gen(): + self.started.set() + await self.release_submit.wait() + self.requests.append(kwargs) + if submitted is not None and not submitted.done(): + submitted.set_result(None) + yield _text_result("ok") + + return _gen() + + engine = _HoldSubmitEngine() + handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0) + ws = TimedWebSocket() + task = asyncio.create_task(handler.handle_session(ws)) + ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False}) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": _b64(_make_jpeg())}) + assert await _poll(engine.started.is_set) + ws.put({"type": "video.query", "text": "describe"}) + await asyncio.sleep(0.05) + assert engine.requests == [] + assert engine.aborted == [] + engine.release_submit.set() + assert await _poll(lambda: engine.aborted) + assert engine.warmup_ids() + assert engine.aborted[0] == engine.warmup_ids()[0] + assert await _poll(lambda: engine.query_ids()) + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + + +@pytest.mark.asyncio +async def test_cancelled_query_reaches_stage0_before_abort(): + """Cancel during query submit must not abort a request the engine never received.""" + + class _HoldQueryEngine(RecordingReuseEngine): + def __init__(self): + super().__init__() + self.query_started = asyncio.Event() + self.release_submit = asyncio.Event() + + def generate(self, **kwargs): + if str(kwargs.get("request_id", "")).startswith("video-warmup-"): + return super().generate(**kwargs) + submitted = kwargs.get("submitted") + + async def _gen(): + self.query_started.set() + await self.release_submit.wait() + self.requests.append(kwargs) + if submitted is not None and not submitted.done(): + submitted.set_result(None) + yield _text_result("ok") + + return _gen() + + engine = _HoldQueryEngine() + handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0) + ws = TimedWebSocket() + task = asyncio.create_task(handler.handle_session(ws)) + ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False}) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": _b64(_make_jpeg())}) + assert await _poll(lambda: engine.warmup_ids()) + ws.put({"type": "video.query", "text": "first"}) + assert await _poll(engine.query_started.is_set) + ws.put({"type": "video.query", "text": "second"}) + await asyncio.sleep(0.05) + assert engine.query_ids() == [] + assert engine.aborted == [] + engine.release_submit.set() + assert await _poll(lambda: engine.query_ids()) + first_query = engine.query_ids()[0] + assert await _poll(lambda: first_query in engine.aborted) + assert await _poll(lambda: len(engine.query_ids()) == 2) + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + + +@pytest.mark.asyncio +async def test_audio_session_never_starts_warmup(): + engine = RecordingReuseEngine() + handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0) + + ws = TimedWebSocket() + task = asyncio.create_task(handler.handle_session(ws)) + ws.put({"type": "session.config", "model": "test", "modalities": ["text", "audio"], "enable_frame_filter": False}) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": _b64(_make_jpeg())}) + ws.put({"type": "video.frame", "data": _b64(_make_jpeg(9, 9, 9))}) + await asyncio.sleep(0.3) + + assert engine.warmup_ids() == [] + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + + +@pytest.mark.asyncio +async def test_warmup_starts_only_after_frames_ready(monkeypatch): + decode_started = threading.Event() + release_decode = threading.Event() + + def gated_decode(raw_bytes: bytes): + decode_started.set() + release_decode.wait(timeout=5.0) + return Image.open(io.BytesIO(raw_bytes)).convert("RGB") + + monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode) + engine = RecordingReuseEngine() + handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0) + + ws = TimedWebSocket() + task = asyncio.create_task(handler.handle_session(ws)) + ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False}) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": _b64(_make_jpeg())}) + assert await _poll(decode_started.is_set) + await asyncio.sleep(0.1) + assert engine.warmup_ids() == [] + + release_decode.set() + assert await _poll(lambda: engine.warmup_ids()) + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + + +@pytest.mark.asyncio +async def test_query_cancelling_warmup_keeps_prewarm_alive(monkeypatch): + """Cancel warmup during prefill must not drop a later frame still decoding.""" + frame1 = _b64(_make_jpeg(1, 1, 1)) + frame2 = _b64(_make_jpeg(2, 2, 2)) + raw2 = base64.b64decode(frame2) + decode_started = threading.Event() + release_decode = threading.Event() + + def gated_decode(raw_bytes: bytes): + if raw_bytes == raw2: + decode_started.set() + release_decode.wait(timeout=5.0) + return Image.open(io.BytesIO(raw_bytes)).convert("RGB") + + monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode) + engine = RecordingReuseEngine(warmup_delay=10.0) + handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0) + + ws = TimedWebSocket() + task = asyncio.create_task(handler.handle_session(ws)) + ws.put({"type": "session.config", "model": "test", "modalities": ["text"], "enable_frame_filter": False}) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": frame1}) + assert await _poll(lambda: engine.warmup_ids()) + + ws.put({"type": "video.frame", "data": frame2}) + assert await _poll(decode_started.is_set) + + ws.put({"type": "video.query", "text": "describe"}) + await asyncio.sleep(0.1) + release_decode.set() + + assert await _poll(lambda: "response.text.done" in ws.sent_types()) + assert len(handler.query_contents) == 1 + image_parts = [p for p in handler.query_contents[0] if p.get("type") == "image_pil"] + assert len(image_parts) == 2 + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + + +@pytest.mark.asyncio +async def test_query_keeps_frames_evicted_while_pinned(monkeypatch): + """Regression: frames FIFO-evicted while an in-flight query awaits prewarm + must keep their PIL entries and stay in the query prompt.""" + frame_a = _b64(_make_jpeg(1, 1, 1)) + frame_a2 = _b64(_make_jpeg(2, 2, 2)) + frame_b = _b64(_make_jpeg(3, 3, 3)) + raw_a = base64.b64decode(frame_a) + decoded: list[bytes] = [] + release_others = threading.Event() + + def gated_decode(raw_bytes: bytes): + if raw_bytes != raw_a: + release_others.wait(timeout=5.0) + decoded.append(raw_bytes) + return Image.open(io.BytesIO(raw_bytes)).convert("RGB") + + monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", gated_decode) + engine = RecordingReuseEngine() + handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0) + + ws = TimedWebSocket() + task = asyncio.create_task(handler.handle_session(ws)) + ws.put( + { + "type": "session.config", + "model": "test", + "modalities": ["text"], + "max_frames": 2, + "enable_frame_filter": False, + } + ) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": frame_a}) + assert await _poll(lambda: raw_a in decoded) + ws.put({"type": "video.frame", "data": frame_a2}) # decode gated -> query will wait + + ws.put({"type": "video.query", "text": "describe"}) # pins [a, a2] + await asyncio.sleep(0.1) + ws.put({"type": "video.frame", "data": frame_b}) # buffer full -> evicts pinned a + await asyncio.sleep(0.1) + release_others.set() + + assert await _poll(lambda: "response.text.done" in ws.sent_types()) + assert len(handler.query_contents) == 1 + uuids = [p["uuid"] for p in handler.query_contents[0] if p.get("type") == "image_pil"] + expected = [hashlib.md5(base64.b64decode(f), usedforsecurity=False).hexdigest() for f in (frame_a, frame_a2)] + assert uuids == expected + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + + +async def _query_image_uuids(modalities: list[str], frames: list[str], *, max_frames: int) -> list[str]: + engine = RecordingReuseEngine() + handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0) + ws = TimedWebSocket() + task = asyncio.create_task(handler.handle_session(ws)) + ws.put( + { + "type": "session.config", + "model": "test", + "modalities": modalities, + "max_frames": max_frames, + "num_frames": max_frames, + "enable_frame_filter": False, + } + ) + await asyncio.sleep(0) + for frame in frames: + ws.put({"type": "video.frame", "data": frame}) + ws.put({"type": "video.query", "text": "describe"}) + assert await _poll(lambda: handler.query_contents) + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + return [part["uuid"] for part in handler.query_contents[0] if part.get("type") == "image_pil"] + + +def _frame_uuid(frame_b64: str) -> str: + return hashlib.md5(base64.b64decode(frame_b64), usedforsecurity=False).hexdigest() + + +@pytest.mark.asyncio +async def test_below_max_frames_incremental_matches_legacy_window(): + frames = [_b64(_make_jpeg(i, 0, 0)) for i in range(3)] + expected = [_frame_uuid(frame) for frame in frames] + assert await _query_image_uuids(["text"], frames, max_frames=4) == expected + assert await _query_image_uuids(["text", "audio"], frames, max_frames=4) == expected + + +@pytest.mark.asyncio +async def test_incremental_full_window_drops_oldest_half(): + frames = [_b64(_make_jpeg(i, 1, 0)) for i in range(5)] + kept = await _query_image_uuids(["text"], frames, max_frames=4) + legacy = await _query_image_uuids(["text", "audio"], frames, max_frames=4) + assert kept == [_frame_uuid(frame) for frame in frames[2:]] + assert legacy == [_frame_uuid(frame) for frame in frames[1:]] + + +async def _run_frame_audio_query(modalities: list[str], frame: str, pcm: bytes) -> PromptRecordingHandler: + engine = RecordingReuseEngine() + handler = PromptRecordingHandler(chat_service=object(), engine_client=engine, idle_timeout=5.0) + ws = TimedWebSocket() + task = asyncio.create_task(handler.handle_session(ws)) + ws.put( + { + "type": "session.config", + "model": "test", + "modalities": modalities, + "enable_frame_filter": False, + } + ) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": frame}) + if "audio" not in modalities: + assert await _poll(lambda: engine.warmup_ids()) + else: + await asyncio.sleep(0.05) + assert engine.warmup_ids() == [] + ws.put({"type": "audio.chunk", "data": _b64(pcm)}) + ws.put({"type": "video.query", "text": "describe"}) + assert await _poll(lambda: engine.query_ids()) + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + return handler + + +@pytest.mark.asyncio +async def test_reuse_query_with_audio_matches_legacy_prompt(): + """Incremental and legacy query prompts match; both set the processor kwarg.""" + frame = _b64(_make_jpeg()) + pcm = b"\x01\x00\x02\x00" + reuse = await _run_frame_audio_query(["text"], frame, pcm) + legacy = await _run_frame_audio_query(["text", "audio"], frame, pcm) + + assert reuse.warmup_contents + assert all(part.get("type") != "input_audio" for part in reuse.warmup_contents[0]) + assert [part.get("type") for part in reuse.warmup_contents[0]] == ["image_pil"] + + assert reuse.query_contents and legacy.query_contents + reuse_q = reuse.query_contents[0] + legacy_q = legacy.query_contents[0] + assert [part.get("type") for part in reuse_q] == ["image_pil", "input_audio", "text"] + assert [part.get("type") for part in legacy_q] == ["image_pil", "input_audio", "text"] + assert reuse_q[0]["uuid"] == legacy_q[0]["uuid"] + assert reuse_q[1] == legacy_q[1] + assert reuse_q[2] == legacy_q[2] == {"type": "text", "text": "describe"} + + assert reuse.mm_processor_kwargs[0] == {"use_audio_in_video": True} + assert reuse.mm_processor_kwargs[1] == {"use_audio_in_video": True} + assert legacy.mm_processor_kwargs[0] == {"use_audio_in_video": True} diff --git a/tests/entrypoints/openai_api/test_video_audio_delta.py b/tests/entrypoints/openai_api/test_video_audio_delta.py index 123ee7424..52cd6d892 100644 --- a/tests/entrypoints/openai_api/test_video_audio_delta.py +++ b/tests/entrypoints/openai_api/test_video_audio_delta.py @@ -174,7 +174,9 @@ async def test_real_producer_deltas_reach_video_client( produced.append(OmniRequestOutput.from_stage_output(result, final_output_type="audio")) class Engine: - async def generate(self, *, prompt, request_id, output_modalities, sampling_params_list=None): + async def generate(self, *, prompt, request_id, output_modalities, sampling_params_list=None, submitted=None): + if submitted is not None and not submitted.done(): + submitted.set_result(None) if explicit_sampling: assert sampling_params_list is not None assert sampling_params_list[0].output_kind == RequestOutputKind.DELTA diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 34fd76f22..2a3253599 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -128,6 +128,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient): reasoning_ended: bool | None = None, reasoning_parser_kwargs: dict[str, Any] | None = None, arrival_time: float | None = None, + submitted: asyncio.Future[None] | None = None, ) -> AsyncGenerator[OmniRequestOutput, None]: """Generate outputs for the given prompt(s) asynchronously. @@ -296,6 +297,9 @@ class AsyncOmni(AsyncOmniBase, EngineClient): if admitting: await self._release_generate_admission() admitting = False + # Callers wait on this before abort so Stage0 observes the sender cache. + if submitted is not None and not submitted.done(): + submitted.set_result(None) # Refresh gauges on arrival. self._publish_request_gauges(len(self.request_states)) diff --git a/vllm_omni/entrypoints/openai/serving_video_stream.py b/vllm_omni/entrypoints/openai/serving_video_stream.py index 5f7334bd2..7cabbc496 100644 --- a/vllm_omni/entrypoints/openai/serving_video_stream.py +++ b/vllm_omni/entrypoints/openai/serving_video_stream.py @@ -66,29 +66,10 @@ class QwenOmniStreamingVideoHandler(OmniStreamingVideoHandlerBase): ) -> tuple[list[dict[str, Any]], dict[str, Any]]: prewarmed = prewarmed_frames or {} if frame_indices is None: - frame_indices = self._sample_frame_indices(frame_buffer, config.num_frames, prewarmed) - user_content: list[dict] = [] - for index in frame_indices: - frame_b64 = frame_buffer[index] - cached = prewarmed.get(frame_b64) - if cached is not None: - # The shared selector excludes failed decodes before this point. - assert isinstance(cached, tuple) - pil, pil_uuid = cached - user_content.append( - { - "type": "image_pil", - "image_pil": pil, - "uuid": pil_uuid, - } - ) - else: - user_content.append( - { - "type": "image_url", - "image_url": {"url": f"data:image/jpeg;base64,{frame_b64}"}, - } - ) + frame_indices = self._select_frame_indices(config, frame_buffer, prewarmed) + frames = [frame_buffer[index] for index in frame_indices] + + user_content: list[dict] = self._build_frame_image_parts(frames, prewarmed) if len(audio_buffer) > 0: wav_b64 = self._pcm_to_wav_b64(bytes(audio_buffer)) @@ -107,14 +88,7 @@ class QwenOmniStreamingVideoHandler(OmniStreamingVideoHandlerBase): user_message: dict[str, Any] = {"role": "user", "content": user_content} - messages: list[dict[str, Any]] = [] - if config.system_prompt: - messages.append({"role": "system", "content": config.system_prompt}) - - recent_history = message_history[-2:] if len(message_history) > 2 else message_history - for hist_msg in recent_history: - messages.append(self._text_only_message(hist_msg)) - + messages = self._history_prefix_messages(config, message_history) messages.append(user_message) return messages, user_message diff --git a/vllm_omni/entrypoints/openai/video_stream_base.py b/vllm_omni/entrypoints/openai/video_stream_base.py index e5893545d..d15da9451 100644 --- a/vllm_omni/entrypoints/openai/video_stream_base.py +++ b/vllm_omni/entrypoints/openai/video_stream_base.py @@ -38,7 +38,7 @@ import wave from collections.abc import Mapping from dataclasses import dataclass from enum import Enum -from typing import Any, Final, Protocol, TypeAlias, runtime_checkable +from typing import Any, Final, Protocol, TypeAlias, TypeGuard, runtime_checkable import torch from fastapi import WebSocket, WebSocketDisconnect @@ -77,6 +77,21 @@ def _decode_frame_bytes(raw_bytes: bytes) -> Any: return Image.open(io.BytesIO(raw_bytes)).convert("RGB") +def _usable_prewarmed_frame(cached: PrewarmedFrame | None) -> TypeGuard[tuple[Any, str]]: + return cached is not None and cached is not _BAD_FRAME + + +def _stage0_prefix_caching_enabled(engine_client: Any) -> bool: + """Read stage-0 (thinker) prefix caching from the engine. Called once + at handler init: incremental prefill is a server-side capability, not a + per-request choice.""" + engine = getattr(engine_client, "engine", None) + stage_vllm_configs = getattr(engine, "stage_vllm_configs", None) + stage0 = stage_vllm_configs[0] if stage_vllm_configs else None + cache_config = getattr(stage0, "cache_config", None) + return bool(getattr(cache_config, "enable_prefix_caching", False)) + + @runtime_checkable class VideoStreamPipelineHooks(Protocol): """Pipeline-specific hooks for streaming video handlers.""" @@ -108,6 +123,20 @@ class VideoStreamPipelineHooks(Protocol): """Update session state after a successful turn.""" ... + def build_engine_prompt_prefix( + self, + config: "StreamingVideoSessionConfig", + frame_buffer: list[str], + message_history: list[dict[str, Any]], + prewarmed_frames: Mapping[str, PrewarmedFrame], + ) -> list[dict[str, Any]] | None: + """Prefix of the next :meth:`build_engine_prompt` through the last vision token. + + History plus arriving video frames, without query text or input audio. + Return None when there are no usable frames (or to disable incremental prefill). + """ + ... + @dataclass(frozen=True) class VideoStreamTurnTrigger: @@ -144,7 +173,10 @@ class StreamingVideoSessionConfig(BaseModel): ) use_audio_in_video: bool = Field( default=True, - description="Interleave audio chunks with video frames when audio input is present.", + description=( + "Pass use_audio_in_video to the processor when this is true and " + "(incremental prefill is active, or this query has input audio)." + ), ) sampling_params_list: list[dict[str, Any]] | None = Field( default=None, @@ -194,10 +226,77 @@ class OmniStreamingVideoHandler: ) -> None: raise NotImplementedError + @staticmethod + def _build_frame_image_parts( + frames: list[str], + prewarmed_frames: Mapping[str, PrewarmedFrame] | None, + ) -> list[dict[str, Any]]: + """Decoded frame parts. Prefix and query both wait for PIL prewarm so + every frame is ``image_pil`` plus the arrival-time uuid.""" + prewarmed = prewarmed_frames or {} + parts: list[dict[str, Any]] = [] + for frame_b64 in frames: + cached = prewarmed.get(frame_b64) + if not _usable_prewarmed_frame(cached): + continue + pil, mm_uuid = cached + parts.append( + { + "type": "image_pil", + "image_pil": pil, + "uuid": mm_uuid, + } + ) + return parts + + def _history_prefix_messages( + self, + config: StreamingVideoSessionConfig, + message_history: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + """Shared prompt prefix (system + compressed text-only history). + + :meth:`build_engine_prompt_prefix` and :meth:`build_engine_prompt` + MUST build this identically: any divergence before the frame parts + voids the warmed KV from that token onward. + """ + messages: list[dict[str, Any]] = [] + if config.system_prompt: + messages.append({"role": "system", "content": config.system_prompt}) + recent_history = message_history[-2:] if len(message_history) > 2 else message_history + for hist_msg in recent_history: + messages.append(self._text_only_message(hist_msg)) + return messages + + def build_engine_prompt_prefix( + self, + config: StreamingVideoSessionConfig, + frame_buffer: list[str], + message_history: list[dict[str, Any]], + prewarmed_frames: Mapping[str, PrewarmedFrame], + ) -> list[dict[str, Any]] | None: + """Prefix of the next :meth:`build_engine_prompt` through the last vision token. + + History plus arriving video frames, without query text or input audio. + A mismatch voids the warmed KV from the first diverging token. + Return None when there are no usable frames (or override to disable). + """ + frame_parts = self._build_frame_image_parts(frame_buffer, prewarmed_frames) + if not frame_parts: + return None + + messages = self._history_prefix_messages(config, message_history) + messages.append({"role": "user", "content": frame_parts}) + return messages + def create_message_history(self, config: StreamingVideoSessionConfig) -> Any: """Per-session conversation state (default: empty OpenAI-style list).""" return [] + def _incremental_prefill_active(self, config: StreamingVideoSessionConfig) -> bool: + """Incremental frame prefill: stage-0 prefix caching on, text-only output.""" + return self._incremental_prefill_supported and "audio" not in config.modalities + def on_frame_buffered( self, raw_bytes: bytes, @@ -219,6 +318,7 @@ class OmniStreamingVideoHandler: self._idle_timeout = idle_timeout self._config_timeout = config_timeout self._engine_client = engine_client + self._incremental_prefill_supported = _stage0_prefix_caching_enabled(engine_client) async def handle_session(self, websocket: WebSocket) -> None: """Main session loop for a single WebSocket connection.""" @@ -244,6 +344,53 @@ class OmniStreamingVideoHandler: interrupt_event = asyncio.Event() prewarm_tasks: set[asyncio.Task[Any]] = set() query_task: asyncio.Task[Any] | None = None + session_alive = True + # Session-owned decode readiness. Warmup/query wait on these + # Events, not on the prewarm tasks, so cancelling an engine + # request cannot cascade into JPEG decode. + frame_ready: dict[str, asyncio.Event] = {} + # In-flight warmup/query snapshots keep their PIL entries if FIFO + # evicts the live buffer while they await readiness. + pinned_frame_refs: dict[str, int] = {} + + def _drop_frame_cache(frame_b64: str) -> None: + frame_pil_cache.pop(frame_b64, None) + frame_ready.pop(frame_b64, None) + + def _pin_frames(frames: list[str]) -> None: + for frame_b64 in frames: + pinned_frame_refs[frame_b64] = pinned_frame_refs.get(frame_b64, 0) + 1 + + def _unpin_frames(frames: list[str]) -> None: + for frame_b64 in frames: + remaining = pinned_frame_refs.get(frame_b64, 0) - 1 + if remaining <= 0: + pinned_frame_refs.pop(frame_b64, None) + # Settle the eviction this pin deferred: the frame left + # the buffer, so its cache entry is orphaned. + if frame_b64 not in frame_buffer: + _drop_frame_cache(frame_b64) + else: + pinned_frame_refs[frame_b64] = remaining + + def _snapshot_prewarmed(frames: list[str]) -> dict[str, Any]: + return {frame_b64: frame_pil_cache[frame_b64] for frame_b64 in frames if frame_b64 in frame_pil_cache} + + def _frames_ready(frames: list[str]) -> bool: + for frame_b64 in frames: + event = frame_ready.get(frame_b64) + if event is None or not event.is_set(): + return False + return True + + async def _await_frames_ready(frames: list[str]) -> None: + waits: list[Any] = [] + for frame_b64 in dict.fromkeys(frames): + event = frame_ready.get(frame_b64) + if event is not None and not event.is_set(): + waits.append(event.wait()) + if waits: + await asyncio.gather(*waits) msg_queue: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue(maxsize=_MAX_MSG_QUEUE) @@ -294,15 +441,95 @@ class OmniStreamingVideoHandler: interrupt_event.set() prev_was_interrupted = True logger.info("Interrupt signaled for %s", active_request_id) + if query_task is not None and not query_task.done(): + query_task.cancel() + await asyncio.gather(query_task, return_exceptions=True) + query_task = None + # Abort after the task reaches Stage0 so a cancel during preprocess still submits. if abort_now and self._engine_client: try: await self._engine_client.abort(active_request_id) except Exception: logger.debug("Abort failed for %s", active_request_id, exc_info=True) - if query_task is not None and not query_task.done(): - query_task.cancel() - await asyncio.gather(query_task, return_exceptions=True) - query_task = None + + # --- Incremental prefill: keep the prefix cache warmed to the latest + # buffered frame so a query only pays for its own text suffix. --- + warmup_task: asyncio.Task[Any] | None = None + warmup_request_id: str | None = None + warmed_signature: tuple[Any, ...] | None = None + + def _context_signature() -> tuple[Any, ...]: + history_len = len(message_history) if isinstance(message_history, list) else None + return (history_len, tuple(frame_buffer)) + + def _maybe_start_warmup() -> None: + nonlocal warmup_task, warmup_request_id + if not session_alive: + return + if self._engine_client is None or not self._incremental_prefill_active(config): + return + if active_request_id is not None: + return + if warmup_task is not None and not warmup_task.done(): + return + if not frame_buffer or _context_signature() == warmed_signature: + return + frames = list(frame_buffer) + # Decode is session-owned: do not start prefill until PIL is + # already in cache, so this task is only engine work. + if not _frames_ready(frames): + return + signature = _context_signature() + request_id = f"video-warmup-{uuid.uuid4().hex[:12]}" + warmup_request_id = request_id + _pin_frames(frames) + + async def _run_warmup() -> None: + nonlocal warmup_request_id, warmed_signature + cancelled = False + try: + messages = self.build_engine_prompt_prefix( + config, frames, message_history, _snapshot_prewarmed(frames) + ) + if not messages: + warmed_signature = signature + return + await self._prefill_context(config, messages, request_id) + warmed_signature = signature + except asyncio.CancelledError: + cancelled = True + raise + except Exception: + # Mark attempted so a persistently failing state does + # not hot-loop; the next frame retriggers naturally. + warmed_signature = signature + logger.debug("Context warmup failed for %s", request_id, exc_info=True) + finally: + _unpin_frames(frames) + if warmup_request_id == request_id: + warmup_request_id = None + # Skip on cancel: query already owns the session and will + # call_soon after it finishes. Restarting here can spawn a + # second warmup with the same frames before active_request_id + # is set. + if not cancelled: + asyncio.get_running_loop().call_soon(_maybe_start_warmup) + + warmup_task = asyncio.create_task(_run_warmup()) + + async def _cancel_warmup() -> None: + nonlocal warmup_task, warmup_request_id + request_id = warmup_request_id + if warmup_task is not None and not warmup_task.done(): + warmup_task.cancel() + await asyncio.gather(warmup_task, return_exceptions=True) + warmup_task = None + warmup_request_id = None + if request_id and self._engine_client: + try: + await self._engine_client.abort(request_id) + except Exception: + logger.debug("Warmup abort failed for %s", request_id, exc_info=True) async def _start_query_turn(*, query_text: str) -> None: """Schedule a new inference turn from the current buffers.""" @@ -311,28 +538,43 @@ class OmniStreamingVideoHandler: await _cancel_active_query() if not frame_buffer: + await _cancel_warmup() await self._send_error(websocket, "No frames buffered") return + # Claim the session before cancelling warmup so a leftover + # call_soon from a just-finished warmup hits the 405 gate. + request_id = f"video-{uuid.uuid4().hex[:12]}" + active_request_id = request_id + interrupt_event.clear() + # Free prefill capacity for the query; completed warmup blocks + # stay in the prefix cache and are reused by it. + await _cancel_warmup() + if prev_was_interrupted and prev_request_id and self._engine_client: try: await self._engine_client.abort(prev_request_id) except Exception: pass prev_was_interrupted = False - - request_id = f"video-{uuid.uuid4().hex[:12]}" - active_request_id = request_id - interrupt_event.clear() query_frames = list(frame_buffer) query_frame_metadata = list(frame_metadata) query_audio_buffer = bytearray(audio_buffer) audio_buffer.clear() - query_prewarmed_frames = dict(frame_pil_cache) + # Legacy stride positions ignore decode status. Waiting on the + # whole buffer would stall the query on a frame it will not use. + if self._incremental_prefill_active(config): + wait_indices = list(range(len(query_frames))) + else: + wait_indices = self._candidate_frame_indices(query_frames, config.num_frames) + wait_frames = [query_frames[index] for index in wait_indices] + _pin_frames(wait_frames) async def _run_query() -> None: nonlocal active_request_id, prev_request_id try: + await _await_frames_ready(wait_frames) + query_prewarmed_frames = _snapshot_prewarmed(query_frames) process_kwargs: dict[str, Any] = {} if any(metadata.get("frame_id") for metadata in query_frame_metadata): process_kwargs["frame_metadata"] = query_frame_metadata @@ -349,19 +591,25 @@ class OmniStreamingVideoHandler: **process_kwargs, ) finally: + _unpin_frames(wait_frames) if active_request_id == request_id: prev_request_id = request_id active_request_id = None + # Warm frames that arrived during the turn (fires after + # this task completes; gated on no active query). + asyncio.get_running_loop().call_soon(_maybe_start_warmup) query_task = asyncio.create_task(_run_query()) async def _processor() -> None: """Process enqueued messages.""" nonlocal active_request_id, prev_request_id, prev_was_interrupted, query_task + nonlocal session_alive while True: msg = await msg_queue.get() if msg is None: + session_alive = False await _cancel_active_query(abort_now=True) return @@ -376,8 +624,8 @@ class OmniStreamingVideoHandler: ] frame_buffer[:] = [frame_buffer[index] for index in retained_indices] frame_metadata[:] = [frame_metadata[index] for index in retained_indices] - if frame_pil_cache.get(frame_data) is _BAD_FRAME: - frame_pil_cache.pop(frame_data, None) + if frame_pil_cache.get(frame_data) is _BAD_FRAME and frame_data not in pinned_frame_refs: + _drop_frame_cache(frame_data) if removed: await self._send_error(websocket, "Frame decode failed") @@ -410,10 +658,24 @@ class OmniStreamingVideoHandler: max_buf = config.max_frames dropped_frame_id: str | None = None if len(frame_buffer) >= max_buf: - dropped = frame_buffer.pop(0) - dropped_metadata = frame_metadata.pop(0) - dropped_frame_id = dropped_metadata.get("frame_id") - frame_pil_cache.pop(dropped, None) + # Legacy drops one frame. Incremental drops the oldest half so later frames append. + drop_count = max(1, max_buf // 2) if self._incremental_prefill_active(config) else 1 + for _ in range(drop_count): + dropped = frame_buffer.pop(0) + dropped_metadata = frame_metadata.pop(0) + if dropped_frame_id is None: + dropped_frame_id = dropped_metadata.get("frame_id") + # Duplicate frames share one cache entry keyed by + # b64: keep it while any duplicate remains buffered. + if dropped not in pinned_frame_refs and dropped not in frame_buffer: + _drop_frame_cache(dropped) + if drop_count > 1: + logger.warning( + "incremental prefill dropped %d oldest frames at max_frames=%d; " + "buffered_frames is the remaining suffix", + drop_count, + max_buf, + ) frame_buffer.append(frame_data) frame_metadata.append( { @@ -435,17 +697,23 @@ class OmniStreamingVideoHandler: ) # Prewarm: decode PIL off the event loop so query-time chat_template # can skip base64+Image.open. uuid=md5 lets mm_cache dedupe identical frames. - if frame_data not in frame_pil_cache: + if frame_data not in frame_pil_cache and frame_data not in frame_ready: mm_uuid = hashlib.md5(raw_bytes, usedforsecurity=False).hexdigest() + frame_ready[frame_data] = asyncio.Event() async def _prewarm(b64: str, b: bytes, u: str) -> None: + cancelled = False try: pil = await asyncio.to_thread(_decode_frame_bytes, b) - # The frame may have been evicted while decoding. - if b64 in frame_buffer: + # The frame may have been evicted (and unpinned) + # while decoding; frame_ready tracks liveness. + if b64 in frame_ready: frame_pil_cache[b64] = (pil, u) + except asyncio.CancelledError: + cancelled = True + raise except Exception: - if b64 not in frame_buffer: + if b64 not in frame_ready: return frame_pil_cache[b64] = _BAD_FRAME logger.warning("prewarm decode failed for frame (len=%d)", len(b)) @@ -455,10 +723,18 @@ class OmniStreamingVideoHandler: logger.warning( "frame decode failure event dropped because message queue is full" ) + finally: + event = frame_ready.get(b64) + if event is not None: + event.set() + if not cancelled: + asyncio.get_running_loop().call_soon(_maybe_start_warmup) task = asyncio.create_task(_prewarm(frame_data, raw_bytes, mm_uuid)) prewarm_tasks.add(task) task.add_done_callback(prewarm_tasks.discard) + elif frame_data in frame_pil_cache: + frame_ready.setdefault(frame_data, asyncio.Event()).set() is_generating = active_request_id is not None or ( query_task is not None and not query_task.done() @@ -471,6 +747,8 @@ class OmniStreamingVideoHandler: ) ): await _start_query_turn(query_text="") + else: + _maybe_start_warmup() elif msg_type == "audio.chunk": data_b64 = msg.get("data", "") @@ -501,6 +779,7 @@ class OmniStreamingVideoHandler: await _start_query_turn(query_text=query_text) elif msg_type == "video.done": + session_alive = False if query_task is not None and not query_task.done(): await asyncio.gather(query_task, return_exceptions=True) query_task = None @@ -520,6 +799,7 @@ class OmniStreamingVideoHandler: try: await _processor() finally: + session_alive = False reader_task.cancel() try: await reader_task @@ -529,6 +809,7 @@ class OmniStreamingVideoHandler: t.cancel() if prewarm_tasks: await asyncio.gather(*prewarm_tasks, return_exceptions=True) + await _cancel_warmup() if query_task is not None and not query_task.done(): await _cancel_active_query(abort_now=True) @@ -643,9 +924,10 @@ class OmniStreamingVideoHandler: ChatCompletionRequest, ) + reuse_active = self._incremental_prefill_active(config) # Select once from this turn's snapshot so prompt images and the # consumed event refer to exactly the same frames, including decode failures. - frame_indices = self._sample_frame_indices(frame_buffer, config.num_frames, prewarmed_frames) + frame_indices = self._select_frame_indices(config, frame_buffer, prewarmed_frames) messages, user_message = self.build_engine_prompt( config, frame_buffer, @@ -665,7 +947,9 @@ class OmniStreamingVideoHandler: "continue_final_message": False, "add_special_tokens": False, } - if config.use_audio_in_video and len(audio_buffer) > 0: + # Incremental: same kwargs as warmup so frame mm hashes stay a prefix. + # Legacy: only when this query actually has input audio. + if config.use_audio_in_video and (reuse_active or len(audio_buffer) > 0): request_kwargs["mm_processor_kwargs"] = { "use_audio_in_video": True, } @@ -680,97 +964,142 @@ class OmniStreamingVideoHandler: await self._send_error(websocket, f"Failed to build request: {e}") return - try: - engine_prompt = await self._preprocess_to_engine_prompt(chat_request) - except Exception as e: - await self._send_error(websocket, f"Preprocess failed: {e}") - return - decoded_ready_ts_ms = _time.monotonic() * 1000 - selected_metadata = [frame_metadata[index] for index in frame_indices] if frame_metadata else [] - model_selected_ts_ms = _time.monotonic() * 1000 + submitted = asyncio.get_running_loop().create_future() - await websocket.send_json({"type": "response.start"}) - text_parts: list[str] = [] - text_done_sent = False - audio_chunk_count = 0 - # Count emitted tensors only to identify the first audio emission. - # Default and explicit sampling parameters use DELTA audio outputs. - audio_chunks_drained = 0 - previous_text = "" - interrupted = False - frames_consumed_sent = False - t_start = _time.monotonic() - t_first_text = None - t_first_audio = None + async def _drive_query() -> None: + decoded_ready_ts_ms = _time.monotonic() * 1000 + selected_metadata = [frame_metadata[index] for index in frame_indices] if frame_metadata else [] + model_selected_ts_ms = _time.monotonic() * 1000 - # Wire-level async-chunk switch. "off" means - # buffer all deltas server-side and flush once at the end; the engine - # pipeline still overlaps internally. - async_chunk_mode = video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK - streaming = async_chunk_mode == "on" - audio_tail_tensors: list[Any] = [] + if not interrupt_event.is_set(): + await websocket.send_json({"type": "response.start"}) + text_parts: list[str] = [] + text_done_sent = False + audio_chunk_count = 0 + # Count emitted tensors only to identify the first audio emission. + # Default and explicit sampling parameters use DELTA audio outputs. + audio_chunks_drained = 0 + previous_text = "" + interrupted = False + frames_consumed_sent = False + t_start = _time.monotonic() + t_first_text = None + t_first_audio = None - try: - result_gen = engine_client.generate( - prompt=engine_prompt, - request_id=request_id, - output_modalities=config.modalities, - **sampling_kwargs, - ) + # Wire-level async-chunk switch. "off" means + # buffer all deltas server-side and flush once at the end; the engine + # pipeline still overlaps internally. + async_chunk_mode = video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK + streaming = async_chunk_mode == "on" + audio_tail_tensors: list[Any] = [] - async for output in result_gen: - # Soft interrupt: drain without sending - if interrupt_event.is_set(): - if not interrupted: - logger.info("Generation interrupted — draining") - interrupted = True - continue + try: + assert engine_client is not None + result_gen = engine_client.generate( + prompt=engine_prompt, + request_id=request_id, + output_modalities=config.modalities, + submitted=submitted, + **sampling_kwargs, + ) - if not isinstance(output, OmniRequestOutput): - continue + async for output in result_gen: + # Soft interrupt: drain without sending + if interrupt_event.is_set(): + if not interrupted: + logger.info("Generation interrupted — draining") + interrupted = True + continue - if not frames_consumed_sent and frame_metadata: - await websocket.send_json( - { - "type": "video.frames.consumed", - "request_id": request_id, - "model_selected_ts_ms": model_selected_ts_ms, - "frame_ids": [ - metadata["frame_id"] - for metadata in selected_metadata - if isinstance(metadata.get("frame_id"), str) - ], - "frames": [ - { - "frame_id": metadata.get("frame_id"), - "pts_ms": metadata.get("pts_ms"), - "source_pts_ms": metadata.get("source_pts_ms"), - "quality_profile": metadata.get("quality_profile"), - "receiver_received_ts_ms": metadata.get("receiver_received_ts_ms"), - "decoded_ready_ts_ms": decoded_ready_ts_ms, - } - for metadata in selected_metadata - ], - "latest_pts_ms": selected_metadata[-1].get("pts_ms") if selected_metadata else None, - } - ) - frames_consumed_sent = True + if not isinstance(output, OmniRequestOutput): + continue - out_type = getattr(output, "final_output_type", "text") + if not frames_consumed_sent and frame_metadata: + await websocket.send_json( + { + "type": "video.frames.consumed", + "request_id": request_id, + "model_selected_ts_ms": model_selected_ts_ms, + "frame_ids": [ + metadata["frame_id"] + for metadata in selected_metadata + if isinstance(metadata.get("frame_id"), str) + ], + "frames": [ + { + "frame_id": metadata.get("frame_id"), + "pts_ms": metadata.get("pts_ms"), + "source_pts_ms": metadata.get("source_pts_ms"), + "quality_profile": metadata.get("quality_profile"), + "receiver_received_ts_ms": metadata.get("receiver_received_ts_ms"), + "decoded_ready_ts_ms": decoded_ready_ts_ms, + } + for metadata in selected_metadata + ], + "latest_pts_ms": selected_metadata[-1].get("pts_ms") if selected_metadata else None, + } + ) + frames_consumed_sent = True - if out_type == "audio": - if streaming and not text_done_sent: - full_text = "".join(text_parts) - await websocket.send_json({"type": "response.text.done", "text": full_text}) - text_done_sent = True + out_type = getattr(output, "final_output_type", "text") - if t_first_audio is None: - t_first_audio = _time.monotonic() - audio_chunk_count += 1 - if streaming: - b64, audio_chunks_drained = self._extract_audio_delta_b64( + if out_type == "audio": + if streaming and not text_done_sent: + full_text = "".join(text_parts) + await websocket.send_json({"type": "response.text.done", "text": full_text}) + text_done_sent = True + + if t_first_audio is None: + t_first_audio = _time.monotonic() + audio_chunk_count += 1 + if streaming: + b64, audio_chunks_drained = self._extract_audio_delta_b64( + output, + audio_chunks_drained, + ) + if b64: + await websocket.send_json( + { + "type": "response.output_audio.delta", + "data": b64, + "format": "wav", + } + ) + else: + audio_data = self._get_audio_data(output) + if audio_data is not None: + if isinstance(audio_data, list): + audio_tail_tensors.extend(audio_data) + else: + audio_tail_tensors.append(audio_data) + else: + delta_text, previous_text = self._extract_text_delta( output, - audio_chunks_drained, + previous_text, + ) + if delta_text: + if t_first_text is None: + t_first_text = _time.monotonic() + text_parts.append(delta_text) + if streaming: + await websocket.send_json({"type": "response.text.delta", "delta": delta_text}) + + if not text_done_sent: + full_text = "".join(text_parts) + await websocket.send_json({"type": "response.text.done", "text": full_text}) + text_done_sent = True + + if not streaming and audio_tail_tensors: + try: + coalesced = audio_tail_tensors[0] + if len(audio_tail_tensors) != 1: + coalesced = torch.cat(audio_tail_tensors, dim=-1) + tail_np = self._tensor_to_1d_np(coalesced) + b64, _ = self._encode_tail( + tail_np, + 0, + new_drained=len(audio_tail_tensors), + is_first=True, ) if b64: await websocket.send_json( @@ -780,76 +1109,45 @@ class OmniStreamingVideoHandler: "format": "wav", } ) - else: - audio_data = self._get_audio_data(output) - if audio_data is not None: - if isinstance(audio_data, list): - audio_tail_tensors.extend(audio_data) - else: - audio_tail_tensors.append(audio_data) - else: - delta_text, previous_text = self._extract_text_delta( - output, - previous_text, - ) - if delta_text: - if t_first_text is None: - t_first_text = _time.monotonic() - text_parts.append(delta_text) - if streaming: - await websocket.send_json({"type": "response.text.delta", "delta": delta_text}) + except Exception: + logger.exception("Failed to coalesce off-path audio") - if not text_done_sent: + if audio_chunk_count > 0: + await websocket.send_json({"type": "response.output_audio.done"}) + + response_text = "".join(text_parts) + self.on_turn_complete(message_history, user_message, response_text) + + t_end = _time.monotonic() + logger.info( + "[TIMING] mode=%s total=%.2fs first_text=%.2fs first_audio=%.2fs audio_chunks=%d", + async_chunk_mode, + t_end - t_start, + (t_first_text - t_start) if t_first_text else -1, + (t_first_audio - t_start) if t_first_audio else -1, + audio_chunk_count, + ) + + except Exception: + logger.exception("Engine query failed") + await self._send_error(websocket, "Query processing failed") + + if not text_done_sent and not interrupt_event.is_set(): full_text = "".join(text_parts) await websocket.send_json({"type": "response.text.done", "text": full_text}) - text_done_sent = True - if not streaming and audio_tail_tensors: - try: - coalesced = ( - audio_tail_tensors[0] if len(audio_tail_tensors) == 1 else torch.cat(audio_tail_tensors, dim=-1) - ) - tail_np = self._tensor_to_1d_np(coalesced) - b64, _ = self._encode_tail( - tail_np, - 0, - new_drained=len(audio_tail_tensors), - is_first=True, - ) - if b64: - await websocket.send_json( - { - "type": "response.output_audio.delta", - "data": b64, - "format": "wav", - } - ) - except Exception: - logger.exception("Failed to coalesce off-path audio") - - if audio_chunk_count > 0: - await websocket.send_json({"type": "response.output_audio.done"}) - - response_text = "".join(text_parts) - self.on_turn_complete(message_history, user_message, response_text) - - t_end = _time.monotonic() - logger.info( - "[TIMING] mode=%s total=%.2fs first_text=%.2fs first_audio=%.2fs audio_chunks=%d", - async_chunk_mode, - t_end - t_start, - (t_first_text - t_start) if t_first_text else -1, - (t_first_audio - t_start) if t_first_audio else -1, - audio_chunk_count, - ) - - except Exception: - logger.exception("Engine query failed") - await self._send_error(websocket, "Query processing failed") - - if not text_done_sent: - full_text = "".join(text_parts) - await websocket.send_json({"type": "response.text.done", "text": full_text}) + prompt_task = asyncio.ensure_future(self._preprocess_to_engine_prompt(chat_request)) + try: + engine_prompt = await asyncio.shield(prompt_task) + except asyncio.CancelledError: + # Sender cache is updated in preprocess; submit before the caller aborts. + engine_prompt = await asyncio.shield(prompt_task) + await self._stop_after_submit(asyncio.ensure_future(_drive_query()), submitted) + raise + except Exception as e: + await self._send_error(websocket, f"Preprocess failed: {e}") + return + await self._await_until_submitted(asyncio.ensure_future(_drive_query()), submitted) # ------------------------------------------------------------------ # Audio helpers @@ -1009,6 +1307,105 @@ class OmniStreamingVideoHandler: return text[len(previous_text) :], text return text, text + # ------------------------------------------------------------------ + # Incremental prefill (context warmup) + # ------------------------------------------------------------------ + + async def _stop_after_submit(self, drive: asyncio.Future[Any], submitted: asyncio.Future[None]) -> None: + """Stage0 must observe the request before the caller aborts it.""" + if not submitted.done(): + submit_wait = asyncio.ensure_future(asyncio.shield(submitted)) + await asyncio.wait({submit_wait, drive}, return_when=asyncio.FIRST_COMPLETED) + if not submitted.done(): + submit_wait.cancel() + exc = drive.exception() if drive.done() else None + if exc is not None: + raise exc + if not drive.done(): + drive.cancel() + await asyncio.gather(drive, return_exceptions=True) + + async def _await_until_submitted(self, drive: asyncio.Future[Any], submitted: asyncio.Future[None]) -> None: + """On cancel, wait until Stage0 has accepted the request before aborting.""" + try: + # Shield so cancelling this task does not cancel generate before Stage0 submit. + await asyncio.shield(drive) + except asyncio.CancelledError: + await self._stop_after_submit(drive, submitted) + raise + + async def _drive_warmup_generate( + self, + engine_client: Any, + engine_prompt: Any, + request_id: str, + modalities: list[str], + *, + stop_after_submit: bool = False, + ) -> None: + """max_tokens=1 generate; cancellation still submits before abort.""" + from vllm import SamplingParams + + submitted = asyncio.get_running_loop().create_future() + result_gen = engine_client.generate( + prompt=engine_prompt, + request_id=request_id, + sampling_params=SamplingParams(max_tokens=1, temperature=0.0), + output_modalities=modalities, + submitted=submitted, + ) + + async def _drive() -> None: + async for _ in result_gen: + pass + + drive = asyncio.ensure_future(_drive()) + if stop_after_submit: + await self._stop_after_submit(drive, submitted) + return + await self._await_until_submitted(drive, submitted) + + async def _prefill_context( + self, + config: StreamingVideoSessionConfig, + messages: list[dict[str, Any]], + request_id: str, + ) -> None: + """max_tokens=1 generate so arriving frames land in the prefix cache.""" + from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ) + + engine_client = self._engine_client + if engine_client is None: + return + + request_kwargs: dict[str, Any] = { + "model": config.model or "default", + "messages": messages, + "stream": True, + "modalities": config.modalities, + "add_generation_prompt": True, + "continue_final_message": False, + "add_special_tokens": False, + } + if config.use_audio_in_video: + request_kwargs["mm_processor_kwargs"] = { + "use_audio_in_video": True, + } + chat_request = ChatCompletionRequest(**request_kwargs) + prompt_task = asyncio.ensure_future(self._preprocess_to_engine_prompt(chat_request)) + try: + engine_prompt = await asyncio.shield(prompt_task) + except asyncio.CancelledError: + # Sender cache is updated in preprocess; submit before the caller aborts. + engine_prompt = await asyncio.shield(prompt_task) + await self._drive_warmup_generate( + engine_client, engine_prompt, request_id, config.modalities, stop_after_submit=True + ) + raise + await self._drive_warmup_generate(engine_client, engine_prompt, request_id, config.modalities) + # ------------------------------------------------------------------ # Preprocessing # ------------------------------------------------------------------ @@ -1043,6 +1440,18 @@ class OmniStreamingVideoHandler: except Exception: pass + @staticmethod + def _candidate_frame_indices(frame_buffer: list[str], num_frames: int) -> list[int]: + """Stride positions for a legacy query, including the last frame. + + Ignores decode status so a frame still decoding is waited on, not skipped. + """ + n_buf = len(frame_buffer) + if n_buf <= num_frames: + return list(range(n_buf)) + stride = max(1, n_buf // num_frames) + return [index * stride for index in range(num_frames - 1)] + [n_buf - 1] + @staticmethod def _sample_frame_indices( frame_buffer: list[str], @@ -1050,13 +1459,26 @@ class OmniStreamingVideoHandler: prewarmed_frames: Mapping[str, PrewarmedFrame], ) -> list[int]: """Stride-sample with the last frame, then drop known bad frames without refilling.""" - n_buf = len(frame_buffer) - if n_buf <= num_frames: - indices = list(range(n_buf)) - else: - stride = max(1, n_buf // num_frames) - indices = [index * stride for index in range(num_frames - 1)] + [n_buf - 1] - return [index for index in indices if prewarmed_frames.get(frame_buffer[index]) is not _BAD_FRAME] + indices = OmniStreamingVideoHandler._candidate_frame_indices(frame_buffer, num_frames) + return [index for index in indices if _usable_prewarmed_frame(prewarmed_frames.get(frame_buffer[index]))] + + def _select_frame_indices( + self, + config: StreamingVideoSessionConfig, + frame_buffer: list[str], + prewarmed_frames: Mapping[str, PrewarmedFrame], + ) -> list[int]: + """This turn's frame selection, shared by the prompt and the consumed event.""" + if self._incremental_prefill_active(config): + # Frames are sampled at arrival time only (EVS + max_frames), so + # the buffer already is the warmed sequence. Query-time + # subsampling would pick a subset, which is not a prefix. + return [ + index + for index, frame_b64 in enumerate(frame_buffer) + if _usable_prewarmed_frame(prewarmed_frames.get(frame_b64)) + ] + return self._sample_frame_indices(frame_buffer, config.num_frames, prewarmed_frames) @staticmethod async def _send_frame_ack( From 0e6791e9c4925c3539352cff2c2738bd48a1a65a Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Fri, 2 Oct 2026 10:22:52 -0400 Subject: [PATCH 16/22] [Bugfix] Initialize eager state in Higgs test fixture (#8426) Signed-off-by: Nick Cao Co-authored-by: Codex --- tests/model_executor/models/higgs_audio_v3/test_mrv2_state.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/model_executor/models/higgs_audio_v3/test_mrv2_state.py b/tests/model_executor/models/higgs_audio_v3/test_mrv2_state.py index e0210295a..ccf2eb5ef 100644 --- a/tests/model_executor/models/higgs_audio_v3/test_mrv2_state.py +++ b/tests/model_executor/models/higgs_audio_v3/test_mrv2_state.py @@ -13,6 +13,7 @@ from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import HiggsAudioV3TalkerForConditionalGeneration, ) from vllm_omni.model_executor.models.higgs_audio_v3.model_state import HiggsModelState +from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState from vllm_omni.worker_v2.model_states.intermediate_buffer import OmniIntermediateBuffer pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -34,6 +35,8 @@ def make_state(): _audio_continuation_id=99, update_decode_step_metadata=Mock(), ) + state._stream_pos = {} + state._eager_state = EagerMTPState(state) return state From 13c66da75219d8c5ab752618039d254dd9a19eb6 Mon Sep 17 00:00:00 2001 From: Yueqian Lin Date: Sat, 3 Oct 2026 01:12:29 +1000 Subject: [PATCH 17/22] [Perf][AuK] Precompute adaLN modulations per request and run the conv position embedding as one GEMM (#8328) Signed-off-by: Yueqian Lin --- .../models/auk/test_auk_cudagraph.py | 15 +- .../models/auk/test_auk_dit_fastpaths.py | 139 ++++++++++++++++ .../models/auk/test_auk_step_context.py | 4 +- .../diffusion/models/auk/auk_transformer.py | 152 +++++++++++++++--- .../diffusion/models/auk/cudagraph_wrapper.py | 49 ++++-- .../diffusion/models/auk/pipeline_auk.py | 15 +- 6 files changed, 329 insertions(+), 45 deletions(-) create mode 100644 tests/diffusion/models/auk/test_auk_dit_fastpaths.py diff --git a/tests/diffusion/models/auk/test_auk_cudagraph.py b/tests/diffusion/models/auk/test_auk_cudagraph.py index 663cdf8df..c140b8d6c 100644 --- a/tests/diffusion/models/auk/test_auk_cudagraph.py +++ b/tests/diffusion/models/auk/test_auk_cudagraph.py @@ -61,7 +61,10 @@ def test_single_request_graph_wrapper_cpu_falls_back_to_eager(cfg_strength: floa # The conditioning is prepared once per request and reused by every step. assert prepare_spy.call_count == 1 - assert prepare_spy.call_args.args[-1] is (cfg_strength >= 1e-5) + # _prepare(x, x_mask, text, c_mask, ref, ref_mask, uses_cfg, timesteps) + assert prepare_spy.call_args.args[6] is (cfg_strength >= 1e-5) + # The time grid reaches prepare, so the adaLN modulations are computed once. + assert prepare_spy.call_args.args[7].numel() == 2 assert step_spy.call_count == 2 capture_spy.assert_not_called() assert not wrapper._cache @@ -84,7 +87,8 @@ def test_graph_inputs_use_bounded_length_buckets() -> None: assert padded[4].shape == (1, 100, 4) assert padded[5].shape == (1, 100) assert [mask.sum().item() for mask in (padded[1], padded[3], padded[5])] == [65, 65, 51] - assert wrapper._key(padded[0], padded[2], padded[4], False) == (96, 96, 100, False) + assert wrapper._key(padded[0], padded[2], padded[4], False) == (96, 96, 100, False, 0) + assert wrapper._key(padded[0], padded[2], padded[4], False, 32) == (96, 96, 100, False, 32) assert wrapper.max_graphs == 32 @@ -153,18 +157,21 @@ def test_single_request_graph_replay_matches_eager_and_updates_inputs(cfg_streng inputs["ref"], inputs["ref_mask"], ) - key = wrapper._key(bucketed[0], bucketed[2], bucketed[4], cfg_strength >= 1e-5) + # sample_latents passes the grid, so the key carries its step count (two steps here). + key = wrapper._key(bucketed[0], bucketed[2], bucketed[4], cfg_strength >= 1e-5, 2) assert key in wrapper._cache # The static context was refreshed for this request's conditioning. entry = wrapper._cache[key] - fresh = wrapper._prepare(*bucketed, cfg_strength >= 1e-5) + grid = torch.tensor([0.0, 0.4], device="cuda") + fresh = wrapper._prepare(*bucketed, cfg_strength >= 1e-5, grid) for static, want in zip(entry.static_ctx.tensors(), fresh.tensors(), strict=True): if want is None: assert static is None else: torch.testing.assert_close(static, want) torch.testing.assert_close(entry.static_timestep, torch.tensor(0.4, device="cuda")) + assert int(entry.static_step) == 1 assert len(wrapper._cache) == 1 diff --git a/tests/diffusion/models/auk/test_auk_dit_fastpaths.py b/tests/diffusion/models/auk/test_auk_dit_fastpaths.py new file mode 100644 index 000000000..97d97e619 --- /dev/null +++ b/tests/diffusion/models/auk/test_auk_dit_fastpaths.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The AuK DiT fast paths keep the reference numerics. + +Two paths are checked: the grouped conv position embedding as one batched +GEMM, and the adaLN modulations precomputed for a whole time grid, eager and +under the per-step CUDA graph. +""" + +import pytest +import torch + +from vllm_omni.diffusion.models.auk import auk_transformer as T +from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, build_time_grid +from vllm_omni.diffusion.models.auk.cudagraph_wrapper import AuKCUDAGraphWrapper + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_dit(device: str = "cpu") -> AuKTransformer: + torch.manual_seed(5) + return ( + AuKTransformer( + dim=32, + heads=2, + dim_head=16, + ff_mult=2, + latent_dim=4, + text_hidden_dim=8, + num_layers=2, + num_single_layers=2, + ) + .eval() + .to(device) + ) + + +def _reference_conv_pos_embed(module: T.ConvPosEmbedding, x: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor: + """The channels-first conv stack as the reference implementation writes it.""" + keep = None if mask is None else mask.unsqueeze(1) + x = x.transpose(1, 2) + if keep is not None: + x = x.masked_fill(~keep, 0.0) + for layer in module.conv1d: + x = layer(x) + if keep is not None and isinstance(layer, torch.nn.Conv1d): + x = x.masked_fill(~keep, 0.0) + return x.transpose(1, 2) + + +@torch.inference_mode() +def test_grouped_conv_as_batched_gemm_matches_conv1d() -> None: + torch.manual_seed(0) + conv = torch.nn.Conv1d(64, 64, 31, groups=16, padding=15) + x = torch.randn(2, 37, 64) + expected = conv(x.transpose(1, 2)).transpose(1, 2) + torch.testing.assert_close(T._grouped_conv1d_btc(x, conv), expected, atol=1e-5, rtol=1e-5) + + +@torch.inference_mode() +@pytest.mark.parametrize("masked", [False, True]) +def test_conv_pos_embedding_matches_the_channels_first_stack(masked: bool) -> None: + torch.manual_seed(1) + module = T.ConvPosEmbedding(64).eval() + x = torch.randn(2, 37, 64) + mask = None + if masked: + mask = torch.ones(2, 37, dtype=torch.bool) + mask[1, 30:] = False + torch.testing.assert_close(module(x, mask), _reference_conv_pos_embed(module, x, mask), atol=1e-5, rtol=1e-5) + + +@torch.inference_mode() +@pytest.mark.parametrize("cfg_infer", [False, True]) +def test_modulation_table_matches_per_step_modulation(cfg_infer: bool) -> None: + dit = _make_dit() + text = torch.randn(1, 5, 8) + ref = torch.randn(1, 3, 4) + x = torch.randn(1, 6, 4) + grid = build_time_grid(nfe=4, sway_sampling_coef=-1.0, t_grid=None, device="cpu") + plain = dit.prepare(text, target_len=6, ref=ref, cfg_infer=cfg_infer) + tabled = dit.prepare(text, target_len=6, ref=ref, cfg_infer=cfg_infer, timesteps=grid[:-1]) + assert tabled.modulation is not None and tabled.modulation.shape[0] == 4 + for i in range(4): + expected = dit.step(x, grid[i], plain) + torch.testing.assert_close(dit.step(x, grid[i], tabled, step_index=i), expected, atol=1e-5, rtol=1e-5) + index = torch.tensor(i) + torch.testing.assert_close(dit.step(x, grid[i], tabled, step_index=index), expected, atol=1e-5, rtol=1e-5) + + +def _padded_inputs(device: str) -> dict[str, torch.Tensor]: + generator = torch.Generator(device="cpu").manual_seed(3) + text = torch.randn(1, 7, 8, generator=generator) + c_mask = torch.ones(1, 7, dtype=torch.bool) + c_mask[:, 5:] = False + text[:, 5:] = 0.0 + ref = torch.randn(1, 6, 4, generator=generator) + ref_mask = torch.ones(1, 6, dtype=torch.bool) + ref_mask[:, 4:] = False + mask = torch.ones(1, 9, dtype=torch.bool) + mask[:, 7:] = False + inputs = {"text": text, "c_mask": c_mask, "ref": ref, "ref_mask": ref_mask, "mask": mask} + return {name: value.to(device) for name, value in inputs.items()} + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA") +@torch.inference_mode() +@pytest.mark.parametrize("cfg_strength", [0.0, 2.0]) +def test_graph_with_time_grid_matches_eager_with_time_grid(cfg_strength: float) -> None: + dit = _make_dit("cuda") + inputs = _padded_inputs("cuda") + grid = build_time_grid(nfe=4, sway_sampling_coef=-1.0, t_grid=None, device="cuda") + x = torch.randn(1, 9, 4, device="cuda") + wrapper = AuKCUDAGraphWrapper(dit) + common = dict( + text=inputs["text"], + c_mask=inputs["c_mask"], + ref=inputs["ref"], + ref_mask=inputs["ref_mask"], + cfg_strength=cfg_strength, + timesteps=grid[:-1], + ) + for i in range(4): + graph = wrapper(x=x, timestep=grid[i], new_request=i == 0, step_index=i, **common) + ctx = dit.prepare( + inputs["text"], + target_len=9, + c_mask=inputs["c_mask"], + ref=inputs["ref"], + ref_mask=inputs["ref_mask"], + cfg_infer=cfg_strength >= 1e-5, + timesteps=grid[:-1], + ) + eager = wrapper._step(x, grid[i], ctx, cfg_strength, i) + torch.testing.assert_close(graph, eager, atol=1e-4, rtol=1e-4) + # One graph serves every step of the grid. + assert len(wrapper._cache) == 1 + assert next(iter(wrapper._cache))[-1] == 4 diff --git a/tests/diffusion/models/auk/test_auk_step_context.py b/tests/diffusion/models/auk/test_auk_step_context.py index e70d52510..26f722fe3 100644 --- a/tests/diffusion/models/auk/test_auk_step_context.py +++ b/tests/diffusion/models/auk/test_auk_step_context.py @@ -127,7 +127,9 @@ def test_hoisted_euler_loop_matches_per_step_forward(cfg_strength: float) -> Non v = dit(x, time=timesteps[i], **inputs) x = x + (timesteps[i + 1] - timesteps[i]) * v - torch.testing.assert_close(got, x, rtol=0, atol=0) + # sample_latents precomputes the adaLN modulations of the whole grid in one + # GEMM per layer, whose rows round differently from one-row GEMMs in fp32. + torch.testing.assert_close(got, x, rtol=0, atol=1e-6) assert dit.text_cond is None and dit.text_uncond is None diff --git a/vllm_omni/diffusion/models/auk/auk_transformer.py b/vllm_omni/diffusion/models/auk/auk_transformer.py index 4056c72eb..fb146b101 100644 --- a/vllm_omni/diffusion/models/auk/auk_transformer.py +++ b/vllm_omni/diffusion/models/auk/auk_transformer.py @@ -153,6 +153,33 @@ class TimeEmbedding(nn.Module): return self.time_mlp(hidden.to(self.time_mlp[0].weight.dtype)) +def _grouped_conv1d_btc(x: torch.Tensor, conv: nn.Conv1d) -> torch.Tensor: + """``conv`` applied to channels-last ``x`` ``[B, T, C]``, returning ``[B, T, C]``. + + A same-padded grouped convolution written as one im2col copy and one + batched GEMM over the groups. cuDNN runs a grouped conv of this shape as + one kernel per group plus two layout transposes, which at AuK's sizes + (1536 channels in 16 groups, kernel 31, a few hundred frames) costs about + ten times the arithmetic. + """ + batch, frames, channels = x.shape + groups = conv.groups + kernel = conv.kernel_size[0] + per_group = channels // groups + pad = conv.padding[0] + padded = F.pad(x, (0, 0, pad, pad)) + # unfold gives a [B, T, C, K] view; one copy lays it out as [B, T, G, C/G * K]. + cols = padded.unfold(1, kernel, 1).reshape(batch * frames, groups, per_group * kernel) + # conv.weight [C_out, C/G, K] -> [G, C/G * K, C_out/G]; the reduction order matches the columns. + weight = conv.weight.reshape(groups, channels // groups, per_group * kernel).transpose(1, 2).to(cols.dtype) + # [G, B*T, C/G*K] @ [G, C/G*K, C_out/G] -> [G, B*T, C_out/G] + out = torch.bmm(cols.transpose(0, 1), weight) + out = out.transpose(0, 1).reshape(batch, frames, channels) + if conv.bias is not None: + out = out + conv.bias.to(out.dtype) + return out + + class ConvPosEmbedding(nn.Module): """Depthwise-grouped conv stack that adds local position information.""" @@ -169,17 +196,19 @@ class ConvPosEmbedding(nn.Module): ) def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor: - keep = None if mask is None else mask.unsqueeze(1) - x = x.transpose(1, 2) + keep = None if mask is None else mask.unsqueeze(-1) if keep is not None: x = x.masked_fill(~keep, 0.0) for layer in self.conv1d: - x = layer(x) - # Re-zero padding after each convolution so it cannot bleed into - # valid frames through the kernel window. - if keep is not None and isinstance(layer, nn.Conv1d): - x = x.masked_fill(~keep, 0.0) - return x.transpose(1, 2) + if isinstance(layer, nn.Conv1d): + x = _grouped_conv1d_btc(x, layer) + # Re-zero padding after each convolution so it cannot bleed into + # valid frames through the kernel window. + if keep is not None: + x = x.masked_fill(~keep, 0.0) + else: + x = layer(x) + return x class AudioEmbedding(nn.Module): @@ -205,9 +234,12 @@ class AdaLayerNorm(nn.Module): self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) def forward( - self, x: torch.Tensor, emb: torch.Tensor + self, x: torch.Tensor, emb: torch.Tensor | None, mod: torch.Tensor | None = None ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.linear(self.silu(emb)).chunk(6, dim=1) + """``mod`` is the precomputed ``linear(silu(emb))``; ``emb`` is then unused.""" + if mod is None: + mod = self.linear(self.silu(emb)) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=1) x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] return x, gate_msa, shift_mlp, scale_mlp, gate_mlp @@ -221,8 +253,11 @@ class AdaLayerNormFinal(nn.Module): self.linear = nn.Linear(dim, dim * 2) self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) - def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: - scale, shift = self.linear(self.silu(emb)).chunk(2, dim=1) + def forward(self, x: torch.Tensor, emb: torch.Tensor | None, mod: torch.Tensor | None = None) -> torch.Tensor: + """``mod`` is the precomputed ``linear(silu(emb))``; ``emb`` is then unused.""" + if mod is None: + mod = self.linear(self.silu(emb)) + scale, shift = mod.chunk(2, dim=1) return self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] @@ -383,9 +418,11 @@ class DoubleBlock(nn.Module): c_rope: RopeCache, c_mask: torch.Tensor | None, bias: torch.Tensor | None = None, + mod_c: torch.Tensor | None = None, + mod_x: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, t) - norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, t) + norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, t, mod_c) + norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, t, mod_x) x_attn, c_attn = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope, c_mask=c_mask, bias=bias) @@ -416,8 +453,9 @@ class SingleBlock(nn.Module): mask: torch.Tensor | None, rope: RopeCache, bias: torch.Tensor | None = None, + mod: torch.Tensor | None = None, ) -> torch.Tensor: - norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, t) + norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, t, mod) x = x + gate_msa[:, None] * self.attn(x=norm, mask=mask, rope=rope, bias=bias) norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] return x + gate_mlp[:, None] * self.ff(norm) @@ -491,6 +529,28 @@ class AuKTransformer(nn.Module): self.norm_out = AdaLayerNormFinal(dim) self.proj_out = nn.Linear(dim, latent_dim) + # Widths of the adaLN modulations in the order modulation_table() lays them out. + self._modulation_widths = [module.linear.out_features for module in self._adaln_modules()] + + def _adaln_modules(self) -> list[nn.Module]: + """Every timestep-conditioned norm, in the order :meth:`step` consumes them.""" + modules: list[nn.Module] = [] + for block in self.transformer_blocks: + modules += [block.attn_norm_c, block.attn_norm_x] + modules += [block.attn_norm for block in self.single_transformer_blocks] + modules.append(self.norm_out) + return modules + + def modulation_table(self, timesteps: torch.Tensor) -> torch.Tensor: + """adaLN modulations for every timestep, ``[len(timesteps), sum of widths]``. + + The modulations depend only on the timestep, and a request's time grid + is known before its first step. Computing them for the whole grid at + once reads each modulation weight once per request instead of once per + step, which at batch 1 is what those matrix-vector products cost. + """ + emb = F.silu(self.time_embed(timesteps.reshape(-1).float())) + return torch.cat([module.linear(emb) for module in self._adaln_modules()], dim=1) def project_text(self, text: torch.Tensor) -> torch.Tensor: """Project LLM hidden states ``[B, nt, text_hidden_dim]`` to model width.""" @@ -520,6 +580,7 @@ class AuKTransformer(nn.Module): drop_text: bool = False, cfg_infer: bool = False, cache: bool = False, + timesteps: torch.Tensor | None = None, ) -> "AuKStepContext": """Everything a denoise step needs that does not depend on ``x`` or the time. @@ -528,6 +589,9 @@ class AuKTransformer(nn.Module): prompt embedding, the padding masks and biases and the rotary tables are not recomputed on every step. Arguments are those of :meth:`forward`; ``target_len`` is the target frame count ``n``. + ``timesteps`` are the times the steps will run at (the grid without + its end point); when given, their adaLN modulations are precomputed + and :meth:`step` selects them by ``step_index``. """ batch = text.shape[0] if c_mask is None: @@ -593,25 +657,44 @@ class AuKTransformer(nn.Module): rope_text=_rope_cos_sin(self.rotary_embed(text_len, c_mask)), rope_single=_rope_cos_sin(self.rotary_embed(text_len + audio_len, single_mask)), branches=branches, + modulation=None if timesteps is None else self.modulation_table(timesteps), ) - def step(self, x: torch.Tensor, time: torch.Tensor, ctx: "AuKStepContext") -> torch.Tensor: + def step( + self, + x: torch.Tensor, + time: torch.Tensor, + ctx: "AuKStepContext", + step_index: torch.Tensor | int | None = None, + ) -> torch.Tensor: """Predict the velocity of ``x`` at ``time`` under a :meth:`prepare` context. The target embedding does not depend on the CFG branch, so it is - computed once and shared by both branches. + computed once and shared by both branches. When the context carries a + modulation table, ``step_index`` selects this step's row of it and + ``time`` is unused; the row broadcasts over the CFG branches. """ - if time.ndim == 0: - time = time.repeat(x.shape[0]) - t = self.time_embed(time) target = self.audio_embed(x, mask=ctx.target_mask) + if ctx.modulation is not None and step_index is not None: + if isinstance(step_index, int): + row = ctx.modulation[step_index : step_index + 1] + else: + row = ctx.modulation.index_select(0, step_index.reshape(1)) + mods = list(torch.split(row, self._modulation_widths, dim=1)) + t = None + else: + if time.ndim == 0: + time = time.repeat(x.shape[0]) + t = self.time_embed(time) + if ctx.branches == 2: + t = torch.cat((t, t), dim=0) + mods = [None] * len(self._modulation_widths) if ctx.branches == 2: - t = torch.cat((t, t), dim=0) target = torch.cat((target, target), dim=0) audio = target if ctx.prompt is None else torch.cat([ctx.prompt, target], dim=1) c = ctx.c - for block in self.transformer_blocks: + for i, block in enumerate(self.transformer_blocks): c, audio = block( audio, c, @@ -621,14 +704,24 @@ class AuKTransformer(nn.Module): c_rope=ctx.rope_text, c_mask=ctx.c_mask, bias=ctx.joint_bias, + mod_c=mods[2 * i], + mod_x=mods[2 * i + 1], ) h = torch.cat([c, audio], dim=1) - for block in self.single_transformer_blocks: - h = block(h, t, mask=ctx.single_mask, rope=ctx.rope_single, bias=ctx.single_bias) + offset = 2 * len(self.transformer_blocks) + for i, block in enumerate(self.single_transformer_blocks): + h = block( + h, + t, + mask=ctx.single_mask, + rope=ctx.rope_single, + bias=ctx.single_bias, + mod=mods[offset + i], + ) h = h[:, h.shape[1] - x.shape[1] :] - return self.proj_out(self.norm_out(h, t)) + return self.proj_out(self.norm_out(h, t, mods[-1])) def forward( self, @@ -706,6 +799,8 @@ class AuKStepContext: rope_text: RopeCache rope_single: RopeCache branches: int + #: adaLN modulations per step, ``[steps, widths]``, when the time grid was known. + modulation: torch.Tensor | None = None def tensors(self) -> list[torch.Tensor | None]: """Every tensor field in a fixed order, rope pairs flattened.""" @@ -721,6 +816,7 @@ class AuKStepContext: *self.rope_audio, *self.rope_text, *self.rope_single, + self.modulation, ] def copy_(self, other: "AuKStepContext") -> None: @@ -752,6 +848,7 @@ class AuKStepContext: rope_text=(self.rope_text[0].clone(), self.rope_text[1].clone()), rope_single=(self.rope_single[0].clone(), self.rope_single[1].clone()), branches=self.branches, + modulation=_clone(self.modulation), ) @@ -881,6 +978,8 @@ def sample_latents( timestep=t[i], cfg_strength=cfg_strength, new_request=i == 0, + timesteps=t[:-1], + step_index=i, ) x = x + (t[i + 1] - t[i]) * velocity finally: @@ -922,9 +1021,10 @@ def _sample_latents( ref_mask=ref_mask, cfg_infer=guided, cache=guided, + timesteps=timesteps[:-1], ) for i in range(timesteps.shape[0] - 1): - v = dit.step(x, timesteps[i], ctx) + v = dit.step(x, timesteps[i], ctx, step_index=i) if guided: v_cond, v_uncond = v.chunk(2, dim=0) v = v_cond + (v_cond - v_uncond) * cfg_strength diff --git a/vllm_omni/diffusion/models/auk/cudagraph_wrapper.py b/vllm_omni/diffusion/models/auk/cudagraph_wrapper.py index 7d34557f6..8bf09c2e4 100644 --- a/vllm_omni/diffusion/models/auk/cudagraph_wrapper.py +++ b/vllm_omni/diffusion/models/auk/cudagraph_wrapper.py @@ -24,16 +24,19 @@ class _GraphEntry: static_ctx: AuKStepContext static_timestep: torch.Tensor static_cfg: torch.Tensor | None + static_step: torch.Tensor | None static_out: torch.Tensor class AuKCUDAGraphWrapper: """Replay one DiT denoise step and leave Euler scheduling to the caller. - The graph is keyed by the target, text and reference sequence lengths plus - whether the CFG branch is enabled. Timestep and the CFG strength are + The graph is keyed by the target, text and reference sequence lengths, + whether the CFG branch is enabled and, when the caller passes the time + grid, its length. Timestep (or the step index), and the CFG strength are mutable scalar buffers, so all Euler steps and CFG values within one path - reuse one graph. + reuse one graph. With the time grid, the adaLN modulations of every step + are part of the per-request context and the graph selects its row. Only :meth:`AuKTransformer.step` is captured. The per-request context (text projection, reference embedding, padding biases, rotary tables) is @@ -85,8 +88,9 @@ class AuKCUDAGraphWrapper: text: torch.Tensor, ref: torch.Tensor, uses_cfg: bool, - ) -> tuple[int, int, int, bool]: - return (x.shape[1], text.shape[1], ref.shape[1], uses_cfg) + steps: int = 0, + ) -> tuple[int, int, int, bool, int]: + return (x.shape[1], text.shape[1], ref.shape[1], uses_cfg, steps) def _retire_graph_generation_if_full(self) -> None: """Retire all graphs together so none outlive shared workspaces.""" @@ -103,6 +107,7 @@ class AuKCUDAGraphWrapper: ref: torch.Tensor, ref_mask: torch.Tensor, uses_cfg: bool, + timesteps: torch.Tensor | None = None, ) -> AuKStepContext: return self.dit.prepare( text, @@ -112,6 +117,7 @@ class AuKCUDAGraphWrapper: ref=ref, ref_mask=ref_mask, cfg_infer=uses_cfg, + timesteps=timesteps, ) def _step( @@ -120,9 +126,10 @@ class AuKCUDAGraphWrapper: timestep: torch.Tensor, ctx: AuKStepContext, cfg_strength: torch.Tensor | float, + step_index: torch.Tensor | int | None = None, ) -> torch.Tensor: """One guided (or unguided) velocity; the CFG combine is part of the graph.""" - velocity = self.dit.step(x, timestep, ctx) + velocity = self.dit.step(x, timestep, ctx, step_index=step_index) if ctx.branches == 1: return velocity conditional, unconditional = velocity.chunk(2, dim=0) @@ -140,6 +147,8 @@ class AuKCUDAGraphWrapper: timestep: torch.Tensor, cfg_strength: float, new_request: bool = True, + timesteps: torch.Tensor | None = None, + step_index: int | None = None, ) -> torch.Tensor: """Velocity of ``x`` at ``timestep``. @@ -147,26 +156,33 @@ class AuKCUDAGraphWrapper: (``text``, ``c_mask``, ``ref``, ``ref_mask``) is then prepared anew; the later steps of the request pass ``False`` and reuse it. The default prepares on every call, which is always correct. + + ``timesteps`` (the times every step of the request runs at) and + ``step_index`` (this step's position in them) let the adaLN + modulations be computed once per request instead of once per step. """ uses_cfg = cfg_strength >= 1e-5 + if timesteps is None or step_index is None: + timesteps, step_index = None, None if not self.enabled or x.device.type != "cuda" or torch.cuda.is_current_stream_capturing(): if new_request or self._eager_ctx is None: - self._eager_ctx = self._prepare(x, None, text, c_mask, ref, ref_mask, uses_cfg) - return self._step(x, timestep, self._eager_ctx, cfg_strength) + self._eager_ctx = self._prepare(x, None, text, c_mask, ref, ref_mask, uses_cfg, timesteps) + return self._step(x, timestep, self._eager_ctx, cfg_strength, step_index) target_frames = x.shape[1] x, x_mask, text, c_mask, ref, ref_mask = self._bucket_inputs(x, text, c_mask, ref, ref_mask) - key = self._key(x, text, ref, uses_cfg) + steps = 0 if timesteps is None else int(timesteps.numel()) + key = self._key(x, text, ref, uses_cfg, steps) entry = self._cache.get(key) ctx = None if entry is None or new_request or key != self._active_key: - ctx = self._prepare(x, x_mask, text, c_mask, ref, ref_mask, uses_cfg) + ctx = self._prepare(x, x_mask, text, c_mask, ref, ref_mask, uses_cfg, timesteps) if entry is None: # Same Retirement as #6587, details discussed in #7469 self._retire_graph_generation_if_full() try: assert ctx is not None - entry = self._capture(x, ctx, timestep, cfg_strength=cfg_strength) + entry = self._capture(x, ctx, timestep, cfg_strength=cfg_strength, step_index=step_index) except Exception: self._cache.clear() self._active_key = None @@ -186,6 +202,8 @@ class AuKCUDAGraphWrapper: entry.static_x.copy_(x) entry.static_timestep.copy_(timestep) + if entry.static_step is not None: + entry.static_step.fill_(step_index) if entry.static_cfg is not None: entry.static_cfg.fill_(cfg_strength) entry.graph.replay() @@ -198,10 +216,14 @@ class AuKCUDAGraphWrapper: timestep: torch.Tensor, *, cfg_strength: float, + step_index: int | None = None, ) -> _GraphEntry: static_x = x.clone() static_ctx = ctx.clone() static_timestep = timestep.clone() + static_step = None + if static_ctx.modulation is not None and step_index is not None: + static_step = torch.full((), step_index, device=static_x.device, dtype=torch.long) static_cfg = None if ctx.branches == 2: static_cfg = torch.empty((), device=static_x.device, dtype=torch.float32) @@ -209,13 +231,13 @@ class AuKCUDAGraphWrapper: scale = static_cfg if static_cfg is not None else cfg_strength # Warm-up runs any lazy torch.compile outside the capture. for _ in range(3): - self._step(static_x, static_timestep, static_ctx, scale) + self._step(static_x, static_timestep, static_ctx, scale, static_step) if self._pool_handle is None: self._pool_handle = torch.cuda.graph_pool_handle() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, pool=self._pool_handle): # The CFG scalar stays a mutable buffer, so cfg=2 and cfg=3 share this graph. - static_out = self._step(static_x, static_timestep, static_ctx, scale) + static_out = self._step(static_x, static_timestep, static_ctx, scale, static_step) logger.info( "Captured AuK DiT single-step CUDA graph: target_frames=%d text_tokens=%d ref_frames=%d cfg=%s", @@ -230,6 +252,7 @@ class AuKCUDAGraphWrapper: static_ctx=static_ctx, static_timestep=static_timestep, static_cfg=static_cfg, + static_step=static_step, static_out=static_out, ) diff --git a/vllm_omni/diffusion/models/auk/pipeline_auk.py b/vllm_omni/diffusion/models/auk/pipeline_auk.py index 5e74c491a..da5d8d1e5 100644 --- a/vllm_omni/diffusion/models/auk/pipeline_auk.py +++ b/vllm_omni/diffusion/models/auk/pipeline_auk.py @@ -27,7 +27,12 @@ from vllm.logger import init_logger from vllm_omni.diffusion.compile import regionally_compile from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device -from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, dit_state_dict, sample_latents +from vllm_omni.diffusion.models.auk.auk_transformer import ( + AuKTransformer, + build_time_grid, + dit_state_dict, + sample_latents, +) from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE from vllm_omni.diffusion.models.auk.cudagraph_wrapper import AuKCUDAGraphWrapper from vllm_omni.diffusion.models.auk.vae_cudagraph import AuKVAEDecodeGraph @@ -267,6 +272,12 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp text = torch.zeros(1, _WARMUP_TEXT_TOKENS, self.text_hidden_dim, device=self.device, dtype=self.dtype) c_mask = torch.ones(text.shape[:2], dtype=torch.bool, device=self.device) timestep = torch.zeros((), device=self.device, dtype=torch.float32) + # The graphs are keyed by the number of steps, so warm the variant's default grid. + if self.is_flash: + grid = torch.tensor(self.flash_t_grid, device=self.device, dtype=torch.float32) + else: + sway = None if self.default_sway is None else float(self.default_sway) + grid = build_time_grid(nfe=self.default_nfe, sway_sampling_coef=sway, t_grid=None, device=self.device) start = time.perf_counter() with torch.inference_mode(), self._dit_autocast(): for ref_frames in _WARMUP_REF_FRAMES: @@ -283,6 +294,8 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp timestep=timestep, cfg_strength=cfg, new_request=True, + timesteps=grid[:-1], + step_index=0, ) if self.device.type != "cpu": torch.accelerator.synchronize(self.device) From 4025bb30f670c83649a4dce6779227d01f2b58ec Mon Sep 17 00:00:00 2001 From: Yueqian Lin Date: Sat, 3 Oct 2026 01:43:55 +1000 Subject: [PATCH 18/22] [Perf][AuK] Graph the encoder prefill, cache reference latents and fuse the codec activation (#8305) Signed-off-by: Yueqian Lin --- .../models/auk/test_auk_vae_fastpath.py | 69 ++++++- .../diffusion/models/auk/test_pipeline_auk.py | 46 +++++ .../test_auk_encoder_cudagraph.py | 175 ++++++++++++++++++ vllm_omni/deploy/auk.yaml | 14 +- .../diffusion/models/auk/alias_free_triton.py | 148 +++++++++++++++ vllm_omni/diffusion/models/auk/auk_vae.py | 39 ++++ .../diffusion/models/auk/pipeline_auk.py | 31 +++- 7 files changed, 515 insertions(+), 7 deletions(-) create mode 100644 tests/e2e/offline_inference/test_auk_encoder_cudagraph.py create mode 100644 vllm_omni/diffusion/models/auk/alias_free_triton.py diff --git a/tests/diffusion/models/auk/test_auk_vae_fastpath.py b/tests/diffusion/models/auk/test_auk_vae_fastpath.py index 8a1b39aa1..bec028bc7 100644 --- a/tests/diffusion/models/auk/test_auk_vae_fastpath.py +++ b/tests/diffusion/models/auk/test_auk_vae_fastpath.py @@ -12,7 +12,8 @@ import pytest import torch from tests.helpers.mark import hardware_test -from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE, LowPass, SnakeBeta, Upsample +from vllm_omni.diffusion.models.auk import auk_vae +from vllm_omni.diffusion.models.auk.auk_vae import AliasFreeActivation, AuKVAE, LowPass, SnakeBeta, Upsample from vllm_omni.diffusion.models.auk.vae_cudagraph import AuKVAEDecodeGraph, plan_tiles from vllm_omni.platforms import current_omni_platform @@ -246,6 +247,72 @@ def test_tiled_decode_matches_the_whole_decode() -> None: AuKVAEDecodeGraph(vae, tile_frames=60) +@pytest.mark.cpu +@torch.inference_mode() +def test_fused_activation_is_skipped_off_cuda() -> None: + activation = AliasFreeActivation(SnakeBeta(4, alpha_logscale=True)) + x = torch.randn(1, 4, 16) + assert not activation._fused_eligible(x) + assert activation(x).shape == (1, 4, 16) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="The fused activation kernels require CUDA") +@torch.inference_mode() +def test_fused_activation_is_skipped_without_triton(monkeypatch: pytest.MonkeyPatch) -> None: + activation = AliasFreeActivation(SnakeBeta(4, alpha_logscale=True)).to("cuda") + x = torch.randn(1, 4, 16, device="cuda") + assert activation._fused_eligible(x) + monkeypatch.setattr(auk_vae, "HAS_TRITON", False) + assert not activation._fused_eligible(x) + assert activation(x).shape == (1, 4, 16) + + +def _activation_input(layout: str, length: int) -> torch.Tensor: + """A ``[2, 6, length]`` input, contiguous or as a strided view of a larger buffer.""" + if layout == "contiguous": + x = torch.randn(2, 6, length, device="cuda") + elif layout == "transposed": + x = torch.randn(2, length, 6, device="cuda").transpose(1, 2) + elif layout == "sliced": + x = torch.randn(2, 8, 2 * length, device="cuda")[:, 1:7, ::2] + else: + raise ValueError(layout) + assert x.shape == (2, 6, length) + if length > 1: # contiguity of a single-sample view depends on the layout + assert x.is_contiguous() == (layout == "contiguous") + return x * 3 + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="The fused activation kernels require CUDA") +@pytest.mark.parametrize("layout", ["contiguous", "transposed", "sliced"]) +@pytest.mark.parametrize("causal", [False, True]) +# The kernels tile the output in blocks of 1024 samples: the 2x oversampled +# signal crosses a block boundary at 512 input samples and the decimated one at +# 1024, so both sides of each boundary are covered, plus short and odd lengths. +@pytest.mark.parametrize("length", [1, 7, 511, 512, 513, 1000, 1023, 1024, 1025]) +@torch.inference_mode() +def test_fused_activation_matches_the_eager_modules(layout: str, causal: bool, length: int) -> None: + torch.manual_seed(0) + activation = AliasFreeActivation(SnakeBeta(6, alpha_logscale=True), causal=causal).to("cuda") + with torch.no_grad(): + activation.act.alpha.normal_(0.0, 0.3) + activation.act.beta.normal_(0.0, 0.3) + x = _activation_input(layout, length) + original = x.clone() + + eager = activation.downsample(activation.act(activation.upsample(x))) + assert activation._fused_eligible(x) + fused = activation(x) + + # Same fp32 arithmetic, only the order in which the taps are summed differs. + assert fused.shape == eager.shape + torch.testing.assert_close(fused, eager, atol=1e-5, rtol=1e-5) + # Neither path writes into its (possibly strided) input. + torch.testing.assert_close(x, original, atol=0, rtol=0) + + @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph replay requires CUDA") @torch.inference_mode() diff --git a/tests/diffusion/models/auk/test_pipeline_auk.py b/tests/diffusion/models/auk/test_pipeline_auk.py index 67d38e09a..01cc035f3 100644 --- a/tests/diffusion/models/auk/test_pipeline_auk.py +++ b/tests/diffusion/models/auk/test_pipeline_auk.py @@ -291,6 +291,18 @@ class TestRequestParsing: with pytest.raises(ValueError, match="AuK max_dit_graphs must be a positive integer"): build_pipeline(model_config={"max_dit_graphs": max_dit_graphs}) + @pytest.mark.parametrize( + "model_config, expected", [({}, 16), ({"auk_ref_cache_size": 0}, 0), ({"auk_ref_cache_size": 4}, 4)] + ) + def test_ref_cache_size_is_taken_from_model_config(self, build_pipeline, model_config, expected): + pipeline, _ = build_pipeline(model_config=model_config) + assert pipeline._ref_cache_size == expected + + @pytest.mark.parametrize("size", [-1, 1.9, True, False, "4", None]) + def test_invalid_ref_cache_size_is_rejected(self, build_pipeline, size): + with pytest.raises(ValueError, match="AuK auk_ref_cache_size must be a non-negative integer"): + build_pipeline(model_config={"auk_ref_cache_size": size}) + def test_pipeline_declares_audio_output(self, build_pipeline): pipeline, _ = build_pipeline() @@ -492,6 +504,40 @@ class TestRequestParsing: assert call["sample"] is True assert isinstance(call["generator"], torch.Generator) + def test_reference_latents_are_cached_by_content(self, build_pipeline): + pipeline, calls = build_pipeline() + clip = _silence(1.0) + + for _ in range(3): + pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1)) + # A different clip is a new entry, the first one still hits. + other = (np.full_like(clip[0], 0.25), clip[1]) + pipeline.forward(_batch(_prompt(audio=other, knobs={"gen_seconds": 1.0}), seed=1)) + pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1)) + + assert len(pipeline.vae.encode_calls) == 2 + assert calls[0]["ref"] is calls[1]["ref"] is calls[4]["ref"] + + def test_posterior_draws_and_disabled_cache_always_encode(self, build_pipeline): + pipeline, _ = build_pipeline() + for _ in range(2): + pipeline.forward( + _batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "vae_sample": True}), seed=1) + ) + assert len(pipeline.vae.encode_calls) == 2 + + pipeline, _ = build_pipeline(model_config={"auk_ref_cache_size": 0}) + for _ in range(2): + pipeline.forward(_batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0}), seed=1)) + assert len(pipeline.vae.encode_calls) == 2 + + def test_reference_cache_is_bounded(self, build_pipeline): + pipeline, _ = build_pipeline(model_config={"auk_ref_cache_size": 2}) + for level in (0.1, 0.2, 0.3): + clip = (np.full(SAMPLE_RATE, level, dtype=np.float32), SAMPLE_RATE) + pipeline.forward(_batch(_prompt(audio=clip, knobs={"gen_seconds": 1.0}), seed=1)) + assert len(pipeline._ref_cache) == 2 + def test_one_request_per_forward(self, build_pipeline): pipeline, _ = build_pipeline() batch = _batch(_prompt(knobs={"gen_seconds": 1.0}), seed=1) diff --git a/tests/e2e/offline_inference/test_auk_encoder_cudagraph.py b/tests/e2e/offline_inference/test_auk_encoder_cudagraph.py new file mode 100644 index 000000000..72c1a73dd --- /dev/null +++ b/tests/e2e/offline_inference/test_auk_encoder_cudagraph.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""AuK encoder stage: FULL CUDA graph replay matches eager, request by request. + +``auk.yaml`` captures the encoder prefill as FULL CUDA graphs padded to a few +token buckets. A replay reuses static input buffers and the graph pool, so a +request could in principle read conditioning left behind by an earlier, longer +prompt. This test runs the same request sequence through an eager encoder and +through the graphed one and compares the text condition stage 0 hands to +stage 1 (the fused thinker hidden states) for every request. The sequence +alternates prompts of different lengths inside one capture bucket, crosses +into a larger bucket, and then returns to the first prompt. + +Like ``test_auk.py`` it needs an assembled checkpoint directory: + + VLLM_OMNI_AUK_MODEL_DIR=/path/to/auk-omni python -m pytest tests/e2e/offline_inference/test_auk_encoder_cudagraph.py +""" + +from __future__ import annotations + +import bisect +import os +from pathlib import Path +from typing import Any + +import pytest +import torch +import yaml + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.model_executor.stage_input_processors import auk as auk_input_processor +from vllm_omni.model_extras.auk import auk_prompt, auk_sampling_params + +MODEL_DIR_ENV = "VLLM_OMNI_AUK_MODEL_DIR" + +# Text-only instruct prompts: the token count, and so the capture bucket, +# follows the text length. SHORT_A and SHORT_B differ in length but share a +# bucket; LONG lands in a larger one. The test checks these properties on the +# recorded condition lengths, so a tokenizer change fails loudly instead of +# silently shrinking the coverage. +_TEMPLATE = 'Generate speech based on the following description: "{desc}". The content to speak is: "{text}".' +SHORT_A = _TEMPLATE.format( + desc="A calm young woman speaking warmly.", + text="Welcome back, how was your day? I kept some soup warm for you, and the kettle is on if you want tea.", +) +SHORT_B = _TEMPLATE.format( + desc="An older man with a deep, slightly raspy voice, speaking slowly.", + text=( + "The train to the coast leaves at seven, so please be on time. Bring a warm coat, " + "some cash for the ferry, and the blue umbrella from the hall, because the forecast says rain." + ), +) +LONG = _TEMPLATE.format( + desc=( + "A cheerful radio host in her thirties, bright and energetic, with clear articulation, " + "a slight smile in her voice, and a brisk but relaxed pace, recorded close to the microphone" + ), + text=( + "Good morning and welcome to the weekend edition. Today we are talking about small gardens, " + "the plants that thrive on a sunny balcony, and the neighbours who turned an empty lot into " + "a place where the whole street now meets on Saturday afternoons to trade seedlings and stories. " + "Stay with us, because after the news we will open the phone lines and hear from you." + ), +) +SEQUENCE = [SHORT_A, SHORT_B, SHORT_A, LONG, SHORT_B, SHORT_A] + +# The DiT output is not inspected; keep stage 1 as cheap as it goes (it also +# runs eager in both configs, see _deploy). +GEN_SECONDS = 1.0 +NFE = 1 + +# bf16 hidden states. Graph replay pads the prefill to the bucket size but runs +# the same kernels on the real tokens, so the two paths agree to bf16 rounding. +# Measured on H20: cosine >= 0.99998 and max-abs error <= 2.9e-3 of the +# condition's max magnitude for every request; the bounds leave about 3x +# headroom for other GPUs. +MIN_COSINE = 0.9999 +MAX_REL_ABS = 1e-2 + +_model_dir = os.environ.get(MODEL_DIR_ENV) +pytestmark = [ + pytest.mark.slow, + pytest.mark.tts, + pytest.mark.skipif(not _model_dir, reason=f"set {MODEL_DIR_ENV} to an assembled AuK directory"), +] + + +def _deploy(tmp_path: Path, *, graph_encoder: bool) -> str: + """``auk.yaml`` with stage 1 eager, and the encoder stage eager or graphed as shipped. + + Only the encoder differs between the two runs. The conditions compared here + are recorded before stage 1 runs, so its DiT compile, graph capture and + codec bucket warm-up would only add startup time; eager stage 1 skips them + and still takes the real stage handoff. + """ + with open(get_deploy_config_path("auk.yaml")) as f: + deploy = yaml.safe_load(f) + for stage in deploy["stages"]: + if stage["stage_id"] == 1: + stage["enforce_eager"] = True + elif not graph_encoder: + stage["enforce_eager"] = True + stage.pop("compilation_config", None) + path = tmp_path / f"auk_{'graph' if graph_encoder else 'eager'}_encoder.yaml" + path.write_text(yaml.safe_dump(deploy)) + return str(path) + + +def _graph_capture_sizes() -> list[int]: + with open(get_deploy_config_path("auk.yaml")) as f: + deploy = yaml.safe_load(f) + stage = next(s for s in deploy["stages"] if s["stage_id"] == 0) + assert stage.get("enforce_eager") is False, "auk.yaml no longer graphs the encoder; update this test" + return sorted(stage["compilation_config"]["cudagraph_capture_sizes"]) + + +def _run_sequence(deploy_config: str, monkeypatch: pytest.MonkeyPatch) -> list[torch.Tensor]: + """Run SEQUENCE one request at a time and return each request's text condition.""" + conditions: list[torch.Tensor] = [] + original = auk_input_processor.encoder2dit + + def recording_encoder2dit(*args: Any, **kwargs: Any) -> dict[str, Any] | None: + result = original(*args, **kwargs) + if result is not None: + conditions.append(result["prompt_embeds"].detach().to("cpu", torch.float32).clone()) + return result + + # The orchestrator runs in this process and resolves the stage input + # function by module attribute at startup, so the patch must precede Omni. + with monkeypatch.context() as patch: + patch.setattr(auk_input_processor, "encoder2dit", recording_encoder2dit) + with OmniRunner(str(Path(_model_dir).resolve()), deploy_config=deploy_config) as runner: + for text in SEQUENCE: + prompt = auk_prompt(text, None, gen_seconds=GEN_SECONDS) + outputs = runner.omni.generate(prompt, auk_sampling_params(nfe=NFE, cfg=2.0, seed=0)) + assert outputs, "no outputs returned" + assert len(conditions) == len(SEQUENCE), f"recorded {len(conditions)} conditions for {len(SEQUENCE)} requests" + return conditions + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_encoder_graph_replay_matches_eager(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + capture_sizes = _graph_capture_sizes() + eager = _run_sequence(_deploy(tmp_path, graph_encoder=False), monkeypatch) + graphed = _run_sequence(_deploy(tmp_path, graph_encoder=True), monkeypatch) + + lengths = {text: cond.shape[0] for text, cond in zip(SEQUENCE, eager)} + for i, (ref, got) in enumerate(zip(eager, graphed)): + assert got.shape == ref.shape, f"request {i}: shape {tuple(got.shape)} != eager {tuple(ref.shape)}" + assert torch.isfinite(got).all(), f"request {i}: non-finite condition" + cosine = torch.nn.functional.cosine_similarity(got.flatten(), ref.flatten(), dim=0).item() + rel_abs = ((got - ref).abs().max() / ref.abs().max()).item() + assert cosine >= MIN_COSINE and rel_abs <= MAX_REL_ABS, ( + f"request {i} ({lengths[SEQUENCE[i]]} tokens): cosine {cosine:.6f}, max-abs error {rel_abs:.2e} " + "of the condition's magnitude" + ) + + # The sequence must exercise what it claims to: two different lengths in + # one bucket, a second bucket, and every prompt inside the graphed range. + assert max(lengths.values()) <= capture_sizes[-1], f"{lengths} exceeds the largest capture size" + bucket = {text: capture_sizes[bisect.bisect_left(capture_sizes, n)] for text, n in lengths.items()} + assert lengths[SHORT_A] != lengths[SHORT_B] and bucket[SHORT_A] == bucket[SHORT_B], (lengths, bucket) + assert bucket[LONG] > bucket[SHORT_A], (lengths, bucket) + + # Repeats of a prompt reproduce its first condition exactly, eager or graphed: + # nothing from the requests in between leaks into the replay. + for run in (eager, graphed): + first: dict[str, torch.Tensor] = {} + for text, cond in zip(SEQUENCE, run): + if text in first: + torch.testing.assert_close(cond, first[text], atol=0, rtol=0) + else: + first[text] = cond diff --git a/vllm_omni/deploy/auk.yaml b/vllm_omni/deploy/auk.yaml index 989e64965..555342027 100644 --- a/vllm_omni/deploy/auk.yaml +++ b/vllm_omni/deploy/auk.yaml @@ -20,14 +20,20 @@ enable_chunked_prefill: false stages: # Stage 0: Qwen2.5-Omni thinker as a frozen text encoder. - # Prefill only — one sampled token, discarded downstream. eager because the - # fused-hidden forward walks the decoder layers itself and so bypasses the - # compiled graph. + # Prefill only: one sampled token, discarded downstream. The fused-hidden + # forward walks the decoder layers itself, so torch.compile has nothing to + # hook (mode 0); the whole prefill is instead captured as FULL CUDA graphs + # padded to the sizes below, which removes most of its launch overhead + # (about 20 ms per request on H200). Longer prompts run eager. - stage_id: 0 devices: "0" max_num_seqs: 1 gpu_memory_utilization: 0.45 - enforce_eager: true + enforce_eager: false + compilation_config: + mode: 0 + cudagraph_mode: FULL + cudagraph_capture_sizes: [64, 128, 192, 256, 384, 512] max_num_batched_tokens: 32768 # AuK takes text plus at most one clip; no image or video tower is built. limit_mm_per_prompt: {audio: 1} diff --git a/vllm_omni/diffusion/models/auk/alias_free_triton.py b/vllm_omni/diffusion/models/auk/alias_free_triton.py new file mode 100644 index 000000000..2d37e859e --- /dev/null +++ b/vllm_omni/diffusion/models/auk/alias_free_triton.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Triton kernels for the AuK codec's alias-free SnakeBeta activation. + +The eager module oversamples with a grouped transposed convolution, applies +SnakeBeta, and decimates with a grouped strided convolution. Both grouped +convolutions carry one 12-tap filter shared by every channel, and cuDNN and +ATen serve them with generic depthwise kernels that dominate the decode (more +than half of a 12 s decode on H200). Here the oversampling is written as its +polyphase form fused with the activation, and the decimation as a direct +strided FIR, each one pass over memory per row. The arithmetic is the same +as the eager module in fp32; only the summation order of the taps differs. + +The ops are registered with ``torch.library.triton_op`` so the compiled +decode buckets trace them instead of breaking the graph. +""" + +from __future__ import annotations + +import torch +from vllm.triton_utils import HAS_TRITON, tl, triton + +__all__ = ["alias_free_snake"] + +_BLOCK = 1024 + +if HAS_TRITON: + + @triton.jit + def _upsample_snake_kernel( + x_ptr, + out_ptr, + taps_ptr, + exp_alpha_ptr, + inv_beta_ptr, + channels, + in_len, + out_len, + ratio: tl.constexpr, + taps: tl.constexpr, + pad: tl.constexpr, + crop: tl.constexpr, + block: tl.constexpr, + ): + # out[m] = snake(ratio * sum_k taps[k] * xpad[(m + crop - k) / ratio]) over the k + # with (m + crop - k) divisible by ratio, where xpad is x replicate-padded by pad + # on both sides. That is conv_transpose1d(stride=ratio) followed by the crop. + row = tl.program_id(0) + offs = tl.program_id(1) * block + tl.arange(0, block) + mask = offs < out_len + x_row = x_ptr + row.to(tl.int64) * in_len + acc = tl.zeros((block,), dtype=tl.float32) + for k in tl.static_range(taps): + pos = offs + (crop - k) + hit = (pos % ratio) == 0 + src = pos // ratio - pad + src = tl.minimum(tl.maximum(src, 0), in_len - 1) + w = tl.load(taps_ptr + k) + val = tl.load(x_row + src, mask=mask & hit, other=0.0) + acc += tl.where(hit, w * val, 0.0) + u = acc * ratio + channel = row % channels + alpha = tl.load(exp_alpha_ptr + channel) + inv_beta = tl.load(inv_beta_ptr + channel) + s = tl.sin(u * alpha) + u = u + inv_beta * (s * s) + tl.store(out_ptr + row.to(tl.int64) * out_len + offs, u, mask=mask) + + @triton.jit + def _lowpass_kernel( + x_ptr, + out_ptr, + taps_ptr, + in_len, + out_len, + stride: tl.constexpr, + taps: tl.constexpr, + pad_left: tl.constexpr, + block: tl.constexpr, + ): + # out[n] = sum_k taps[k] * x[clamp(stride * n + k - pad_left)], i.e. replicate + # padding followed by a strided conv1d. + row = tl.program_id(0) + offs = tl.program_id(1) * block + tl.arange(0, block) + mask = offs < out_len + x_row = x_ptr + row.to(tl.int64) * in_len + acc = tl.zeros((block,), dtype=tl.float32) + for k in tl.static_range(taps): + src = offs * stride + (k - pad_left) + src = tl.minimum(tl.maximum(src, 0), in_len - 1) + acc += tl.load(taps_ptr + k) * tl.load(x_row + src, mask=mask, other=0.0) + tl.store(out_ptr + row.to(tl.int64) * out_len + offs, acc, mask=mask) + + @torch.library.triton_op("vllm_omni::auk_alias_free_snake", mutates_args=()) + def alias_free_snake( + x: torch.Tensor, + up_taps: torch.Tensor, + down_taps: torch.Tensor, + exp_alpha: torch.Tensor, + inv_beta: torch.Tensor, + up_ratio: int, + up_pad: int, + up_crop: int, + down_stride: int, + down_pad_left: int, + down_pad_right: int, + ) -> torch.Tensor: + """Upsample by ``up_ratio``, SnakeBeta, low-pass and decimate: ``[B, C, T] -> [B, C, T_out]``.""" + batch, channels, length = x.shape + x = x.contiguous() + up_len = length * up_ratio + oversampled = torch.empty((batch, channels, up_len), device=x.device, dtype=torch.float32) + rows = batch * channels + torch.library.wrap_triton(_upsample_snake_kernel)[(rows, triton.cdiv(up_len, _BLOCK))]( + x, + oversampled, + up_taps, + exp_alpha, + inv_beta, + channels, + length, + up_len, + ratio=up_ratio, + taps=up_taps.numel(), + pad=up_pad, + crop=up_crop, + block=_BLOCK, + ) + taps = down_taps.numel() + out_len = (up_len + down_pad_left + down_pad_right - taps) // down_stride + 1 + out = torch.empty((batch, channels, out_len), device=x.device, dtype=torch.float32) + torch.library.wrap_triton(_lowpass_kernel)[(rows, triton.cdiv(out_len, _BLOCK))]( + oversampled, + out, + down_taps, + up_len, + out_len, + stride=down_stride, + taps=taps, + pad_left=down_pad_left, + block=_BLOCK, + ) + return out + +else: # pragma: no cover + + def alias_free_snake(*args, **kwargs): + raise RuntimeError("Triton is not available") diff --git a/vllm_omni/diffusion/models/auk/auk_vae.py b/vllm_omni/diffusion/models/auk/auk_vae.py index 2322cfbc9..4ad20dd87 100644 --- a/vllm_omni/diffusion/models/auk/auk_vae.py +++ b/vllm_omni/diffusion/models/auk/auk_vae.py @@ -32,7 +32,9 @@ from torch import nn from torch.nn.utils import remove_weight_norm as _fold_weight_norm from torch.nn.utils import weight_norm as _apply_weight_norm from vllm.logger import init_logger +from vllm.triton_utils import HAS_TRITON +from vllm_omni.diffusion.models.auk.alias_free_triton import alias_free_snake from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta as _SharedSnakeBeta logger = init_logger(__name__) @@ -319,8 +321,45 @@ class AliasFreeActivation(nn.Module): self.downsample = Downsample(down_ratio, down_kernel_size, causal=causal) def forward(self, x: torch.Tensor) -> torch.Tensor: + if self._fused_eligible(x): + return self._fused_forward(x) return self.downsample(self.act(self.upsample(x))) + def _fused_eligible(self, x: torch.Tensor) -> bool: + """Whether ``x`` can take the Triton kernels in alias_free_triton (CUDA fp32, AuK's filter geometry).""" + up = self.upsample + taps = up.filter.numel() + return ( + x.is_cuda + and x.dtype == torch.float32 + and x.dim() == 3 + and isinstance(self.act, SnakeBeta) + and HAS_TRITON + # The kernel indexes the crop from the start of the transposed + # conv output, so the crop must cover the filter's reach, and the + # cropped length must be exactly ratio * T. + and up.pad_left >= taps - 1 + and (2 * up.pad - 1) * up.ratio + taps == up.pad_left + up.pad_right + ) + + def _fused_forward(self, x: torch.Tensor) -> torch.Tensor: + up, low, act = self.upsample, self.downsample.lowpass, self.act + if not act._cached: + act.precompute_exp_cache() + return alias_free_snake( + x, + up.filter.reshape(-1), + low.filter.reshape(-1), + act._exp_alpha, + act._inv_beta, + up.ratio, + up.pad, + up.pad_left, + low.stride, + low.pad_left, + low.pad_right, + ) + class AmpBlock(nn.Module): """Anti-aliased multi-periodicity residual block: alternating alias-free snake and dilated convs.""" diff --git a/vllm_omni/diffusion/models/auk/pipeline_auk.py b/vllm_omni/diffusion/models/auk/pipeline_auk.py index da5d8d1e5..21b639411 100644 --- a/vllm_omni/diffusion/models/auk/pipeline_auk.py +++ b/vllm_omni/diffusion/models/auk/pipeline_auk.py @@ -11,9 +11,11 @@ both, and the BigVGAN-flow decoder renders the waveform. from __future__ import annotations +import hashlib import json import os import time +from collections import OrderedDict from collections.abc import Iterable from typing import Any, ClassVar @@ -64,6 +66,10 @@ _FLASH_CFG = 0.0 _WARMUP_TARGET_FRAMES = (150, 300, 600) _WARMUP_REF_FRAMES = (0, 150) _WARMUP_TEXT_TOKENS = 96 +# Reference clips whose posterior-mean latents are kept. Voice cloning reuses a +# small set of reference voices, so repeat requests skip the VAE encode. +# ``model_config.auk_ref_cache_size`` overrides it; 0 disables the cache. +_REF_CACHE_SIZE = 16 # Decorrelates the VAE posterior draw from the initial latent for the same request seed. _VAE_SEED_OFFSET = 0x5EED_0A0C @@ -157,6 +163,9 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp max_dit_graphs = od_config.model_config.get("max_dit_graphs", 32) if isinstance(max_dit_graphs, bool) or not isinstance(max_dit_graphs, int) or max_dit_graphs < 1: raise ValueError("AuK max_dit_graphs must be a positive integer") + ref_cache_size = od_config.model_config.get("auk_ref_cache_size", _REF_CACHE_SIZE) + if isinstance(ref_cache_size, bool) or not isinstance(ref_cache_size, int) or ref_cache_size < 0: + raise ValueError("AuK auk_ref_cache_size must be a non-negative integer (0 disables the cache)") self.device = get_local_device() self.dtype = getattr(od_config, "dtype", None) or torch.bfloat16 @@ -218,6 +227,8 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp if model_config.get("auk_vae_tile_frames") is not None: vae_decode_kwargs["tile_frames"] = int(model_config["auk_vae_tile_frames"]) self.vae_decode = AuKVAEDecodeGraph(self.vae, enabled=not od_config.enforce_eager, **vae_decode_kwargs) + self._ref_cache_size = ref_cache_size + self._ref_cache: OrderedDict[str, torch.Tensor] = OrderedDict() logger.info( "AuK pipeline ready: variant=%s dtype=%s latent_dim=%d hop=%d sample_rate=%d", @@ -399,8 +410,24 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp if audio is None: return torch.zeros(1, 0, self.latent_dim, device=self.device, dtype=torch.float32) - wav = self._prepare_waveform(_unwrap_single(audio)).to(self.device) - return self.vae.encode(wav, sample=sample, generator=generator) + wav = self._prepare_waveform(_unwrap_single(audio)) + # The posterior mean is a pure function of the prepared waveform, so it + # is cached by content. A posterior draw depends on the generator and + # is always recomputed. Callers must not modify the returned tensor. + key = None + if not sample and self._ref_cache_size: + digest = hashlib.blake2b(wav.numpy().tobytes(), digest_size=16).hexdigest() + key = f"{wav.shape[-1]}:{digest}" + cached = self._ref_cache.get(key) + if cached is not None: + self._ref_cache.move_to_end(key) + return cached + latents = self.vae.encode(wav.to(self.device), sample=sample, generator=generator) + if key is not None: + self._ref_cache[key] = latents + while len(self._ref_cache) > self._ref_cache_size: + self._ref_cache.popitem(last=False) + return latents def _target_frames(self, gen_seconds: Any, ref_frames: int) -> int: """Target latent length, capped by the remaining context.""" From 2ba5074ca039a94698e7a021243d98577430d84f Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Fri, 2 Oct 2026 08:49:50 -0700 Subject: [PATCH 19/22] [Model] Optimize CosyVoice3 packed streaming with bounded HiFT graphs (#8420) Signed-off-by: Sy03 <1370724210@qq.com> --- tests/config/test_environment_variables.py | 4 +- .../cosyvoice3/test_cosyvoice3_hift_graph.py | 41 +++++++++ .../config/environment_variable_inventory.py | 2 + ...3_packed_streaming_optimized_standard.yaml | 24 ++++++ .../cosyvoice3/code2wav_core/hifigan.py | 13 +++ .../cosyvoice3/code2wav_core/hift_graph.py | 85 +++++++++++++++++++ .../cosyvoice3/code2wav_core/packed_dit.py | 6 +- .../models/cosyvoice3/cosyvoice3_code2wav.py | 7 ++ 8 files changed, 178 insertions(+), 4 deletions(-) create mode 100644 tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_graph.py create mode 100644 vllm_omni/deploy/cosyvoice3_packed_streaming_optimized_standard.yaml create mode 100644 vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hift_graph.py diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index e1f6aa651..eb3db91ec 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: 32, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29, - EnvironmentVariableCategory.MODEL_SPECIFIC: 75, + EnvironmentVariableCategory.MODEL_SPECIFIC: 77, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21, EnvironmentVariableCategory.INTERNAL: 7, } @@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC ) assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == { - ModelEnvironmentVariableDisposition.PROMOTE: 46, + ModelEnvironmentVariableDisposition.PROMOTE: 48, ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5, ModelEnvironmentVariableDisposition.EXTERNAL: 0, ModelEnvironmentVariableDisposition.INTERNALIZE: 17, diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_graph.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_graph.py new file mode 100644 index 000000000..f96947cf9 --- /dev/null +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_graph.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""HiFT graphs consume live inputs and return independently owned waveforms.""" + +import pytest +import torch + +from tests.helpers.mark import hardware_test +from tests.model_executor.models.cosyvoice3.test_cosyvoice3_incremental_hift import _make_hift + +pytestmark = [pytest.mark.core_model] + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("finalize", [False, True]) +@torch.inference_mode() +def test_hift_graph_preserves_current_and_previous_request(finalize): + hift = _make_hift().cuda().eval() + hift.remove_weight_norm() + hift.f0_predictor.remove_weight_norm() + hift.enable_decode_graphs() + generator = torch.Generator(device="cuda").manual_seed(17) + saved: list[tuple[torch.Tensor, torch.Tensor]] = [] + for _ in range(5): + mel = torch.randn(1, 80, 30, device="cuda", generator=generator) + source = torch.randn(1, 1, 30 * 480, device="cuda", generator=generator) + expected = hift._decode_eager(mel, source, finalize=finalize) + actual = hift.decode(mel, source, finalize=finalize) + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6) + for previous, snapshot in saved: + assert torch.equal(previous, snapshot) + if hift._decode_graphs.graphs: + saved.append((actual, actual.clone())) + assert len(hift._decode_graphs.graphs) == 1 + hift._decode_graphs.max_graphs = 1 + mel = torch.randn(1, 80, 31, device="cuda", generator=generator) + source = torch.randn(1, 1, 31 * 480, device="cuda", generator=generator) + actual = hift.decode(mel, source, finalize=finalize) + expected = hift._decode_eager(mel, source, finalize=finalize) + assert torch.equal(actual, expected) + assert len(hift._decode_graphs.graphs) == 1 diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index ee7f599f2..06dcb3ef4 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -148,11 +148,13 @@ _MODEL_PROMOTE = ( "COSMOS3_SOUND_TOKENIZER_CONFIG_PATH", "COSMOS3_SOUND_TOKENIZER_PATH", "COSYVOICE3_BATCH_FLOW", + "COSYVOICE3_CACHED_ISTFT", "COSYVOICE3_ESTIMATOR_ONNX", "COSYVOICE3_FLOW_CUDAGRAPH", "COSYVOICE3_FLOW_GRAPH", "COSYVOICE3_FLOW_GRAPH_MAX_TOTAL", "COSYVOICE3_FULL_RESPONSE_OPTIMIZATIONS", + "COSYVOICE3_HIFT_GRAPH", "COSYVOICE3_PACKED_STREAMING", "COSYVOICE3_REFERENCE_PREFETCH", "COSYVOICE3_REFERENCE_WORKERS", diff --git a/vllm_omni/deploy/cosyvoice3_packed_streaming_optimized_standard.yaml b/vllm_omni/deploy/cosyvoice3_packed_streaming_optimized_standard.yaml new file mode 100644 index 000000000..0d738a0b1 --- /dev/null +++ b/vllm_omni/deploy/cosyvoice3_packed_streaming_optimized_standard.yaml @@ -0,0 +1,24 @@ +# Opt-in single-Hopper streaming throughput profile: AR32 / codec32, Flow10. +# Use --hf-overrides '{"cosyvoice3_sampling_mode": "standard"}' and stream=true. +# F0 and HiFT stay FP32; only deterministic HiFT decode is graphed. +# HiFT graphs capture on the third use of each exact shape, adding latency then. +# The codec graph pool is not capped by gpu_memory_utilization; leave memory headroom. +base_config: cosyvoice3_packed_streaming.yaml +stages: + - stage_id: 0 + max_num_seqs: 32 + default_sampling_params: + max_tokens: 2048 + repetition_penalty: 1.21 + temperature: 0.7 + top_k: 20 + top_p: 0.8 + - stage_id: 1 + max_num_seqs: 32 + env: + COSYVOICE3_TRT: "0" + COSYVOICE3_BATCH_FLOW: "1" + COSYVOICE3_FULL_RESPONSE_OPTIMIZATIONS: "0" + COSYVOICE3_PACKED_STREAMING: "1" + COSYVOICE3_CACHED_ISTFT: "1" + COSYVOICE3_HIFT_GRAPH: "1" 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 a353675bf..a1e31c070 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py @@ -966,7 +966,20 @@ class CausalHiFTGenerator(HiFTGenerator): self.conv_pre_look_right = conv_pre_look_right self.f0_predictor = f0_predictor + def enable_decode_graphs(self) -> None: + """Opt into exact-shape CUDA replay after frozen weights are loaded.""" + from .hift_graph import HiFTDecodeGraphs + + self.enable_cached_istft() + self._decode_graphs = HiFTDecodeGraphs(self._decode_eager) + def decode(self, x: torch.Tensor, s: torch.Tensor = torch.zeros(1, 1, 0), finalize: bool = True) -> torch.Tensor: + runner = getattr(self, "_decode_graphs", None) + if runner is not None: + return runner.run(x, s, finalize) + return self._decode_eager(x, s, finalize=finalize) + + def _decode_eager(self, x: torch.Tensor, s: torch.Tensor, finalize: bool = True) -> torch.Tensor: s_stft_real, s_stft_imag = self._stft(s.squeeze(1)) if finalize is True: x = self.conv_pre(x) diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hift_graph.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hift_graph.py new file mode 100644 index 000000000..588497310 --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hift_graph.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Bounded CUDA graph replay for HiFT's deterministic waveform decoder. + +F0, sine/noise generation, phase carry, and per-request streaming state remain +outside the graph. Each replay consumes current mel and source tensors and +returns an owned waveform, because multiple requests can reuse one shape. +""" + +from collections import OrderedDict +from collections.abc import Callable +from dataclasses import dataclass + +import torch + + +@dataclass +class _DecodeGraph: + graph: torch.cuda.CUDAGraph + mel: torch.Tensor + source: torch.Tensor + output: torch.Tensor + + +class HiFTDecodeGraphs: + """Capture recurrent exact shapes, with at most 32 resident graphs. + + ``decode`` must have warmed its ISTFT envelope before capture and remain + deterministic for its mel, source and finalize arguments. Uncommon shapes, + CPU execution and nested CUDA capture use that same eager implementation. + """ + + def __init__(self, decode: Callable[..., torch.Tensor], max_graphs: int = 32) -> None: + self.decode = decode + self.max_graphs = max_graphs + self.graphs: dict[tuple, _DecodeGraph] = {} + self.seen: OrderedDict[tuple, int] = OrderedDict() + self.pool = None + + @torch.inference_mode() + def run(self, mel: torch.Tensor, source: torch.Tensor, finalize: bool) -> torch.Tensor: + """Replay a recurring layout without aliasing another request's output.""" + if not mel.is_cuda or not source.is_cuda or torch.cuda.is_current_stream_capturing(): + return self.decode(mel, source, finalize=finalize) + autocast = torch.is_autocast_enabled("cuda") + key = ( + mel.shape, + source.shape, + mel.device, + mel.dtype, + source.dtype, + finalize, + torch.backends.cudnn.allow_tf32, + torch.backends.cuda.matmul.allow_tf32, + autocast, + torch.get_autocast_dtype("cuda") if autocast else None, + ) + state = self.graphs.get(key) + if state is None: + if len(self.graphs) >= self.max_graphs: + return self.decode(mel, source, finalize=finalize) + count = self.seen.pop(key, 0) + 1 + self.seen[key] = count + if len(self.seen) > 128: + self.seen.popitem(last=False) + if count < 3: + return self.decode(mel, source, finalize=finalize) + static_mel, static_source = mel.clone(), source.clone() + stream = torch.cuda.Stream(device=mel.device) + stream.wait_stream(torch.cuda.current_stream(mel.device)) + with torch.cuda.stream(stream): + self.decode(static_mel, static_source, finalize=finalize) + torch.cuda.current_stream(mel.device).wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=self.pool): + output = self.decode(static_mel, static_source, finalize=finalize) + if self.pool is None: + self.pool = graph.pool() + state = _DecodeGraph(graph, static_mel, static_source, output) + self.graphs[key] = state + state.mel.copy_(mel) + state.source.copy_(source) + state.graph.replay() + # A later request may replay this shape before the caller consumes it. + return state.output.clone() diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/packed_dit.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/packed_dit.py index ff6535b87..579f7d32e 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/packed_dit.py +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/packed_dit.py @@ -170,7 +170,7 @@ class PackedDiT: t, rows, attention, - max_seqlen_q=attention.page_table.shape[1], + max_seqlen_q=attention.max_seqlen_q, block_modulation=block_modulation, final_modulation=final_modulation, rope=rope, @@ -771,7 +771,9 @@ class FlowGraphRunner: attention.heads, attention.head_dim = self.heads, self.head_dim attention.cache_seqlens, attention.cu_seqlens_q = state.cache_seqlens, state.cu_seqlens_q attention.page_table = state.page_table - attention.max_seqlen_q = state.page_table.shape[1] + # Streaming queries cover one training chunk, even when their KV + # prefix and the packed graph allocation contain thousands of frames. + attention.max_seqlen_q = estimator.dit.static_chunk_size if key[0] else state.page_table.shape[1] mark_packed_compile_metadata(rows, attention) mu_cfg = torch.cat((state.mu, torch.zeros_like(state.mu)), dim=1) cond_cfg = torch.cat((state.cond, torch.zeros_like(state.cond)), dim=1) diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py index 6fd13c173..96ec400c6 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py @@ -769,4 +769,11 @@ class CosyVoice3Code2Wav(nn.Module): if f0_folded == 0: logger.warning("HiFT F0 predictor had no weight-norm layers to fold; check config drift") self.hift.eval() + if device.type == "cuda" and cosyvoice3_packed_streaming_enabled(): + if os.getenv("COSYVOICE3_CACHED_ISTFT", "0") == "1": + self.hift.enable_cached_istft() + logger.info("CosyVoice3 streaming HiFT: cached ISTFT enabled") + if os.getenv("COSYVOICE3_HIFT_GRAPH", "0") == "1": + self.hift.enable_decode_graphs() + logger.info("CosyVoice3 streaming HiFT: decode CUDA graphs enabled") logger.info(f"Loaded hift weights from {hift_path}") From 527982d886db4da88f205a576cbc0333f9b8eac5 Mon Sep 17 00:00:00 2001 From: Zheng Wengang Date: Fri, 2 Oct 2026 17:14:39 +0100 Subject: [PATCH 20/22] [Bugfix] Skip async-chunk SHM put for stage-0-final requests (#7245) Signed-off-by: ZhengWG Co-authored-by: amy-why-3459 --- .buildkite/cuda/test-merge.yml | 24 +++ .../test_omni_ar_scheduler_kv_transfer.py | 181 +++++++++++++++++- .../test_omni_ar_scheduler_stale_drain.py | 2 + .../sched/test_omni_ar_scheduler_streaming.py | 4 + ...mni_scheduler_input_coordinator_cleanup.py | 5 +- .../test_chunk_transfer_adapter.py | 28 ++- .../omni_connectors/test_shm_connector.py | 22 ++- .../test_qwen3_omni_shm_leak.py | 155 +++++++++++++++ tests/engine/test_cfg_companion_lifecycle.py | 5 +- .../engine/test_stage_pool_collective_rpc.py | 40 ++++ tests/helpers/client.py | 4 + tests/helpers/omni_scheduler.py | 28 +++ vllm_omni/core/sched/omni_ar_scheduler.py | 72 +++++-- vllm_omni/core/sched/omni_scheduler_mixin.py | 3 + .../connectors/shm_connector.py | 21 +- .../chunk_transfer_adapter.py | 48 ++++- vllm_omni/engine/orchestrator.py | 34 ++++ vllm_omni/engine/stage_engine_core_proc.py | 15 ++ vllm_omni/engine/stage_pool.py | 27 +++ 19 files changed, 680 insertions(+), 38 deletions(-) create mode 100644 tests/e2e/online_serving/test_qwen3_omni_shm_leak.py create mode 100644 tests/helpers/omni_scheduler.py diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index 104d6df26..4162e60e8 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.buildkite/cuda/test-merge.yml @@ -124,6 +124,30 @@ steps: - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py -m 'advanced_model and cuda' --run-level 'advanced_model' mirror_hardwares: h100_2 + - label: "Omni · Qwen3-Omni Async-Chunk SHM Leak" + source_file_dependencies: + - tests/e2e/online_serving/test_qwen3_omni_shm_leak.py + - tests/helpers/runtime.py + - tests/helpers/stage_config.py + - vllm_omni/core/sched/omni_ar_scheduler.py + - vllm_omni/distributed/omni_connectors/connectors/shm_connector.py + - vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py + - vllm_omni/engine/orchestrator.py + - vllm_omni/engine/stage_engine_core_proc.py + - vllm_omni/engine/stage_pool.py + - vllm_omni/entrypoints/openai/serving_chat.py + - vllm_omni/model_executor/models/qwen3_omni/ + - vllm_omni/model_executor/models/common/qwen3_code_predictor.py + - vllm_omni/model_executor/models/common/snake_activation.py + - vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py + - vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py + - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py + - vllm_omni/deploy/qwen3_omni_moe.yaml + timeout_in_minutes: 40 + commands: + - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_shm_leak.py -m 'advanced_model and cuda' --run-level 'advanced_model' + mirror_hardwares: h100_2 + - label: "Omni · MiniCPM-o 4.5 Test" source_file_dependencies: omni_minicpmo_4_5_function timeout_in_minutes: 50 diff --git a/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py b/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py index cffe3f289..86102c451 100644 --- a/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py +++ b/tests/core/sched/test_omni_ar_scheduler_kv_transfer.py @@ -1,10 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from types import SimpleNamespace import pytest from vllm import SamplingParams from vllm.v1.engine import EngineCoreRequest +from vllm.v1.request import RequestStatus +from tests.helpers.omni_scheduler import bind_omits_transfer_helpers from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler +from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin from vllm_omni.engine.async_engine_utils import apply_omni_final_stage_metadata from vllm_omni.engine.serialization import deserialize_additional_information @@ -25,7 +31,7 @@ def _make_engine_request() -> EngineCoreRequest: ) -def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]: +def _omit_flags(*, force_kv_transfer: bool) -> tuple[bool, bool, dict]: tagged = apply_omni_final_stage_metadata( _make_engine_request(), final_stage_id=0, @@ -37,21 +43,182 @@ def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]: request_id="req", additional_information=tagged.additional_information, ) - result = scheduler._request_omits_kv_transfer_to_next_stage(request) + omits_kv = scheduler._request_omits_kv_transfer_to_next_stage(request) + omits_chunk = scheduler._request_omits_chunk_transfer_to_next_stage(request) metadata = deserialize_additional_information(tagged.additional_information) - return result, metadata + return omits_kv, omits_chunk, metadata def test_stage_zero_request_omits_kv_transfer(): - omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=False) + omits_kv, omits_chunk, metadata = _omit_flags(force_kv_transfer=False) - assert omits_transfer + assert omits_kv + assert omits_chunk assert "omni_force_kv_transfer" not in metadata def test_cfg_companion_forces_kv_transfer_without_downstream_payload(): - omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=True) + omits_kv, omits_chunk, metadata = _omit_flags(force_kv_transfer=True) - assert not omits_transfer + assert not omits_kv + assert omits_chunk assert metadata["omni_final_stage_id"] == 0 assert metadata["omni_force_kv_transfer"] is True + + +class _ChunkRequest(SimpleNamespace): + def __hash__(self): + return hash(self.request_id) + + def __eq__(self, other): + return isinstance(other, _ChunkRequest) and other.request_id == self.request_id + + +def _make_chunk_request(final_stage_id: int, *, force_kv_transfer: bool = False): + tagged = apply_omni_final_stage_metadata( + _make_engine_request(), + final_stage_id=final_stage_id, + force_kv_transfer=force_kv_transfer, + ) + request = _ChunkRequest( + request_id="req-shm-leak", + additional_information=tagged.additional_information, + num_in_flight_tokens=1, + num_stale_output_tokens=0, + sampling_params=SimpleNamespace(num_logprobs=None), + has_encoder_inputs=False, + pooling_params=None, + status=RequestStatus.RUNNING, + resumable=False, + num_output_placeholders=0, + spec_token_ids=[], + _output_token_ids=[], + num_computed_tokens=1, + stop_reason=None, + num_nans_in_logits=None, + client_index=0, + trace_headers=None, + ) + request.is_finished = lambda: False + request.get_finished_reason = lambda: None + request.take_prefill_stats = lambda: None + request.take_events = lambda: None + return request + + +def _run_finished_save_step(mocker, request, *, inter_stage_outputs=None): + adapter = mocker.MagicMock() + adapter._confirmed_num_computed_tokens.return_value = 1 + + sched = mocker.MagicMock() + sched.requests = {request.request_id: request} + sched.perf_metrics = None + sched.defer_block_free = False + sched.structured_output_manager.should_advance.return_value = False + sched._update_request_with_output.return_value = ([42], True) + sched._process_kv_transfer_trigger.return_value = False + sched._handle_stopped_request.return_value = True + + def _free_request(req, delay_free_blocks=False): + sched._omits_kv_transfer_cache.pop(req.request_id, None) + return None, None + + sched._free_request.side_effect = _free_request + sched.chunk_transfer_adapter = adapter + sched.running = [request] + sched.waiting_for_transfer_free = set() + sched.transfer_triggered_requests = set() + sched.active_kv_transfers = set() + sched.pending_stop_after_extraction = set() + sched.connector = None + sched.kv_cache_manager.take_events.return_value = None + sched.kv_cache_manager.estimate_cached_tokens.return_value = 0 + sched.finished_req_ids_dict = {} + sched.make_stats.return_value = None + bind_omits_transfer_helpers(sched) + + scheduler_output = mocker.MagicMock() + scheduler_output.num_scheduled_tokens = {request.request_id: 1} + scheduler_output.total_num_scheduled_tokens = 1 + scheduler_output.scheduled_spec_decode_tokens = {} + scheduler_output.num_invalid_spec_tokens = 0 + + model_runner_output = mocker.MagicMock() + model_runner_output.sampled_token_ids = [[42]] + model_runner_output.logprobs = None + model_runner_output.prompt_logprobs_dict = {} + model_runner_output.pooler_output = None + model_runner_output.num_nans_in_logits = None + model_runner_output.kv_connector_output = None + model_runner_output.cudagraph_stats = None + model_runner_output.req_id_to_index = {request.request_id: 0} + model_runner_output.routed_experts = None + model_runner_output.inter_stage_outputs = inter_stage_outputs + + OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output) + return adapter + + +@pytest.mark.parametrize( + ("final_stage_id", "force_kv_transfer", "expect_save"), + [ + pytest.param(0, False, False, id="stage-zero-final"), + pytest.param(1, False, True, id="downstream-final"), + ], +) +def test_finish_saves_async_only_with_downstream_consumer(mocker, final_stage_id, force_kv_transfer, expect_save): + request = _make_chunk_request(final_stage_id, force_kv_transfer=force_kv_transfer) + adapter = _run_finished_save_step(mocker, request) + + assert adapter.save_async.called is expect_save + + +def test_stage_zero_final_skips_inter_stage_output(mocker): + warn = mocker.patch("vllm_omni.core.sched.omni_ar_scheduler.logger.warning") + request = _make_chunk_request(0) + adapter = _run_finished_save_step(mocker, request, inter_stage_outputs=[{"codes": 1}]) + + adapter.save_async.assert_not_called() + assert any("inter_stage_output is present" in str(call) for call in warn.call_args_list) + + +def test_streaming_session_update_invalidates_omits_kv_transfer_cache(): + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler._omits_kv_transfer_cache = {} + scheduler.skipped_waiting = [] + scheduler.log_stats = False + tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0) + session = SimpleNamespace( + request_id="req", + additional_information=tagged.additional_information, + status=RequestStatus.WAITING, + ) + assert scheduler._request_omits_chunk_transfer_to_next_stage(session) + assert "req" in scheduler._omits_kv_transfer_cache + update = SimpleNamespace( + additional_information={"omni_final_stage_id": 1}, + model_intermediate_buffer=None, + arrival_time=0.0, + sampling_params=None, + ) + + OmniSchedulerMixin._finish_streaming_session_update(scheduler, session, update) + + assert scheduler._omits_kv_transfer_cache == {} + assert session.additional_information == {"omni_final_stage_id": 1} + assert not scheduler._request_omits_chunk_transfer_to_next_stage(session) + assert not scheduler._request_omits_kv_transfer_to_next_stage(session) + + +def test_omits_cache_misses_when_payload_object_is_replaced(): + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler._omits_kv_transfer_cache = {} + tagged = apply_omni_final_stage_metadata(_make_engine_request(), final_stage_id=0) + request = SimpleNamespace( + request_id="req", + additional_information=tagged.additional_information, + ) + assert scheduler._request_omits_chunk_transfer_to_next_stage(request) + request.additional_information = {"omni_final_stage_id": 1} + + assert not scheduler._request_omits_chunk_transfer_to_next_stage(request) diff --git a/tests/core/sched/test_omni_ar_scheduler_stale_drain.py b/tests/core/sched/test_omni_ar_scheduler_stale_drain.py index b4eb839a2..e5bd8dd4e 100644 --- a/tests/core/sched/test_omni_ar_scheduler_stale_drain.py +++ b/tests/core/sched/test_omni_ar_scheduler_stale_drain.py @@ -40,6 +40,7 @@ from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler +from tests.helpers.omni_scheduler import bind_omits_transfer_helpers # isort: on @@ -103,6 +104,7 @@ def _make_drain_sched(session: Request) -> MagicMock: sched.kv_cache_manager.estimate_cached_tokens.return_value = 0 sched.finished_req_ids_dict = {} sched.make_stats.return_value = None + bind_omits_transfer_helpers(sched) return sched diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py index 94d796fd6..bdfbadba1 100644 --- a/tests/core/sched/test_omni_ar_scheduler_streaming.py +++ b/tests/core/sched/test_omni_ar_scheduler_streaming.py @@ -26,6 +26,7 @@ from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARS from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( OmniChunkTransferAdapter, ) +from tests.helpers.omni_scheduler import bind_omits_transfer_helpers # isort: on @@ -175,6 +176,7 @@ def _run_resumable_segment_stop( sched.kv_cache_manager.estimate_cached_tokens.return_value = 0 sched.finished_req_ids_dict = {} sched.make_stats.return_value = None + bind_omits_transfer_helpers(sched) scheduler_output = MagicMock(spec=SchedulerOutput) scheduler_output.num_scheduled_tokens = {session.request_id: 1} @@ -312,6 +314,7 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non sched.kv_cache_manager.estimate_cached_tokens.return_value = 0 sched.finished_req_ids_dict = {} sched.make_stats.return_value = None + bind_omits_transfer_helpers(sched) scheduler_output = MagicMock(spec=SchedulerOutput) scheduler_output.num_scheduled_tokens = {session.request_id: 1} @@ -400,6 +403,7 @@ def test_stale_async_frame_is_dropped_before_output_processing() -> None: sched.kv_cache_manager.estimate_cached_tokens.return_value = 0 sched.finished_req_ids_dict = {} sched.make_stats.return_value = None + bind_omits_transfer_helpers(sched) scheduler_output = MagicMock(spec=SchedulerOutput) scheduler_output.num_scheduled_tokens = {session.request_id: 1} diff --git a/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py b/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py index 008e52636..2ecd2772a 100644 --- a/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py +++ b/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations from types import SimpleNamespace @@ -76,7 +79,7 @@ def test_ar_free_request_cleans_input_coordinator_on_normal_free() -> None: scheduler = OmniARScheduler.__new__(OmniARScheduler) scheduler.input_coordinator = coordinator scheduler.chunk_transfer_adapter = None - scheduler._omits_kv_transfer_cache = {"req-free": True} + scheduler._omits_kv_transfer_cache = {"req-free": (1, True, False)} scheduler.encoder_cache_manager = SimpleNamespace(free=lambda request: None) scheduler.finished_req_ids = set() scheduler.finished_req_ids_dict = None diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 8f85366e3..438414854 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -15,6 +15,7 @@ from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats from vllm.v1.request import Request, RequestStatus +from tests.helpers.omni_scheduler import bind_omits_transfer_helpers from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct @@ -95,6 +96,30 @@ def test_shm_abort_reclaims_inflight_put_without_blocking(shm_sender, monkeypatc assert not adapter.code_prompt_token_ids +def test_shm_release_reclaims_undrained_chunks_after_finish(shm_sender): + adapter, connector = shm_sender + ext_id = f"release_{uuid.uuid4().hex}" + request = _req("done", RequestStatus.RUNNING, external_req_id=ext_id) + sibling_key = f"{ext_id}_1_0_0" + assert connector.put("0", "1", sibling_key, "sibling")[0] + for _ in range(3): + adapter.save_async(None, request) + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + # The consumer read only the first chunk before its own request ended. + assert connector.get("0", "1", f"{ext_id}_0_0") is not None + + adapter.release_shm_resources(ext_id) + # The release is queued behind sends; nothing is unlinked on the caller. + assert f"{ext_id}_0_1" in connector._pending_keys + while adapter._pending_save_reqs: + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + + assert adapter._reclaimed_shm_total == 2 + assert not any(k.startswith(f"{ext_id}_0_") for k in connector._pending_keys) + assert connector.get("0", "1", f"{ext_id}_0_1") is None + assert connector.get("0", "1", sibling_key)[0] == "sibling" + + @pytest.mark.parametrize("fail_terminal", [False, True]) def test_shm_terminal_keeps_payload_only_after_success(shm_sender, monkeypatch, fail_terminal): adapter, connector = shm_sender @@ -2620,10 +2645,9 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw)) adapter_mock.save_async = lambda *a, **kw: save_calls.append((a, kw)) - from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler - scheduler = mocker.MagicMock() scheduler.chunk_transfer_adapter = adapter_mock + bind_omits_transfer_helpers(scheduler) scheduler.connector = None scheduler.perf_metrics = None scheduler.log_stats = False diff --git a/tests/distributed/omni_connectors/test_shm_connector.py b/tests/distributed/omni_connectors/test_shm_connector.py index 39a0e8327..d0efd06f5 100644 --- a/tests/distributed/omni_connectors/test_shm_connector.py +++ b/tests/distributed/omni_connectors/test_shm_connector.py @@ -213,7 +213,7 @@ class TestCleanup: connector.put("s0", "s1", "cleanup_req_42", data) assert "cleanup_req_42" in connector._pending_keys - connector.cleanup("cleanup_req_42") + assert connector.cleanup("cleanup_req_42") is True assert "cleanup_req_42" not in connector._pending_keys result = connector.get("s0", "s1", "cleanup_req_42", metadata=None) @@ -224,7 +224,7 @@ class TestCleanup: connector.put("s0", "s1", "consumed_req_99", data) connector.get("s0", "s1", "consumed_req_99", metadata=None) - connector.cleanup("consumed_req_99") + assert connector.cleanup("consumed_req_99") is False assert "consumed_req_99" not in connector._pending_keys def test_close_cleans_all_pending(self, connector): @@ -296,6 +296,24 @@ def test_consumed_bad_payload_removes_lock_file(connector, monkeypatch): assert key not in connector._pending_keys +@pytest.mark.parametrize("suffix", ["_1_2", "_1", "_0", "_suffix"]) +def test_cleanup_prefix_only_matches_own_chunk_keys(connector, suffix): + request_id = f"prefix_{uuid.uuid4().hex}" + own_keys = [f"{request_id}_0_{i}" for i in range(3)] + sibling_key = f"{request_id}{suffix}_0_0" + for key in own_keys: + assert connector.put("0", "1", key, key)[0] + assert connector.put("0", "1", sibling_key, "sibling")[0] + # A consumed chunk is no longer counted as reclaimed. + assert connector.get("0", "1", own_keys[0])[0] == own_keys[0] + + assert connector.cleanup_prefix(f"{request_id}_0_") == 2 + for key in own_keys: + assert key not in connector._pending_keys + assert connector.get("0", "1", key) is None + assert connector.get("0", "1", sibling_key)[0] == "sibling" + + # ── Arrival wakeups ─────────────────────────────────────────────────── diff --git a/tests/e2e/online_serving/test_qwen3_omni_shm_leak.py b/tests/e2e/online_serving/test_qwen3_omni_shm_leak.py new file mode 100644 index 000000000..2afd8bfd8 --- /dev/null +++ b/tests/e2e/online_serving/test_qwen3_omni_shm_leak.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Live /dev/shm leak check for Qwen3-Omni ``async_chunk`` (stage-0-final). + +The default L2 Qwen3-Omni job uses ``ci/qwen3_omni_moe.yaml`` with +``async_chunk: false`` plus ``--no-async-chunk``, so it never exercises the +thinker→talker SharedMemoryConnector path. + +A text-only request is tagged ``omni_final_stage_id=0``, so nothing downstream +ever ``get()``s what stage-0 ``put()``s for it. Before #7245 the scheduler still +``put()`` a finished marker on finish and nothing unlinked it, so every such +request left a ``SharedMemoryConnector`` segment plus lockfile in ``/dev/shm``. +Stage-0-final finishes now skip the ``put()`` and the orchestrator reclaims any +undrained segment on request cleanup. + +This module starts the production async-chunk deploy and asserts the three +text-only request ids leave no connector lockfile or POSIX segment in +``/dev/shm``. +""" + +from __future__ import annotations + +import os +import time +from pathlib import Path + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +_USE_PD = os.environ.get("VLLM_TEST_PD_MODE", "0") == "1" +_MODEL = os.environ.get("VLLM_OMNI_TEST_MODEL", "Qwen/Qwen3-Omni-30B-A3B-Instruct") +_SHM_DIR = Path("/dev/shm") +_NUM_TEXT_REQUESTS = 3 +_SETTLE_S = 15.0 + +# CI overlay keeps seq/token limits small; flip async_chunk back on so stage-0 +# still uses SharedMemoryConnector. Production yaml is already async_chunk:true. +_ASYNC_CHUNK_DEPLOY = modify_stage_config( + get_deploy_config_path("ci/qwen3_omni_moe.yaml"), + updates={"async_chunk": True}, +) + +test_params = [ + pytest.param( + OmniServerParams( + model=_MODEL, + stage_config_path=_ASYNC_CHUNK_DEPLOY, + use_stage_cli=True, + ), + id="async_chunk", + ) +] + + +def _system_prompt() -> dict: + return { + "role": "system", + "content": [ + { + "type": "text", + "text": ( + "You are Qwen, a virtual human developed by the Qwen Team, " + "Alibaba Group, capable of perceiving auditory and visual inputs, " + "as well as generating text and speech." + ), + } + ], + } + + +def _shm_used_mib() -> str: + try: + st = os.statvfs(_SHM_DIR) + used = (st.f_blocks - st.f_bavail) * st.f_frsize + return f"{used / (1024 * 1024):.1f}MiB" + except OSError: + return "unknown" + + +def _is_entry_for_request(name: str, request_id: str) -> bool: + """Match a segment or ``shm_`` lockfile keyed by the ``{request_id}-{suffix}`` orchestrator id.""" + return name.removeprefix("shm_").startswith((f"{request_id}-", f"{request_id}_")) + + +def _entries_for_request_ids(request_ids: list[str]) -> list[str]: + if not _SHM_DIR.is_dir(): + pytest.skip("/dev/shm is not available") + found: list[str] = [] + with os.scandir(_SHM_DIR) as entries: + for entry in entries: + if any(_is_entry_for_request(entry.name, request_id) for request_id in request_ids): + found.append(entry.name) + return found + + +def _wait_no_entries_for_request_ids(request_ids: list[str], timeout_s: float = _SETTLE_S) -> list[str]: + """Reclaim is fire-and-forget; poll until these ids are gone or *timeout_s*.""" + deadline = time.monotonic() + timeout_s + time.sleep(1.0) + leftover: list[str] = _entries_for_request_ids(request_ids) + while leftover: + if time.monotonic() >= deadline: + return leftover + time.sleep(2.0) + leftover = _entries_for_request_ids(request_ids) + return leftover + + +@pytest.mark.advanced_model +@pytest.mark.omni +@pytest.mark.skipif(_USE_PD, reason="Temporarily skip PD mode in this test module.") +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_text_only_async_chunk_does_not_leak_shm(omni_server, online_client) -> None: + """Stage-0-final (text-only) requests must not leave connector SHM behind.""" + print(f"[shm-leak] before requests: /dev/shm_used={_shm_used_mib()}", flush=True) + + messages = dummy_messages_from_mix_data( + system_prompt=_system_prompt(), + content_text="What is the capital of China? Answer in 20 words.", + ) + request_config = { + "model": omni_server.model, + "messages": messages, + "stream": False, + "modalities": ["text"], + "key_words": {"text": ["beijing"]}, + } + responses = online_client.send_omni_request(request_config, request_num=_NUM_TEXT_REQUESTS) + produced = [resp.text_content for resp in responses if getattr(resp, "text_content", None)] + assert len(produced) == _NUM_TEXT_REQUESTS, ( + f"Need {_NUM_TEXT_REQUESTS} text completions to exercise stage-0 finish; " + f"got {len(produced)} non-empty texts. A preprocess/engine miss would " + f"leave /dev/shm unchanged and hide a leak." + ) + request_ids = [resp.request_id for resp in responses] + assert all(request_ids) and len(request_ids) == _NUM_TEXT_REQUESTS, ( + f"Need chat completion ids to assert per-request SHM cleanup; got {request_ids}" + ) + + leftover = _wait_no_entries_for_request_ids(request_ids) + print( + f"[shm-leak] after {_NUM_TEXT_REQUESTS} text-only requests: " + f"/dev/shm_used={_shm_used_mib()} leftover={leftover} request_ids={request_ids}", + flush=True, + ) + assert not leftover, ( + f"SharedMemoryConnector leaked /dev/shm entries for {leftover} after stage-0-final request(s) {request_ids}" + ) diff --git a/tests/engine/test_cfg_companion_lifecycle.py b/tests/engine/test_cfg_companion_lifecycle.py index 6b907c3d2..7fa04d191 100644 --- a/tests/engine/test_cfg_companion_lifecycle.py +++ b/tests/engine/test_cfg_companion_lifecycle.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 """Deterministic tests for the CFG-companion output lifecycle. Regression suite for the companion-output race behind the nightly Bagel @@ -94,6 +94,9 @@ class _FakePool: def release_bindings(self, request_ids): pass + async def release_request_resources(self, request_ids): + pass + def _make_orchestrator(num_stages: int = 2) -> Orchestrator: orch = Orchestrator.__new__(Orchestrator) diff --git a/tests/engine/test_stage_pool_collective_rpc.py b/tests/engine/test_stage_pool_collective_rpc.py index 281b4b504..d15df2bdd 100644 --- a/tests/engine/test_stage_pool_collective_rpc.py +++ b/tests/engine/test_stage_pool_collective_rpc.py @@ -154,6 +154,46 @@ def test_abort_requests_does_not_commit_op_state_when_engine_abort_fails(): asyncio.run(run()) +@pytest.mark.cpu +def test_release_request_resources_skips_without_async_chunk(): + async def run() -> None: + call = AsyncMock() + pool = StagePool( + 0, + [SimpleNamespace(call_utility_async=call)], # type: ignore[list-item] + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(async_chunk=False)), + ) + + await pool.release_request_resources(["req-1"]) + call.assert_not_awaited() + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_release_request_resources_times_out_hung_replica_without_blocking_others(monkeypatch): + async def run() -> None: + async def hang(*_args): + await asyncio.Event().wait() + + live = AsyncMock() + pool = StagePool( + 0, + [SimpleNamespace(call_utility_async=hang), SimpleNamespace(call_utility_async=live)], # type: ignore[list-item] + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(async_chunk=True)), + ) + monkeypatch.setattr(pool, "RELEASE_RPC_TIMEOUT_S", 0.2) + + release = asyncio.create_task(pool.release_request_resources(["req-1"])) + await asyncio.sleep(0.05) + live.assert_awaited_once_with("omni_release_request_resources", ["req-1"]) + assert not release.done() + + await asyncio.wait_for(release, timeout=1.0) + + asyncio.run(run()) + + @pytest.mark.cpu @pytest.mark.parametrize("method", ["reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"]) @pytest.mark.parametrize("failure", ["error", "timeout", "missing"]) diff --git a/tests/helpers/client.py b/tests/helpers/client.py index 8637e510d..883659565 100644 --- a/tests/helpers/client.py +++ b/tests/helpers/client.py @@ -97,6 +97,7 @@ def _split_request_config_by_per_output_sizes(cfg: dict[str, Any]) -> list[dict[ class OmniResponse: """Decoded multimodal / chat output from the OpenAI SDK or offline runner (not raw ``requests``).""" + request_id: str | None = None text_content: str | None = None audio_data: list[str] | None = None audio_content: str | None = None @@ -421,6 +422,8 @@ class OnlineOmniClient: text_content = "" audio_data = [] for chunk in chat_completion: + if result.request_id is None: + result.request_id = getattr(chunk, "id", None) for choice in chunk.choices: content = getattr(getattr(choice, "delta", None), "content", None) modality = getattr(chunk, "modality", None) @@ -453,6 +456,7 @@ class OnlineOmniClient: """Wall clock from *before* ``chat.completions.create`` through response parse + local decode.""" result = OmniResponse() try: + result.request_id = getattr(chat_completion, "id", None) audio_data = None text_content = None for choice in chat_completion.choices: diff --git a/tests/helpers/omni_scheduler.py b/tests/helpers/omni_scheduler.py new file mode 100644 index 000000000..e5a4c58a5 --- /dev/null +++ b/tests/helpers/omni_scheduler.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Scheduler test helpers shared by AR / chunk-transfer unit tests.""" + +from __future__ import annotations + +from types import MethodType + +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler + + +def bind_omits_transfer_helpers(sched) -> None: + """Bind the real omit helpers so a MagicMock scheduler does not skip save_async. + + MagicMock is truthy, so an unbound ``_request_omits_chunk_transfer_to_next_stage`` + silently skips the chunk put. + """ + sched._omits_kv_transfer_cache = {} + sched._omni_final_stage_flags = MethodType(OmniARScheduler._omni_final_stage_flags, sched) + sched._request_omits_kv_transfer_to_next_stage = MethodType( + OmniARScheduler._request_omits_kv_transfer_to_next_stage, + sched, + ) + sched._request_omits_chunk_transfer_to_next_stage = MethodType( + OmniARScheduler._request_omits_chunk_transfer_to_next_stage, + sched, + ) diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 7dc0a5d7c..9c6f2e6d2 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -144,8 +144,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # Track requests that have already triggered prefill transfer to avoid duplicates self.transfer_triggered_requests: set[str] = set() - # Cache per-request flag to avoid repeated deserialization of additional_information - self._omits_kv_transfer_cache: dict[str, bool] = {} + # Cache per-request flags to avoid repeated deserialization of + # additional_information. Value is (payload_id, is_stage_zero_final, + # force_kv_transfer); payload_id invalidates the entry when the + # object is replaced. + self._omits_kv_transfer_cache: dict[str, tuple[int, bool, bool]] = {} # KV-wait start ts for the vllm_omni:kv_wait_s metric; see # _emit_kv_wait_output for the engine-core → orchestrator carry. @@ -264,26 +267,40 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): return config.get(key, default) return getattr(config, key, default) if config is not None else default - def _request_omits_kv_transfer_to_next_stage(self, request: Request) -> bool: - """True when this stage-zero-final request does not need downstream KV. - - The result is cached per request to avoid repeated deserialization of - additional_information on every scheduler tick. - """ - rid = request.request_id - cached = self._omits_kv_transfer_cache.get(rid) - if cached is not None: - return cached - + def _omni_final_stage_flags(self, request: Request) -> tuple[bool, bool]: + """Return ``(is_stage_zero_final, force_kv_transfer)`` from request metadata.""" payload = getattr(request, "additional_information", None) if payload is None: - result = False - else: - info = deserialize_additional_information(payload) - result = info.get("omni_final_stage_id") == 0 and not bool(info.get("omni_force_kv_transfer", False)) + return False, False - self._omits_kv_transfer_cache[rid] = result - return result + cache = getattr(self, "_omits_kv_transfer_cache", None) + rid = request.request_id + payload_id = id(payload) + if cache is not None: + cached = cache.get(rid) + if cached is not None and cached[0] == payload_id: + return cached[1], cached[2] + + info = deserialize_additional_information(payload) + is_final = info.get("omni_final_stage_id") == 0 + force_kv = bool(info.get("omni_force_kv_transfer", False)) + if cache is not None: + cache[rid] = (payload_id, is_final, force_kv) + return is_final, force_kv + + def _request_omits_kv_transfer_to_next_stage(self, request: Request) -> bool: + """True when this stage-zero-final request does not need downstream KV.""" + is_final, force_kv = self._omni_final_stage_flags(request) + return is_final and not force_kv + + def _request_omits_chunk_transfer_to_next_stage(self, request: Request) -> bool: + """True when this request has no downstream chunk consumer. + + CFG companions still force KV transfer but are stage-0-final for + ordinary inter-stage chunks, so they omit ``save_async`` here. + """ + is_final, _ = self._omni_final_stage_flags(request) + return is_final def _should_defer_waiting_admission(self) -> bool: return False @@ -669,6 +686,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): confirmed_num_computed_tokens = None boundary_generation = None + # Only when this step might save. Read additional_information + # before _free_request rewrites it. + omits_chunk_transfer = False + if self.chunk_transfer_adapter is not None and (inter_stage_output is not None or stopped): + omits_chunk_transfer = self._request_omits_chunk_transfer_to_next_stage(request) # Capture before resumable stop handling can clear token history. output_token_ids: Any = getattr(request, "output_token_ids", None) if output_token_ids is None: @@ -787,8 +809,16 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # Invariant: EngineCore returns no partial prefill outputs. assert not prompt_logprobs_tensors - if self.chunk_transfer_adapter is not None and ( - inter_stage_output is not None or is_segment_finished or finished + if omits_chunk_transfer and inter_stage_output is not None: + logger.warning( + "Skipping inter-stage chunk for request %s: " + "omni_final_stage_id=0 but inter_stage_output is present", + req_id, + ) + if ( + self.chunk_transfer_adapter is not None + and not omits_chunk_transfer + and (inter_stage_output is not None or is_segment_finished or finished) ): save_kwargs = { "new_token_ids": new_token_ids, diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 77a568d14..6f9f03d6c 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -229,6 +229,9 @@ class OmniSchedulerMixin(_SchedulerMixinBase): def _finish_streaming_session_update(self, session: Request, update: StreamingUpdate) -> None: """Install payload metadata and send an updated session to admission.""" + cache = getattr(self, "_omits_kv_transfer_cache", None) + if cache is not None: + cache.pop(session.request_id, None) session.additional_information = update.additional_information or None session.model_intermediate_buffer = getattr( update, diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index 09854d22d..d0f2cdb1a 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py @@ -307,9 +307,13 @@ class SharedMemoryConnector(OmniConnectorBase): self._metrics["gets"] += 1 return result - def cleanup(self, request_id: str) -> None: - """Unlink the exact key passed to ``put()``, never a request-id prefix.""" + def cleanup(self, request_id: str) -> bool: + """Unlink the exact key passed to ``put()``, never a request-id prefix. + + Returns True when an unconsumed segment was actually unlinked. + """ key = request_id + unlinked = False with self._pending_keys_lock: self._pending_keys.pop(key, None) try: @@ -317,6 +321,7 @@ class SharedMemoryConnector(OmniConnectorBase): seg.close() seg.unlink() logger.debug("cleanup: unlinked unconsumed SHM segment %s", key) + unlinked = True except FileNotFoundError: pass except Exception as e: @@ -327,6 +332,18 @@ class SharedMemoryConnector(OmniConnectorBase): os.remove(lock_file) except OSError: pass + return unlinked + + def cleanup_prefix(self, key_prefix: str) -> int: + """Unlink every tracked key of the form ``{key_prefix}{chunk_id}``. + + Only keys still in ``_pending_keys`` are considered, and the suffix + must be a bare integer chunk id, so another request whose id merely + shares the prefix is never matched. Returns the unlinked count. + """ + with self._pending_keys_lock: + keys = [k for k in self._pending_keys if k.startswith(key_prefix) and k[len(key_prefix) :].isdigit()] + return sum(self.cleanup(key) for key in keys) def close(self) -> None: """Unlink all remaining tracked SHM segments.""" diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index bd6d7abf1..5e4ef6820 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -28,6 +28,8 @@ from .base import OmniTransferAdapterBase logger = get_connector_logger(__name__) +_RECLAIM_WARN_INTERVAL = 1000 + class _SenderGeneration: """Fence one external request generation without blocking cleanup.""" @@ -196,6 +198,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.request_payload = {} self.code_prompt_token_ids: dict[str, list[torch.Tensor]] = defaultdict(list) self.request_ids_mapping: dict[str, str] = {} + # Save-thread only: running count of segments reclaimed after finish. + self._reclaimed_shm_total = 0 self.waiting_for_chunk_waiting_requests: deque[Any] = deque() self.waiting_for_chunk_running_requests: deque[Any] = deque() @@ -688,6 +692,9 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): for chunk_id in range(num_puts): self.connector.cleanup(f"{key_prefix}_{chunk_id}") return + if "release_shm" in task: + self._release_shm_prefix(task["release_shm"]) + return request = task["request"] external_req_id = request.external_req_id sender_token = task.get("sender_token") @@ -979,8 +986,9 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): Called after a terminal chunk is sent or when the scheduler aborts the request before a terminal chunk can be produced. In-flight sends are - cancelled here and reclaim their own state in ``finally``; cleanup - never waits for connector I/O on the scheduler thread. + cancelled here and reclaim adapter state plus leftover SHM in + ``finally`` after ``put()`` returns; cleanup never waits for connector + I/O on the scheduler thread. Idempotent: calling with an already-cleaned or unknown id is safe. """ @@ -1042,6 +1050,42 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): if cached_ic is not None: cached_ic.pop(external_req_id, None) + def release_shm_resources(self, request_id: str) -> None: + """Queue reclaim of inter-stage segments this request left unconsumed. + + Safe only once every stage has finished with the request: the producer + keeps writing until its own request ends, and a consumer that stopped + early never drains the remainder. Successful terminal chunks are kept + by ``_queue_shm_cleanup_locked`` for exactly that consumer, so the + orchestrator -- the only party that knows all stages are done -- drives + this (see ``Orchestrator._cleanup_request_ids``). + + Like abort cleanup, the unlink runs on the save thread, so the + scheduler never performs SHM I/O. + """ + if not isinstance(self.connector, SharedMemoryConnector): + return + external_req_id = self.request_ids_mapping.get(request_id, request_id) + self._pending_save_reqs.append({"release_shm": f"{external_req_id}_{self.connector.stage_id}_"}) + with self._save_cond: + self._save_cond.notify() + + def _release_shm_prefix(self, key_prefix: str) -> None: + reclaimed = self.connector.cleanup_prefix(key_prefix) + if not reclaimed: + return + # Non-zero means a consumer stopped before draining the producer. + # Common on audio requests, so warn only when the running total + # crosses another _RECLAIM_WARN_INTERVAL boundary. + prev = self._reclaimed_shm_total + self._reclaimed_shm_total = prev + reclaimed + if prev // _RECLAIM_WARN_INTERVAL != self._reclaimed_shm_total // _RECLAIM_WARN_INTERVAL: + logger.warning( + "Reclaimed %d unconsumed inter-stage segments so far; " + "a downstream stage finished before draining its producer", + self._reclaimed_shm_total, + ) + def cleanup( self, request_id: str, diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 8f55280cd..fad9a8f52 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -301,6 +301,7 @@ class OrchestratorBase: _transfer_emitter: Any = None _prom_metrics: Any = None _stat_logger: OmniPrometheusStatLogger | None = None + _transfer_release_tasks: set[asyncio.Task] = set() def __init__( self, @@ -348,6 +349,8 @@ class OrchestratorBase: self._pd_bootstrap_addr = pd_config.get("bootstrap_addr") self._pd_prefill_engine_id = pd_config.get("prefill_engine_id") self.request_states: dict[str, OrchestratorRequestState] = {} + # Strong refs for in-flight releases; the loop only weak-refs tasks, so dropping these risks mid-flight GC. + self._transfer_release_tasks: set[asyncio.Task] = set() self._init_metrics_state( stage_pools, running_counter, @@ -690,6 +693,36 @@ class OrchestratorBase: output_msg.finished = index == last_index_by_req[output_msg.request_id] return abort_outputs + def _release_stage_transfer_resources(self, request_ids: list[str]) -> None: + """Drop each stage's inter-stage transfer resources for finished requests. + + This is the only point that knows every stage is done with the request, + so it is the only safe place to reclaim segments a consumer never + drained. Scheduled rather than awaited: reclaim is best-effort and must + not add RPC latency to request teardown. + """ + if not request_ids: + return + + async def _run() -> None: + results = await asyncio.gather( + *(pool.release_request_resources(request_ids) for pool in self.stage_pools), + return_exceptions=True, + ) + for result in results: + if isinstance(result, Exception): + logger.warning("[Orchestrator] release transfer resources failed: %s", result) + + try: + task = asyncio.get_running_loop().create_task(_run()) + self._transfer_release_tasks.add(task) + task.add_done_callback(self._transfer_release_tasks.discard) + except RuntimeError: + logger.warning( + "[Orchestrator] no running event loop; skipped reclaim of transfer resources for %s", + request_ids, + ) + def _release_request_bindings(self, request_ids: list[str]) -> None: """Release all stage-local route bindings for the given request ids.""" for pool in self.stage_pools: @@ -1645,6 +1678,7 @@ class OrchestratorBase: if abort: abort_outputs = await self._abort_request_ids(cleanup_ids) self._release_request_bindings(cleanup_ids) + self._release_stage_transfer_resources(cleanup_ids) for request_id in cleanup_ids: self._pd_kv_params.pop(request_id, None) req_state = self.request_states.pop(request_id, None) diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index c19a5cf75..f9fcd34b9 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -131,6 +131,21 @@ class StageEngineCoreProc(EngineCoreProc): if _bind_native_data_plane_ready_sink(self.model_executor, self.scheduler): logger.info("Bound native MRv2 connector readiness directly to the scheduler inbox.") + def omni_release_request_resources(self, request_ids: list[str]) -> None: + """Release this stage's inter-stage transfer resources for *request_ids*. + + Invoked over the UTILITY channel by the orchestrator once every stage + has finished with the request. Idempotent and safe for unknown ids. + """ + adapter = getattr(getattr(self, "scheduler", None), "chunk_transfer_adapter", None) + if adapter is None: + return + for request_id in request_ids or (): + try: + adapter.release_shm_resources(request_id) + except Exception as e: + logger.debug("omni_release_request_resources(%s) failed: %s", request_id, e) + def preprocess_add_request(self, request: OmniEngineCoreRequest) -> tuple[Any, int]: """Preserve omni payloads when vLLM builds its scheduler request.""" scheduler_request, current_wave = super().preprocess_add_request(request) diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 222cd3b5d..176eb9dc0 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -99,6 +99,8 @@ class StagePool: DISPATCH_WAIT_TIMEOUT_S: float = 10.0 DISPATCH_RETRY_INTERVAL_S: float = 0.1 + # A replica that dies mid-release never answers; without a bound its background release never finishes. + RELEASE_RPC_TIMEOUT_S: float = 5.0 # Only these EngineCore helpers may skip collective_rpc_async. A generic # ``{method}_async`` on AsyncMPClient must not silently drop timeout. _CACHE_RESET_METHODS = frozenset({"reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"}) @@ -135,6 +137,9 @@ class StagePool: self.clients: list[StagePoolClient | None] = list(normalized_clients) self._output_processor = output_processor self._stage_vllm_config = stage_vllm_config + self._has_chunk_transfer_adapter = bool( + getattr(getattr(stage_vllm_config, "model_config", None), "async_chunk", False) + ) self._next_replica_id = 0 self._request_bindings: dict[str, int] = {} self._unavailable_replicas: set[int] = set() @@ -1304,6 +1309,28 @@ class StagePool: return abort_outputs + async def release_request_resources(self, request_ids: list[str]) -> None: + """Ask every live replica to drop transfer resources for *request_ids*. + + Broadcast rather than binding-routed: the orchestrator releases route + bindings as part of the same teardown, so a binding lookup here would + race it. The engine-core handler is idempotent for unknown ids. + """ + if not request_ids or not self._has_chunk_transfer_adapter: + return + ids = list(request_ids) + + async def release(replica_id: int, call: Any) -> None: + try: + await asyncio.wait_for(call("omni_release_request_resources", ids), timeout=self.RELEASE_RPC_TIMEOUT_S) + except Exception as e: + logger.warning( + "[StagePool-%s] release_request_resources on replica %s failed: %r", self.stage_id, replica_id, e + ) + + calls = [(i, getattr(self.clients[i], "call_utility_async", None)) for i in self.live_replica_ids()] + await asyncio.gather(*(release(i, call) for i, call in calls if call is not None)) + async def collective_rpc( self, replica_id: int, From 3bc3f1a7d0a139df23e19566e502276e5d998e49 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Fri, 2 Oct 2026 11:03:14 -0700 Subject: [PATCH 21/22] [Model] Optimize MOSS Local 1.5 MRV2 serving and progressive chunks (#8213) Signed-off-by: Sy03 <1370724210@qq.com> Signed-off-by: Yueqian Lin Co-authored-by: Yueqian Lin --- docs/configuration/environment_variables.md | 6 +- docs/configuration/stage_configs.md | 47 ++ recipes/OpenMOSS/MOSS-TTS.md | 297 ++++++++++ tests/config/test_config_factory.py | 68 ++- tests/config/test_environment_variables.py | 4 +- .../sched/test_generation_batch_coalescing.py | 253 ++++++++ .../sched/test_generation_scheduler_init.py | 102 ++++ .../test_omni_sched_deferred_free_fence.py | 5 + tests/engine/test_stage_engine_core_proc.py | 66 +++ .../openai_api/test_moss_local_prompt.py | 85 +++ .../openai_api/test_tts_adapter.py | 39 ++ tests/entrypoints/test_pd_disaggregation.py | 3 + tests/metrics/test_emit_calls.py | 4 +- tests/metrics/test_stats.py | 39 ++ .../moss_tts/test_codec_async_output.py | 338 +++++++++++ .../moss_tts/test_codec_ring_headroom.py | 146 +++++ .../moss_tts/test_first_chunk_fast_path.py | 547 ++++++++++++++++++ .../moss_tts/test_local_code_rows_snapshot.py | 32 + .../moss_tts/test_local_compiled_sampler.py | 37 ++ .../moss_tts/test_local_determined_tokens.py | 124 ++++ .../models/moss_tts/test_local_eager_mtp.py | 166 ++++++ .../models/moss_tts/test_local_model_state.py | 286 +++++++++ .../moss_tts/test_local_prefill_batch.py | 86 +++ .../moss_tts/test_local_short_attention.py | 82 +++ .../models/moss_tts/test_reference_encoder.py | 233 ++++++++ .../moss_tts/test_reference_encoder_graphs.py | 320 ++++++++++ .../moss_tts/test_shared_reference_encoder.py | 339 +++++++++++ .../models/moss_tts/test_slot_attention.py | 301 ++++++++++ ...t_streaming_decoder_wrapper_concurrency.py | 86 +++ .../moss_tts/test_streaming_terminal_batch.py | 32 +- .../test_moss_raw_async_chunk.py | 80 +++ .../test_moss_raw_empty_terminal.py | 93 +++ .../test_moss_tts_async_chunk.py | 137 +++++ tests/worker/test_omni_connector_mixin.py | 106 +++- tests/worker_v2/test_first_audio_sender.py | 16 + tests/worker_v2/test_omni_ar_model_runner.py | 31 + tests/worker_v2/test_omni_model_state.py | 55 ++ .../config/environment_variable_inventory.py | 18 + .../core/sched/omni_generation_scheduler.py | 144 +++++ vllm_omni/core/sched/omni_scheduler_mixin.py | 14 +- .../core/sched/omni_scheduling_coordinator.py | 7 +- vllm_omni/deploy/moss_tts_local.yaml | 3 + vllm_omni/deploy/moss_tts_local_mrv2.yaml | 30 + .../moss_tts_local_mrv2_high_concurrency.yaml | 53 ++ .../moss_tts_local_mrv2_low_latency.yaml | 66 +++ .../omni_connector_payload_transport.py | 19 +- vllm_omni/engine/stage_engine_core_proc.py | 24 +- vllm_omni/entrypoints/omni_base.py | 2 +- .../entrypoints/openai/tts_adapters/base.py | 13 +- .../openai/tts_adapters/moss_tts.py | 105 +++- vllm_omni/metrics/stats.py | 99 ++-- .../models/moss_tts/audio_tokenizer_v2.py | 131 ++++- .../cuda_graph_streaming_decoder_wrapper.py | 35 +- .../models/moss_tts/first_chunk_fast_path.py | 333 +++++++++++ .../models/moss_tts/local_model_state.py | 447 ++++++++++++++ .../models/moss_tts/local_short_attention.py | 100 ++++ .../moss_tts/modeling_moss_tts_codec.py | 477 +++++++++++++-- .../moss_tts/modeling_moss_tts_local_depth.py | 31 +- .../moss_tts/modeling_moss_tts_talker.py | 26 + .../models/moss_tts/pipeline.py | 2 + .../models/moss_tts/reference_encoder.py | 448 ++++++++++++-- .../moss_tts/reference_encoder_graphs.py | 399 +++++++++++++ .../moss_tts/shared_reference_encoder.py | 334 +++++++++++ .../models/moss_tts/slot_attention.py | 475 +++++++++++++++ vllm_omni/model_executor/output_snapshot.py | 125 ++++ .../stage_input_processors/moss_tts.py | 72 +-- vllm_omni/worker_v2/first_audio_sender.py | 19 +- vllm_omni/worker_v2/model_states/__init__.py | 6 + vllm_omni/worker_v2/omni_ar_model_runner.py | 24 + vllm_omni/worker_v2/output_snapshot.py | 94 +-- 70 files changed, 8519 insertions(+), 347 deletions(-) create mode 100644 tests/core/sched/test_generation_batch_coalescing.py create mode 100644 tests/core/sched/test_generation_scheduler_init.py create mode 100644 tests/entrypoints/openai_api/test_moss_local_prompt.py create mode 100644 tests/model_executor/models/moss_tts/test_codec_async_output.py create mode 100644 tests/model_executor/models/moss_tts/test_codec_ring_headroom.py create mode 100644 tests/model_executor/models/moss_tts/test_first_chunk_fast_path.py create mode 100644 tests/model_executor/models/moss_tts/test_local_code_rows_snapshot.py create mode 100644 tests/model_executor/models/moss_tts/test_local_compiled_sampler.py create mode 100644 tests/model_executor/models/moss_tts/test_local_determined_tokens.py create mode 100644 tests/model_executor/models/moss_tts/test_local_eager_mtp.py create mode 100644 tests/model_executor/models/moss_tts/test_local_model_state.py create mode 100644 tests/model_executor/models/moss_tts/test_local_prefill_batch.py create mode 100644 tests/model_executor/models/moss_tts/test_local_short_attention.py create mode 100644 tests/model_executor/models/moss_tts/test_reference_encoder_graphs.py create mode 100644 tests/model_executor/models/moss_tts/test_shared_reference_encoder.py create mode 100644 tests/model_executor/models/moss_tts/test_slot_attention.py create mode 100644 tests/model_executor/models/moss_tts/test_streaming_decoder_wrapper_concurrency.py create mode 100644 tests/model_executor/stage_input_processors/test_moss_raw_async_chunk.py create mode 100644 tests/model_executor/stage_input_processors/test_moss_raw_empty_terminal.py create mode 100644 tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py create mode 100644 vllm_omni/deploy/moss_tts_local_mrv2.yaml create mode 100644 vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency.yaml create mode 100644 vllm_omni/deploy/moss_tts_local_mrv2_low_latency.yaml create mode 100644 vllm_omni/model_executor/models/moss_tts/first_chunk_fast_path.py create mode 100644 vllm_omni/model_executor/models/moss_tts/local_model_state.py create mode 100644 vllm_omni/model_executor/models/moss_tts/local_short_attention.py create mode 100644 vllm_omni/model_executor/models/moss_tts/reference_encoder_graphs.py create mode 100644 vllm_omni/model_executor/models/moss_tts/shared_reference_encoder.py create mode 100644 vllm_omni/model_executor/models/moss_tts/slot_attention.py create mode 100644 vllm_omni/model_executor/output_snapshot.py diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md index 9f0e1cc91..a872eaf3c 100644 --- a/docs/configuration/environment_variables.md +++ b/docs/configuration/environment_variables.md @@ -214,7 +214,7 @@ collection, examples, bug reports, or logs. ## Model-specific variables -The inventory includes 66 variables read by a single model or pipeline family. They are +The inventory includes 93 variables read by a single model or pipeline family. They are not listed as public usage options here because doing so would turn implementation escape hatches into an accidental compatibility contract. @@ -223,10 +223,10 @@ Every audited model-specific name has a migration disposition in the | Disposition | Count | Required outcome | | --- | ---: | --- | -| Promote | 38 | Move a stable setting into typed stage or model configuration. | +| Promote | 64 | Move a stable setting into typed stage or model configuration. | | Request scope | 5 | Move request-varying behavior into a declared request-option schema. | | External | 0 | Retain only when a supported third-party library owns the contract. | -| Internalize | 16 | Keep a debug or diagnostic switch out of public documentation and configuration. | +| Internalize | 17 | Keep a debug or diagnostic switch out of public documentation and configuration. | | Deprecate/remove | 7 | Remove a compatibility escape hatch that has no continuing contract. | The disposition is a migration target, not a statement that the existing diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 2cf7327dd..214eb3187 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -514,3 +514,50 @@ for that comparison; this switch only controls automatic MPS management. MPS does not reserve a GPU. Use only assigned GPUs, explicitly place stages on the intended GPU, and warm the complete pipeline before measuring performance. + +## MOSS-TTS Local 1.5 with Model Runner V2 + +`OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5` can opt into CUDA MRV2 using +the shared runtime introduced for Qwen3-TTS: + +```bash +vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \ + --deploy-config vllm_omni/deploy/moss_tts_local_mrv2.yaml +``` + +This profile inherits the batching, codec graph buckets and 1-frame/15-frame +chunk geometry from `moss_tts_local.yaml`, and selects V2 for both the Local +Talker and codec stages. The Local depth predictor exposes the MRV2 `mtp` +capabilities while retaining its V1 `talker_mtp` implementation, sampling +defaults and explicit request-seed handling. Codec chunks use the native data +plane; the internal Talker retains its final-only orchestrator output policy. + +On CUDA, the profile bounds stage-0 prefill work to 512 tokens per iteration. +The codec retains its own CUDA Graph capture/replay, including batch buckets +through 64, but disables Inductor compilation (`compilation_config.mode: 0`) +of the stateful decoder to avoid lengthy compilation at startup. This is +independent of `enforce_eager`; the profile keeps `enforce_eager: false`. +It explicitly selects `cudagraph_mode: FULL`, which works without Inductor; +the default `FULL_AND_PIECEWISE` would otherwise normalize to `NONE` and clear +the codec's capture buckets when compilation is disabled. + +For sustained C128 serving on a large-memory CUDA GPU, select the +separate `moss_tts_local_mrv2_high_concurrency.yaml` profile. It uses 128 +stream slots per stage, a 32 GiB Talker KV budget, and the codec's +`triton_slot` backend with Inductor compilation and codec-owned CUDA graphs. +The bounded KV budget leaves room for codec state and graphs; it does not +guarantee that 128 maximum-length prompts fit simultaneously. See the +[MOSS recipe](gh-file:recipes/OpenMOSS/MOSS-TTS.md#local-15-mrv2-and-slot-attention) +for activation, backend comparisons, memory requirements and benchmark commands. + +Omitting `--deploy-config`, or selecting `moss_tts_local.yaml`, retains V1. +NPU, XPU, ROCm and MUSA overrides also retain V1. This profile does not enable +MRV2 for MOSS Delay, Realtime or Nano. Local 1.5 outputs 48 kHz stereo audio; +set `VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000` and +`VLLM_OMNI_BENCH_AUDIO_CHANNELS=2` when benchmarking raw PCM. + +Event-driven orchestration remains independently selectable with +`VLLM_OMNI_EVENT_DRIVEN_ORCH=0` or `1`. Keep the runner and deployment identical +when comparing these modes. Model-runner selection does not change the +orchestration default or enable experimental reference encoding, chunk ramps, +generation-output draining or MPS. diff --git a/recipes/OpenMOSS/MOSS-TTS.md b/recipes/OpenMOSS/MOSS-TTS.md index 8d37039d8..c9ec7670b 100644 --- a/recipes/OpenMOSS/MOSS-TTS.md +++ b/recipes/OpenMOSS/MOSS-TTS.md @@ -35,6 +35,10 @@ the requested `input` speech. Without a nonblank `ref_text`, Local uses audio-reference generation. Reference transcripts must match the reference audio; they are not style instructions. +For the optional CUDA MRV2 runner, see the +[Local 1.5 deployment profile](../../docs/configuration/stage_configs.md#moss-tts-local-15-with-model-runner-v2). +The default Local deployment continues to use V1. + ## References - Offline inference example: [`examples/offline_inference/text_to_speech/moss_tts/`](../../examples/offline_inference/text_to_speech/moss_tts/) @@ -152,3 +156,296 @@ curl -X POST http://localhost:8091/v1/audio/speech \ - No `ref_audio` required or accepted for MOSS-SoundEffect. - Input field maps to the `ambient_sound` parameter in the upstream processor. - Rate: ~12.5 tokens per second; longer descriptions produce longer audio. + +## Local 1.5 MRV2 and slot attention + +`MOSS-TTS-Local-Transformer-v1.5` supports the native CUDA MRV2 pipeline with +an explicit deploy profile. The default Local profile continues to use V1. +Both profiles below preserve 1-frame initial and 15-frame steady codec chunks +and the model's sampling defaults. + +```bash +vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \ + --deploy-config vllm_omni/deploy/moss_tts_local_mrv2.yaml +``` + +This C64 profile bounds Talker prefill to 512 tokens and retains codec-owned +CUDA graphs without compiling the codec with Inductor. + +For sustained high concurrency on a large-memory CUDA GPU, use: + +```bash +CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=4 \ +vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \ + --deploy-config vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency.yaml \ + --stage-init-timeout 1200 --init-timeout 1500 +``` + +The high-concurrency profile places both stages on one GPU, sets each stage's +capacity to 128 and fixes the Talker KV budget at 32 GiB. The Talker uses +mixed FULL CUDA graphs with both its token budget and capture limit set to +512, so new prompt chunks remain within graph coverage. It also enables +`hf_overrides.mrv2_gpu_slot_state`: Local hidden states, audio codes and +continuation control stay in GPU request slots, gathered through MRV2's +batch mapping. Prefill conditioning and explicit per-request sampling seeds +retain their original behavior. Published audio codes own their storage, so +slot reuse cannot overwrite an in-flight output. Set the override to `false` +to compare with the generic Omni model state. The profile selects +`codec_attention_backend: triton_slot`, Inductor mode 3 with combo kernels +disabled, and codec CUDA graph buckets through 128. The configuration was +validated on one H200. The earlier capacity-256 variant used about 110 GiB of +sampled peak GPU memory, including loading and warmup. Reduce the stage capacities, +graph buckets and coalescing target together when adapting it to smaller GPUs. Cold codec +compilation can take several minutes; subsequent starts can reuse the AOT +cache. Graph capture is still performed at startup. + +The high-concurrency profile also coalesces ready codec streams with +`connectors.shm.extra.generation_min_batch_size: 32` and +`generation_max_wait_ms: 12`. The target controls when to stop waiting; it +neither caps execution at 32 nor guarantees every codec group has 32 rows. +First and final chunks share the bounded window at high concurrency. When +fewer than 32 requests remain, dispatch is immediate. Pending outputs are +retired before waiting, with the original deadline preserved, so waiting does +not prevent the previous batch from releasing its in-flight state. Cancellation +and input notifications keep their ordinary scheduler bookkeeping. This option +requires a stateful native MRV2 generation stage with TP1/PP1. Set the wait to +`0` to disable coalescing; other profiles keep the immediate-dispatch default. + +The high-concurrency profile enables two additional Local model-state +optimizations under the CUDA stage-0 `hf_overrides`, alongside +`mrv2_gpu_slot_state: true`: + +- `mrv2_batch_prefill: true` reuses the batch's text embeddings and combines + CPU reference-code slices and their destination positions into one pinned + upload. Host staging is reused only after the upload event completes; + GPU-resident references keep the original path. +- `mrv2_direct_tokens: true` returns the already determined text token after + the Local audio/stop decision. It avoids constructing and sampling the + full text vocabulary. Audio-code and binary-stop sampling are unchanged. + Requests requiring distribution metadata or token constraints, and + unsupported execution modes, retain the normal sampler. + +Other profiles leave both flags disabled. Set both flags explicitly when +comparing their combined effect. Deployment +inheritance replaces a stage's `hf_overrides` mapping, so an overlay must also +retain `mrv2_gpu_slot_state: true`. Measure first-packet latency as well as +throughput: faster prompt preparation can change competition between the +Talker and codec sharing the GPU. + +For the single-H200 capacity-128 throughput baseline, keep both switches on. +In complete Seed-TTS EN1088 runs at client concurrency 128, four combined-switch +rounds pooled 344.11 audio-s/s versus 335.78 for batch prefill alone; direct +tokens alone had no reliable gain. The combined path had 685 ms mean first +audio versus 679 ms across six control rounds, with lower mean completion time. +A fixed-seed 128-row audio sample showed mean Whisper WER of 3.68% versus +3.78% for the control. This sample does not establish speaker similarity or +exclude small quality changes. At concurrency 256, the combined path had +raised mean first-audio latency by about 0.33 s, so other profiles leave the +switches off until their latency and quality are checked. + +The optional connector setting `generation_coalescing_policy: idle_wait` +waits for an inbox notification when no stream is runnable, all receivers +are parked and registered, and no output needs retirement. The first ready +arrival starts a fresh coalescing window, preserving the batch budget. +Control messages, including cancellations, can wait up to two windows +instead of one (24 ms at the 12 ms setting). The default policy remains +`fixed`; neither policy changes the global orchestration default. + +The codec backends differ in state access: + +- `sdpa` uses PyTorch attention after gathering and updating the ring cache. +- `triton` replaces the attention calculation, retaining the ring-cache + gather/copy and explicit mask construction. +- `triton_slot` writes surviving K/V directly into request slots, attends + directly to the ring, and advances the active slot offsets in three + ordered kernels. Graph-padding rows do not advance persistent state. + +The slot path preserves the existing chunk-complete ring semantics, including +retaining the final cache-capacity tokens when a chunk exceeds the ring. +It does not introduce another cache owner or change request-slot lifetime. +Zero-length padding rows skip the attention computation and emit zeros. + +The high-concurrency profile includes intermediate codec batch buckets 6, 12 +and 24. To compare with power-of-two buckets, change the stage-1 +`cudagraph_capture_sizes` to `[1, 2, 4, 8, 16, 32, 64, 128]`. +Keep the maximum bucket equal to the state capacity to retain terminal-tail +coalescing. Smaller buckets reduce padding work but require more graphs; +measure complete serving runs before selecting them for a deployment. + +For a codec output-path comparison, set +`connectors.shm.extra.codec_gpu_stream_output: 0` to select the synchronous +codec output path; `1` selects GPU output snapshots on a private codec +stream. Both configurations use MRV2 and keep streaming audio responses. +The option changes input ownership, metadata staging and snapshot handling +as well as transfer scheduling, so its timing difference is not just D2H +copy time. + +Those chunk-complete semantics truncate the causal window: every ring holds +exactly `context` entries and a whole chunk is written before attention runs, +so token `i` of a `T`-token chunk sees `capacity - T + i + 1` keys. With the +tokenizer-v2 decoder's per-layer contexts (10/10/8/4/2/1 s, i.e. 125/250/400/ +400/400/400 tokens) a 15-frame chunk is 480 tokens at the last transformer, so +its first 80 tokens attend to nothing and no chunk sees the previous one there. +Padding a terminal tail to 15 frames therefore decodes it differently from an +exact-length decode; the difference is deterministic (identical in fp32 and +fp64), not bf16 noise. The connector option `codec_ring_headroom: 1` sizes each +ring as `context + max_chunk_frames * tokens_per_frame` (including ramp shapes); chunked streaming +then reproduces whole-sequence decoding bit-for-bit in fp64 and padded tails +equal exact tails. It increases state/graph memory and capture time; historical C256 experiments +observed about a 2% throughput cost. It is off by default and its perceptual +quality benefit has not been established. Terminal tails are padded +into the regular 15-frame graph bucket in either mode. +The slot kernel is specific to the CUDA tokenizer-v2 decoder; other paths +retain their existing attention implementation. Event-driven orchestration +remains independently selectable and its default is unchanged. + +### Optional progressive chunks + +Set `codec_chunk_ramp` under `connectors.shm.extra` to insert smaller chunks +before steady decoding, for example: + +```yaml +connectors: + shm: + extra: + codec_chunk_ramp: [1, 4, 15] +``` + +The ramp's first entry overrides `initial_codec_chunk_frames`. After the last +entry, chunks use `codec_chunk_frames`. Each request advances independently; +the final partial chunk is flushed and an empty terminal carries no codec +tokens. The codec captures every ramp length, sizes the maximum execution +step for the largest entry, and uses the ramp's first entry for its dedicated +first-chunk graph. Additional shapes increase startup and graph memory. + +Ramps are opt-in. They can reduce early playback underrun while increasing +codec calls, first-audio latency or total generation time. The measured C128 +throughput presets retain 1→15; their published figures do not establish an +additional benefit from enabling ramps. The processor and graph-shape tests +run without GPU or model weights: + +```bash +CUDA_VISIBLE_DEVICES= PYTHONPATH=. python -m pytest -q \ + tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py \ + -m 'core_model and cpu' --run-level core_model +``` + +### Low-latency and reference-encoding options + +`moss_tts_local_mrv2_low_latency.yaml` keeps both stage capacities at 128 and +uses prefill 2048, GPU slot state, batch prefill, direct tokens and eager MTP. +Its codec first-chunk fast path uses dedicated graphs and a stream handoff +before the regular decoder reuses the request slot. A gate limits contention +with regular codec work. Slot waits are bounded at 30 seconds. A failed or +stalled decode raises an error rather than replaying a partially advanced +slot; a failed output enqueue retains owned PCM for regular delivery. +Closing the decoder rejects new jobs. Its regular-batch limit of 16 is internal dispatch +policy; it does not change the client concurrency or stage capacities. + +Reference encoding runs in the API layer, independently of MRV2. Enable +reference graphs explicitly with `VLLM_OMNI_MOSS_REF_GRAPHS=1`; the default +eager encoder is retained when the option is absent. The graph path uses the +loaded tokenizer's encoder/quantizer, length buckets and optional compilation +(`VLLM_OMNI_MOSS_REF_COMPILE=0` disables compilation). It also uses windowed +attention; `VLLM_OMNI_MOSS_REF_ATTN=sdpa` retains the original attention. +Compilation and attention changes need not produce bit-identical codes. + +For multiple API processes, `VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR` enables shared +reference-code storage and, by default, a single encoder host with four +workers. Use a dedicated directory per service, checkpoint and encoding +configuration. Workers have their own graph resources; all graph captures +complete before serving begins. `VLLM_OMNI_MOSS_REF_SHARED_ENCODER=0` retains +separate encoders while sharing codes. More workers/graphs consume memory; +they do not imply more GPUs. Example after selecting an available GPU and +configuring any operator-managed MPS daemon: + +```bash +VLLM_WORKER_MULTIPROC_METHOD=spawn \ +VLLM_OMNI_EVENT_DRIVEN_ORCH=1 \ +VLLM_OMNI_CONNECTOR_RECV_POLL_MS=1 \ +VLLM_OMNI_MOSS_REF_GRAPHS=1 \ +VLLM_OMNI_MOSS_REF_ENCODER_WORKERS=4 \ +VLLM_OMNI_MOSS_REF_BATCH_WINDOW_MS=0 \ +VLLM_OMNI_MOSS_REF_INFLIGHT=4 \ +VLLM_OMNI_MOSS_REF_HOST_WINDOW_MS=2 \ +VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR=/dev/shm/moss-local-service \ +vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 --omni \ + --api-server-count 4 \ + --deploy-config vllm_omni/deploy/moss_tts_local_mrv2_low_latency.yaml \ + --stage-init-timeout 1200 --init-timeout 1500 +``` + +This example is a configurable serving profile, not the complete launch +configuration of a historical benchmark. In particular, the low-latency YAML +uses synchronous stage-1 scheduling; the retained source46 experiment used +asynchronous stage-1 scheduling, FP8 backbone, MPS and additional encoder +environment settings. Do not infer a throughput result from the YAML alone. + +### Cold versus hot reference measurements + +A cold reference misses the encoded-reference cache: the request pays for +reference parsing/preparation and GPU encoding before speech generation. +A hot reference reuses reference codes but still synthesizes new speech; +generated audio is not cached. Pre-registering a reference moves encoding +into registration and can make subsequent synthesis hot, but registration +latency is still part of the first-use cost. + +Model startup/compilation is separate from reference coldness. EN1088 contains +repeated references, so its first pass is not uniformly cold. For an all-unique +cold pass, use its 666 first-occurrence references, keep order, warm execution +with disjoint references, then repeat the same 666 requests. A second pass +through independent API-local caches is not guaranteed hot; verify cache +coverage before comparing it with a shared-cache result. Report requests/s +alongside audio-s/s and retain failed, empty, long and near-silent outputs. + +### Reproduce the serving benchmark + +Use the complete Seed-TTS English test set, including reference audio and text. +Set `SEED_TTS_DATA` to the directory containing `en/meta.lst` and its 1088 +entries. Run the native benchmark after the server is ready: + +```bash +export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=48000 +export VLLM_OMNI_BENCH_AUDIO_CHANNELS=2 +export SEED_TTS_WER_EVAL=0 + +for concurrency in 128; do + for phase in warm r1 r2; do + vllm bench serve --omni \ + --model OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 \ + --backend openai-audio-speech --endpoint /v1/audio/speech \ + --dataset-name seed-tts --dataset-path "$SEED_TTS_DATA" \ + --seed-tts-locale en --disable-shuffle --num-prompts 1088 \ + --num-warmups 0 --ready-check-timeout-sec 0 \ + --output-len 256 --max-concurrency "$concurrency" \ + --request-rate inf --seed 42 \ + --extra-body '{"task_type":"Base","max_new_tokens":256}' \ + --save-result --save-detailed --result-dir results/moss-local \ + --result-filename "c${concurrency}-${phase}.json" + done +done +``` + +Discard the complete `warm` pass and combine measured runs as total generated +audio seconds divided by total benchmark duration. Require 1088 successes, +zero failures and 1088 nonempty-audio metric samples in each run. The output +cap matches the benchmark protocol; success alone does not establish speech +quality or that every sentence ended before the cap. Dataset seed 42 does +not fix an independent sampling seed for every request. + +For an attention-only comparison, copy the high-concurrency YAML beside the +original to preserve relative `base_config` resolution. Change its +codec `compilation_config.mode` to `0` while retaining `cudagraph_mode: FULL`, +and compare `codec_attention_backend: triton` against `triton_slot`. Keep all +other settings, warmup and client concurrency identical. Restart the server +between configurations and use separate result directories. Comparing the +compiled slot profile to uncompiled Triton includes both changes. + +Kernel and MHA regression tests require CUDA, compatible vLLM/Triton packages, +and no model weights: + +```bash +python -m pytest -q tests/model_executor/models/moss_tts/test_slot_attention.py \ + tests/model_executor/models/moss_tts/test_streaming_attention.py \ + -m 'core_model and cuda' --run-level=core_model +``` diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index db97f990d..64f7cd995 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -2928,9 +2928,17 @@ class TestPlatformOverrides: assert deploy.stages[0].engine_extras["block_size"] == 128 - @pytest.mark.parametrize("platform", ["cuda", "npu"]) - def test_recommended_native_runner_platform_defaults(self, platform): - filename, pipeline_key = "qwen3_tts_mrv2.yaml", "qwen3_tts" + @pytest.mark.parametrize( + "filename,pipeline_key", + [ + ("qwen3_tts_mrv2.yaml", "qwen3_tts"), + ("moss_tts_local_mrv2.yaml", "moss_tts_local"), + ("moss_tts_local_mrv2_high_concurrency.yaml", "moss_tts_local"), + ("moss_tts_local_mrv2_low_latency.yaml", "moss_tts_local"), + ], + ) + @pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"]) + def test_recommended_native_runner_platform_defaults(self, platform, filename, pipeline_key): deploy = load_deploy_config(Path(get_deploy_config_path(filename))) deploy = _apply_platform_overrides(deploy, platform=platform) expect_v2 = platform == "cuda" @@ -2946,6 +2954,60 @@ class TestPlatformOverrides: # v2 only engages the native plane on stages declaring support. assert all(ps.supports_native_mrv2_data_plane for ps in pipeline.stages) + def test_moss_local_mrv2_preserves_base_profile_and_variant_scope(self): + pipeline = resolve_pipeline_config("moss_tts_local") + base_path = Path(get_deploy_config_path(pipeline.default_deploy_config_name)) + candidate_path = Path(get_deploy_config_path("moss_tts_local_mrv2.yaml")) + base = load_deploy_config(base_path) + candidate = load_deploy_config(candidate_path) + assert base.model_runner == "v1" + assert candidate.stages == base.stages + assert candidate.connectors == base.connectors + assert candidate.async_chunk == base.async_chunk + cuda = _apply_platform_overrides(candidate, platform="cuda") + stages = merge_pipeline_deploy(pipeline, cuda) + assert stages[0].yaml_engine_args["max_num_batched_tokens"] == 512 + codec_args = stages[1].yaml_engine_args + assert codec_args["max_num_seqs"] == 64 + assert codec_args["enforce_eager"] is False + assert codec_args["compilation_config"] == { + "mode": 0, + "cudagraph_mode": "FULL", + "cudagraph_capture_sizes": [1, 2, 4, 8, 16, 32, 64], + "cudagraph_num_of_warmups": 1, + } + for platform in ("npu", "xpu", "rocm", "musa"): + # Platform resolution mutates the deployment; each server loads + # a fresh profile before applying its own platform overrides. + fallback = _apply_platform_overrides(load_deploy_config(candidate_path), platform=platform) + original = _apply_platform_overrides(load_deploy_config(base_path), platform=platform) + assert fallback.stages == original.stages + for variant in ("moss_tts_delay", "moss_tts_realtime"): + assert not any(stage.supports_native_mrv2_data_plane for stage in resolve_pipeline_config(variant).stages) + + def test_moss_local_high_concurrency_resolves_slot_compile_and_memory_budget(self): + pipeline = resolve_pipeline_config("moss_tts_local") + path = Path(get_deploy_config_path("moss_tts_local_mrv2_high_concurrency.yaml")) + deploy = _apply_platform_overrides(load_deploy_config(path), platform="cuda") + talker, codec = merge_pipeline_deploy(pipeline, deploy) + assert talker.yaml_engine_args["max_num_seqs"] == codec.yaml_engine_args["max_num_seqs"] == 128 + assert talker.yaml_engine_args["hf_overrides"]["mrv2_batch_prefill"] is True + assert talker.yaml_engine_args["hf_overrides"]["mrv2_direct_tokens"] is True + assert talker.yaml_engine_args["hf_overrides"]["mrv2_gpu_slot_state"] is True + assert talker.yaml_engine_args["max_num_batched_tokens"] == 512 + assert talker.yaml_engine_args["kv_cache_memory_bytes"] == 32 * 1024**3 + assert codec.yaml_engine_args["hf_overrides"]["codec_attention_backend"] == "triton_slot" + compilation = codec.yaml_engine_args["compilation_config"] + assert compilation["mode"] == 3 and compilation["backend"] == "inductor" + assert compilation["cudagraph_mode"] == "FULL" + assert compilation["cudagraph_capture_sizes"] == [1, 2, 4, 6, 8, 12, 16, 24, 32, 64, 128] + assert compilation["inductor_compile_config"] == {"combo_kernels": False, "benchmark_combo_kernel": False} + base_path = Path(get_deploy_config_path("moss_tts_local.yaml")) + for platform in ("npu", "xpu", "rocm", "musa"): + actual = _apply_platform_overrides(load_deploy_config(path), platform=platform) + expected = _apply_platform_overrides(load_deploy_config(base_path), platform=platform) + assert actual.stages == expected.stages + @pytest.mark.parametrize("runner,native", [("v1", False), ("v2", False), ("v2", True)]) def test_mrv2_undeclared_transport_warns(self, monkeypatch, runner, native): from unittest.mock import Mock diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index eb3db91ec..58e2a31f1 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: 32, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29, - EnvironmentVariableCategory.MODEL_SPECIFIC: 77, + EnvironmentVariableCategory.MODEL_SPECIFIC: 95, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21, EnvironmentVariableCategory.INTERNAL: 7, } @@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC ) assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == { - ModelEnvironmentVariableDisposition.PROMOTE: 48, + ModelEnvironmentVariableDisposition.PROMOTE: 66, ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5, ModelEnvironmentVariableDisposition.EXTERNAL: 0, ModelEnvironmentVariableDisposition.INTERNALIZE: 17, diff --git a/tests/core/sched/test_generation_batch_coalescing.py b/tests/core/sched/test_generation_batch_coalescing.py new file mode 100644 index 000000000..489e6a819 --- /dev/null +++ b/tests/core/sched/test_generation_batch_coalescing.py @@ -0,0 +1,253 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Bounded readiness coalescing must not allocate or consume request state.""" + +import queue +from types import SimpleNamespace + +import pytest +from vllm.v1.core.sched.interface import PauseState +from vllm.v1.request import RequestStatus + +from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler +from vllm_omni.outputs import OmniConnectorOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_idle_wait_wakes_on_input_without_extending_coalescing_deadline(monkeypatch): + events = [notification("0"), notification("1", "2", "3")] + s, clock = scheduler(monkeypatch, delayed=[(0.004, events[0]), (0.002, events[1])]) + s._generation_coalescing_policy = "idle_wait" + for request in s.requests.values(): + request.status = RequestStatus.WAITING_FOR_CHUNK + assert s._drain_omni_connector_outputs() == events + assert clock[0] == pytest.approx(10.006) + assert s._omni_connector_output_inbox.waits == pytest.approx([0.012, 0.012]) + + +@pytest.mark.parametrize("pending", ["in_flight", "admission", "registration", "terminal", "empty"]) +def test_idle_wait_preserves_completion_and_control_progress(monkeypatch, pending): + s, clock = scheduler(monkeypatch) + s._generation_coalescing_policy = "idle_wait" + for request in s.requests.values(): + request.status = RequestStatus.WAITING_FOR_CHUNK + if pending == "in_flight": + s.requests["0"].num_in_flight_tokens = 1 + elif pending == "admission": + s.requests["0"].status = RequestStatus.WAITING + elif pending == "registration": + s.input_coordinator.pending_chunk_registrations = [object()] + elif pending == "terminal": + s._pending_data_plane_terminal_req_ids = {"0"} + else: + s.requests.clear() + assert s._drain_omni_connector_outputs() == [] + assert clock[0] == 10 + assert not s._omni_connector_output_inbox.waits + + +def test_idle_wait_without_arrivals_is_bounded(monkeypatch): + s, clock = scheduler(monkeypatch) + s._generation_coalescing_policy = "idle_wait" + for request in s.requests.values(): + request.status = RequestStatus.WAITING_FOR_CHUNK + assert s._drain_omni_connector_outputs() == [] + assert clock[0] == pytest.approx(10.012) + + +def notification(*ids, terminal=()): + return OmniConnectorOutput(chunk_ready_req_ids=set(ids), chunk_finished_req_ids=set(terminal)) + + +class Inbox: + def __init__(self, clock, initial=(), delayed=()): + self.clock = clock + self.initial = list(initial) + self.delayed = list(delayed) + self.waits = [] + + def get_nowait(self): + if self.initial: + return self.initial.pop(0) + raise queue.Empty + + def get(self, timeout): + self.waits.append(timeout) + if not self.delayed or self.delayed[0][0] > timeout: + self.clock[0] += timeout + raise queue.Empty + delay, value = self.delayed.pop(0) + self.clock[0] += delay + return value + + +def scheduler(monkeypatch, initial=(), delayed=()): + clock = [10.0] + monkeypatch.setattr("vllm_omni.core.sched.omni_generation_scheduler.time.monotonic", lambda: clock[0]) + s = OmniGenerationScheduler.__new__(OmniGenerationScheduler) + s._generation_max_wait_s = 0.012 + s._generation_min_batch_size = 4 + s._pause_state = PauseState.UNPAUSED + s.requests = { + str(i): SimpleNamespace(request_id=str(i), num_in_flight_tokens=0, is_finished=lambda: False) for i in range(8) + } + s.running = list(s.requests.values()) + s.input_coordinator = SimpleNamespace( + requests_with_ready_chunks=set(), finished_requests=set(), input_terminal_req_ids=set() + ) + s._latest_omni_connector_output = None + s._omni_connector_output_inbox = Inbox(clock, initial, delayed) + return s, clock + + +def test_arrivals_form_target_batch_in_order(monkeypatch): + events = [notification("0"), notification("1", "2"), notification("3")] + s, clock = scheduler(monkeypatch, events[:1], [(0.002, events[1]), (0.003, events[2])]) + assert s._drain_omni_connector_outputs() == events + assert clock[0] == pytest.approx(10.005) + assert not s.input_coordinator.requests_with_ready_chunks + assert all(r.num_in_flight_tokens == 0 for r in s.requests.values()) + + +def test_duplicate_stale_and_finished_notifications_do_not_fill_batch(monkeypatch): + events = [notification("0", "gone", "1", "2"), notification("0")] + s, clock = scheduler(monkeypatch, events[:1], [(0.003, events[1])]) + s.requests["1"].is_finished = lambda: True + s.requests["2"].is_finished = lambda: True + assert s._drain_omni_connector_outputs() == events + assert clock[0] == pytest.approx(10.012) + assert s._omni_connector_output_inbox.waits == pytest.approx([0.012, 0.009]) + + +@pytest.mark.parametrize("reason", ["enough", "empty", "low", "paused", "off"]) +def test_immediate_dispatch(monkeypatch, reason): + event = notification("0") + s, clock = scheduler(monkeypatch, [event]) + if reason == "enough": + event.chunk_ready_req_ids.update(["1", "2", "3"]) + elif reason == "empty": + event.chunk_ready_req_ids.clear() + elif reason == "low": + s.requests = {"0": s.requests["0"]} + elif reason == "paused": + s._pause_state = PauseState.PAUSED_ALL + elif reason == "off": + s._generation_max_wait_s = 0 + assert s._drain_omni_connector_outputs() == [event] + assert not s._omni_connector_output_inbox.waits + assert clock[0] == 10 + + +def test_terminal_arrival_keeps_original_deadline(monkeypatch): + events = [notification("0"), notification("1", terminal=["1"])] + s, clock = scheduler(monkeypatch, events[:1], [(0.002, events[1])]) + assert s._drain_omni_connector_outputs() == events + assert clock[0] == pytest.approx(10.012) + + +def test_pending_ready_and_legacy_notification_preserved(monkeypatch): + s, _ = scheduler(monkeypatch) + s.input_coordinator.requests_with_ready_chunks.update(["0", "1", "2"]) + last = notification("3") + s._latest_omni_connector_output = last + assert s._drain_omni_connector_outputs() == [last] + assert s._latest_omni_connector_output is None + assert not s._omni_connector_output_inbox.waits + + +def test_deadline_is_new_for_each_batch_not_extended_by_arrivals(monkeypatch): + s, clock = scheduler(monkeypatch, [notification("0")], [(0.008, notification("1")), (0.008, notification("2"))]) + assert len(s._drain_omni_connector_outputs()) == 2 + assert clock[0] == pytest.approx(10.012) + s._omni_connector_output_inbox.initial.append(notification("3")) + assert len(s._drain_omni_connector_outputs()) == 2 + assert clock[0] == pytest.approx(10.024) + + +def test_cancelled_notification_is_filtered_before_coordinator(monkeypatch, mocker): + event = notification("0", "cancelled") + event.request_metadata = {"0": {"code_predictor_codes": [7]}, "cancelled": {"code_predictor_codes": [9]}} + s, _ = scheduler(monkeypatch, [event]) + s._generation_max_wait_s = 0 + s.waiting = [] + s.input_coordinator._async_chunk = True + s.input_coordinator.update_request_metadata = mocker.Mock() + s.input_coordinator.process_pending_chunks = mocker.Mock() + s._consume_pending_connector_output("generation") + s.input_coordinator.update_request_metadata.assert_called_once_with( + s.requests, {"0": {"code_predictor_codes": [7]}}, model_mode="generation" + ) + s.input_coordinator.process_pending_chunks.assert_called_once_with(s.waiting, s.running, {"0"}, set()) + + +@pytest.mark.parametrize("target,wait", [(0, 1), (257, 1), (4, -1), (4, float("nan")), (4, float("inf"))]) +def test_invalid_config_rejected(monkeypatch, target, wait): + from vllm_omni.core.sched.omni_generation_scheduler import VLLMScheduler + + def initialize(s, *args, **kwargs): + s.max_num_running_reqs = 256 + s.vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + retains_state_across_chunks=True, + stage_connector_config={"extra": {"generation_min_batch_size": target, "generation_max_wait_ms": wait}}, + ) + ) + + monkeypatch.setattr(VLLMScheduler, "__init__", initialize) + monkeypatch.setattr(OmniGenerationScheduler, "_init_omni_io_scheduling_state", lambda s: None) + with pytest.raises(ValueError): + OmniGenerationScheduler() + + +@pytest.mark.parametrize("terminal", [False, True]) +def test_first_and_terminal_chunks_have_bounded_wait(monkeypatch, terminal): + event = notification("0", terminal=["0"] if terminal else []) + s, clock = scheduler(monkeypatch, [event]) + s.running = s.running[1:] + assert s._drain_omni_connector_outputs() == [event] + assert clock[0] == pytest.approx(10.012) + + +def test_retire_in_flight_output_before_waiting_without_resetting_deadline(monkeypatch): + events = [notification("0"), notification("1")] + s, clock = scheduler(monkeypatch, events[:1]) + s.requests["7"].num_in_flight_tokens = 1 + assert s._drain_omni_connector_outputs() == events[:1] + assert s._generation_defer_batch + assert not s._omni_connector_output_inbox.waits + s.input_coordinator.requests_with_ready_chunks.add("0") + # Simulate EngineCore retiring the previous batch after five milliseconds. + clock[0] += 0.005 + s.requests["7"].num_in_flight_tokens = 0 + s._omni_connector_output_inbox.initial.append(events[1]) + assert s._drain_omni_connector_outputs() == events[1:] + assert not s._generation_defer_batch + assert s._omni_connector_output_inbox.waits == pytest.approx([0.007]) + assert clock[0] == pytest.approx(10.012) + assert s._generation_batch_deadline is None + + +def test_expired_deadline_dispatches_even_with_other_requests_in_flight(monkeypatch): + s, clock = scheduler(monkeypatch, [notification("0")]) + s.requests["7"].num_in_flight_tokens = 1 + s._generation_batch_deadline = clock[0] - 0.001 + assert len(s._drain_omni_connector_outputs()) == 1 + assert not s._generation_defer_batch + assert not s._omni_connector_output_inbox.waits + + +def test_deferred_schedule_keeps_pending_request_and_allocates_nothing(monkeypatch, mocker): + from tests.core.sched.test_generation_scheduler_restore import _chunk_request, _make_generation_scheduler + + request = _chunk_request("pending") + s = _make_generation_scheduler(request) + s._process_pending_omni_inputs = lambda **kwargs: setattr(s, "_generation_defer_batch", True) + allocate = mocker.spy(s.kv_cache_manager, "allocate_slots") + output = s.schedule() + assert output.total_num_scheduled_tokens == 0 + assert not output.num_scheduled_tokens + allocate.assert_not_called() + assert list(s.waiting) == [request] + assert request.num_computed_tokens == request.num_in_flight_tokens == 0 + assert s.chunk_transfer_adapter.restore_called diff --git a/tests/core/sched/test_generation_scheduler_init.py b/tests/core/sched/test_generation_scheduler_init.py new file mode 100644 index 000000000..ade94af17 --- /dev/null +++ b/tests/core/sched/test_generation_scheduler_init.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Exercise generation constructor policy after actual deploy/platform resolution.""" + +from types import SimpleNamespace + +import pytest + +from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.config.pipeline_registry import resolve_pipeline_config +from vllm_omni.config.stage_config import _apply_platform_overrides, load_deploy_config, merge_pipeline_deploy +from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler, VLLMScheduler + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def construct_scheduler(monkeypatch, mocker): + # Only upstream cache allocation and connector construction are omitted. + # Keep the actual Omni I/O initialization, native decision and constructor. + mocker.patch("vllm_omni.core.sched.omni_scheduler_mixin.OmniChunkTransferAdapter") + + def upstream_init(self, config): + self.vllm_config = config + self.max_num_running_reqs = config.scheduler_config.max_num_seqs + + monkeypatch.setattr(VLLMScheduler, "__init__", upstream_init) + return OmniGenerationScheduler + + +def _config(*, native=True, stateful=True, tp=1, pp=1, extras=None, capacity=128): + return SimpleNamespace( + model_config=SimpleNamespace( + use_v2_model_runner=native, + supports_native_mrv2_data_plane=True, + retains_state_across_chunks=stateful, + async_chunk=True, + stage_id=1, + stage_connector_config={"extra": extras or {}}, + ), + parallel_config=SimpleNamespace(tensor_parallel_size=tp, pipeline_parallel_size=pp), + scheduler_config=SimpleNamespace(max_num_seqs=capacity), + ) + + +@pytest.mark.parametrize("profile", ["high_concurrency", "low_latency"]) +@pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"]) +def test_moss_profile_generation_constructor_after_platform_resolution(construct_scheduler, mocker, profile, platform): + deploy = _apply_platform_overrides( + load_deploy_config(get_deploy_config_path(f"moss_tts_local_mrv2_{profile}.yaml")), platform=platform + ) + pipeline = resolve_pipeline_config("moss_tts_local") + codec = merge_pipeline_deploy(pipeline, deploy)[1] + extras = deploy.connectors["shm"]["extra"] + config = _config( + native=codec.yaml_engine_args["use_v2_model_runner"], + stateful=pipeline.stages[1].retains_state_across_chunks, + capacity=codec.yaml_engine_args["max_num_seqs"], + extras=extras, + ) + warning = mocker.patch("vllm_omni.core.sched.omni_generation_scheduler.logger.warning") + scheduler = construct_scheduler(config) + # The raw extras persist through fallback; the live constructor must cope. + assert extras["generation_min_batch_size"] == 32 and extras["generation_max_wait_ms"] == 12 + if platform == "cuda": + assert scheduler._native_data_plane and scheduler.input_coordinator is not None + assert scheduler._generation_min_batch_size == 32 + assert scheduler._generation_max_wait_s == 0.012 + assert scheduler._generation_max_regular_batch == (16 if profile == "low_latency" else 0) + warning.assert_not_called() + if profile == "low_latency": + assert extras["codec_first_chunk_fast_path"] == extras["codec_first_chunk_gate"] == 1 + else: + assert not scheduler._native_data_plane and scheduler.chunk_transfer_adapter is not None + assert scheduler._generation_min_batch_size == 1 + assert scheduler._generation_max_wait_s == scheduler._generation_max_regular_batch == 0 + warning.assert_called_once() + + +@pytest.mark.parametrize("stateful,tp,pp", [(False, 1, 1), (True, 2, 1), (True, 1, 2)]) +def test_native_batch_waiting_still_rejects_unsupported_state_or_topology(construct_scheduler, stateful, tp, pp): + config = _config( + stateful=stateful, tp=tp, pp=pp, extras={"generation_min_batch_size": 32, "generation_max_wait_ms": 12} + ) + with pytest.raises(ValueError, match="stateful native MRV2 TP1/PP1"): + construct_scheduler(config) + + +@pytest.mark.parametrize( + "extras,message", + [ + ({"generation_min_batch_size": 0}, "generation_min_batch_size"), + ({"generation_min_batch_size": 129}, "generation_min_batch_size"), + ({"generation_max_wait_ms": -1}, "generation_max_wait_ms"), + ({"generation_max_wait_ms": float("nan")}, "generation_max_wait_ms"), + ({"generation_max_regular_batch": -1}, "generation_max_regular_batch"), + ({"generation_coalescing_policy": "unknown"}, "generation_coalescing_policy"), + ], +) +def test_non_native_fallback_still_rejects_invalid_values(construct_scheduler, extras, message): + with pytest.raises(ValueError, match=message): + construct_scheduler(_config(native=False, extras=extras)) diff --git a/tests/core/sched/test_omni_sched_deferred_free_fence.py b/tests/core/sched/test_omni_sched_deferred_free_fence.py index 9088c6bc2..c353ced3f 100644 --- a/tests/core/sched/test_omni_sched_deferred_free_fence.py +++ b/tests/core/sched/test_omni_sched_deferred_free_fence.py @@ -126,6 +126,11 @@ def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request sched.max_num_running_reqs = 8 sched._native_data_plane = False sched._retains_state_across_chunks = False + # Unset MagicMock attributes are truthy. Keep optional scheduling modes + # disabled so this fixture actually schedules its one-shot request. + sched._generation_defer_batch = False + sched._generation_first_chunk_only = False + sched._generation_max_regular_batch = 0 sched._pause_state = PauseState.UNPAUSED sched.requests = {request.request_id: request} sched.running = [request] diff --git a/tests/engine/test_stage_engine_core_proc.py b/tests/engine/test_stage_engine_core_proc.py index 217f13e53..19c24b306 100644 --- a/tests/engine/test_stage_engine_core_proc.py +++ b/tests/engine/test_stage_engine_core_proc.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace +from typing import Any from unittest.mock import patch import pytest @@ -25,11 +26,15 @@ def test_first_audio_sink_preserves_first_decoder_and_opts_in_stream_decoder( expected, ): executor = UniProcExecutor.__new__(UniProcExecutor) + executor.vllm_config = SimpleNamespace( + parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1) + ) executor.driver_worker = mocker.Mock() runner = executor.driver_worker.worker.model_runner runner.model.first_frame_decoder = object() if first_decoder else None runner.model.stream_decoder = object() if stream_decoder else None runner.model.stream_first_audio = stream_first_audio + runner._omni_data_plane = None assert _bind_first_audio_sink(executor, mocker.Mock(), mocker.Mock()) is expected assert runner.model_state.set_first_audio_sink.call_count == int(expected) @@ -54,3 +59,64 @@ def test_preprocess_add_request_preserves_omni_fields(): assert current_wave == 3 assert result.external_req_id == "external" assert result.additional_information == {"conditioning": "payload"} + + +def test_codec_uses_generic_first_audio_binding(monkeypatch): + import queue + + import torch + + from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY + from vllm_omni.engine import stage_engine_core_proc as module + + calls: dict[str, Any] = {} + + def hook(*args): + return False + + def bind(sink): + calls["sink"] = sink + return hook + + plane = SimpleNamespace(set_first_chunk_hook=lambda value: calls.update(hook=value)) + model = SimpleNamespace(bind_first_chunk_fast_path=bind) + runner = SimpleNamespace(model=model, get_model=lambda: model, _omni_data_plane=plane) + executor = SimpleNamespace( + vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)), + driver_worker=SimpleNamespace(worker=SimpleNamespace(model_runner=runner)), + ) + monkeypatch.setattr(module, "UniProcExecutor", SimpleNamespace) + outputs: queue.Queue = queue.Queue() + scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=2)}) + assert module._bind_first_audio_sink(executor, outputs, scheduler) + assert calls["hook"] is hook + calls["sink"].prepare(["r"])(["r"], [torch.ones(2)], torch.tensor(24000)) + client, batch = outputs.get_nowait() + assert client == 2 and FIRST_AUDIO_KEY not in batch.outputs[0].multimodal_output + + +def test_first_audio_binding_preserves_talker_marker(monkeypatch): + import queue + + import torch + + from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY + from vllm_omni.engine import stage_engine_core_proc as module + + calls: dict[str, Any] = {} + runner = SimpleNamespace( + model=SimpleNamespace(first_frame_decoder=object()), + model_state=SimpleNamespace(set_first_audio_sink=lambda sink: calls.update(sink=sink)), + ) + executor = SimpleNamespace( + vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=1, pipeline_parallel_size=1)), + driver_worker=SimpleNamespace(worker=SimpleNamespace(model_runner=runner)), + ) + monkeypatch.setattr(module, "UniProcExecutor", SimpleNamespace) + outputs: queue.Queue = queue.Queue() + scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=0)}) + assert module._bind_first_audio_sink(executor, outputs, scheduler) + calls["sink"].prepare(["r"])(["r"], [torch.ones(2)], torch.tensor(24000)) + assert outputs.get_nowait()[1].outputs[0].multimodal_output[FIRST_AUDIO_KEY] + executor.vllm_config.parallel_config.tensor_parallel_size = 2 + assert not module._bind_first_audio_sink(executor, outputs, scheduler) diff --git a/tests/entrypoints/openai_api/test_moss_local_prompt.py b/tests/entrypoints/openai_api/test_moss_local_prompt.py new file mode 100644 index 000000000..8fb4c62a8 --- /dev/null +++ b/tests/entrypoints/openai_api/test_moss_local_prompt.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The Local-v1.5 continuation prompt fast path equals the processor layout.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.entrypoints.openai.tts_adapters.moss_tts import _local_continuation_prompt + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Processor: + """Segment API of the model's processor with a deterministic tokenizer.""" + + def __init__(self): + self.model_config = SimpleNamespace( + n_vq=3, + im_start_token_id=1, + im_end_token_id=2, + audio_start_token_id=3, + audio_pad_token_id=1024, + audio_assistant_slot_token_id=5, + ) + self.encoded = [] + + def _encode_text(self, text): + self.encoded.append(text) + return [100 + ord(c) % 50 for c in text] + + def _user_prompt_prefix_ids(self): + return [1] + self._encode_text("user\n") + self._encode_text("") + + def _user_prompt_after_reference_ids(self, language, fields): + return self._encode_text(f"lang={fields.get('language', language)}") + + def _assistant_prompt_prefix_ids(self): + return self._encode_text("") + [2] + self._encode_text("\n") + [1] + self._encode_text("assistant") + + def reference(self, text, language, codes): + # Mirrors the processor's _build_continuation_codes + _pad (one item). + fields = {} if language is None else {"language": language} + ids = ( + self._user_prompt_prefix_ids() + + self._encode_text("None") + + self._user_prompt_after_reference_ids(language, fields) + + self._encode_text(text) + + self._assistant_prompt_prefix_ids() + + [3] + ) + text_rows = torch.full((len(ids), 4), 1024, dtype=torch.long) + text_rows[:, 0] = torch.tensor(ids) + audio_rows = torch.full((codes.shape[0], 4), 1024, dtype=torch.long) + audio_rows[:, 0] = 5 + audio_rows[:, 1:] = codes + unified = torch.cat([text_rows, audio_rows]) + return unified[:, 0].tolist(), unified[:, 1:].contiguous() + + +@pytest.mark.parametrize("language", [None, "English"]) +def test_matches_processor_and_caches_constant_segments(language): + proc = _Processor() + codes = torch.randint(0, 1024, (7, 3), dtype=torch.int32) + kwargs = {"text": "ref words target words"} + if language is not None: + kwargs["language"] = language + ids, audio = _local_continuation_prompt(proc, kwargs, [codes]) + expected_ids, expected_audio = proc.reference(kwargs["text"], language, codes.long()) + assert ids == expected_ids + assert audio.dtype == torch.int64 and torch.equal(audio, expected_audio) + + proc.encoded.clear() + _local_continuation_prompt(proc, {**kwargs, "text": "other"}, [codes]) + assert proc.encoded == ["other"] + + +def test_unsupported_layouts_fall_back(): + proc = _Processor() + codes = torch.zeros((2, 3), dtype=torch.long) + assert _local_continuation_prompt(proc, {"text": "t", "instruction": "x"}, [codes]) is None + assert _local_continuation_prompt(proc, {"text": "t"}, [codes, codes]) is None + assert _local_continuation_prompt(proc, {"text": "t"}, ["path.wav"]) is None + assert _local_continuation_prompt(proc, {"text": "t"}, [torch.zeros((2, 4), dtype=torch.long)]) is None diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py index 92f13462e..f17fdbb65 100644 --- a/tests/entrypoints/openai_api/test_tts_adapter.py +++ b/tests/entrypoints/openai_api/test_tts_adapter.py @@ -691,6 +691,45 @@ def test_moss_get_processor_places_audio_tokenizer_on_codec_stage_device(mocker) audio_tokenizer.to.assert_called_once() +@pytest.mark.parametrize( + "variant,graphs_env,prepares", + [("local", None, False), ("local", "1", True), ("local", "0", False), ("realtime", "1", False)], +) +def test_moss_warmup_prepares_the_reference_encoder(variant, graphs_env, prepares, mocker, monkeypatch): + if graphs_env is None: + monkeypatch.delenv("VLLM_OMNI_MOSS_REF_GRAPHS", raising=False) + else: + monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", graphs_env) + adapter = _moss_adapter_with_stage_configs([], mocker) + adapter._moss_variant = variant + encoder = mocker.Mock() + mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder) + asyncio.run(adapter.warmup()) + assert encoder.prepare.called is prepares + + +def test_moss_warmup_failure_does_not_stop_startup(mocker, monkeypatch): + monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", "1") + adapter = _moss_adapter_with_stage_configs([], mocker) + encoder = mocker.Mock() + encoder.prepare.side_effect = RuntimeError("capture failed") + mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder) + asyncio.run(adapter.warmup()) + encoder.prepare.assert_called_once_with() + + +def test_moss_shared_encoder_startup_failure_stops_api_startup(mocker, monkeypatch): + monkeypatch.setenv("VLLM_OMNI_MOSS_REF_GRAPHS", "1") + from vllm_omni.model_executor.models.moss_tts.shared_reference_encoder import SharedReferenceEncoderStartupError + + adapter = _moss_adapter_with_stage_configs([], mocker) + encoder = mocker.Mock() + encoder.prepare.side_effect = SharedReferenceEncoderStartupError("still compiling") + mocker.patch.object(adapter, "_get_moss_ref_encoder", return_value=encoder) + with pytest.raises(SharedReferenceEncoderStartupError): + asyncio.run(adapter.warmup()) + + def test_qwen3_tts_metadata(): assert Qwen3TTSAdapter.backend == "ar" assert issubclass(Qwen3TTSAdapter, ARTTSAdapter) diff --git a/tests/entrypoints/test_pd_disaggregation.py b/tests/entrypoints/test_pd_disaggregation.py index e9a3ba54c..bf56da2b3 100644 --- a/tests/entrypoints/test_pd_disaggregation.py +++ b/tests/entrypoints/test_pd_disaggregation.py @@ -288,6 +288,9 @@ def _setup_log_mocks(monkeypatch): def build_and_log_summary(self): return "Fake summary" + def log_timing_summary(self): + pass + monkeypatch.setattr( "vllm_omni.entrypoints.omni.OrchestratorAggregator", _FakeOrchestratorAggregator, diff --git a/tests/metrics/test_emit_calls.py b/tests/metrics/test_emit_calls.py index 02a98fc00..515aaeb8d 100644 --- a/tests/metrics/test_emit_calls.py +++ b/tests/metrics/test_emit_calls.py @@ -111,7 +111,7 @@ class TestFailureCounterWiring: obj.request_states["req-1"] = SimpleNamespace( metrics=SimpleNamespace( e2e_done=set(), - build_and_log_summary=lambda: None, + log_timing_summary=lambda: None, ), failure_recorded=False, ) @@ -127,7 +127,7 @@ class TestFailureCounterWiring: obj.request_states["req-1"] = SimpleNamespace( metrics=SimpleNamespace( e2e_done=set(), - build_and_log_summary=lambda: None, + log_timing_summary=lambda: None, ), failure_recorded=False, ) diff --git a/tests/metrics/test_stats.py b/tests/metrics/test_stats.py index 8b18c5dff..f72cd0867 100644 --- a/tests/metrics/test_stats.py +++ b/tests/metrics/test_stats.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import pytest @@ -155,3 +158,39 @@ def test_build_and_log_summary_multiple_requests() -> None: r2_stage_entry = next(e for e in summary["stage_table"] if e["request_id"] == "r2") assert len(r1_stage_entry["stages"]) == 2 assert len(r2_stage_entry["stages"]) == 1 + + +def test_log_timing_summary_skips_tables_unless_debug(monkeypatch) -> None: + from vllm_omni.metrics import stats + + agg = OrchestratorAggregator(num_stages=1, log_stats=True, wall_start_ts=0.0, final_stage_id_for_e2e=0) + agg.stage_first_ts[0] = 0.0 + agg.stage_last_ts[0] = 0.02 + agg.on_stage_metrics( + 0, + "r1", + StageRequestStats( + batch_id=1, + batch_size=1, + num_tokens_in=3, + num_tokens_out=3, + stage_gen_time_ms=20.0, + rx_transfer_bytes=0, + rx_decode_time_ms=0.0, + rx_in_flight_time_ms=0.0, + stage_stats=StageStats(), + ), + ) + agg.on_finalize_request(0, "r1", req_start_ts=0.0) + tables, lines = [], [] + + def record_table(*args, **kwargs): + tables.append(args) + return "" + + monkeypatch.setattr(stats, "_format_table", record_table) + monkeypatch.setattr(stats.logger, "info", lambda msg, *args: lines.append(msg % args)) + monkeypatch.setattr(stats.logger, "isEnabledFor", lambda level: False) + agg.log_timing_summary() + assert not tables + assert len(lines) == 1 and lines[0].startswith("[OmniTiming] req=r1 total=") diff --git a/tests/model_executor/models/moss_tts/test_codec_async_output.py b/tests/model_executor/models/moss_tts/test_codec_async_output.py new file mode 100644 index 000000000..d855ef9d3 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_codec_async_output.py @@ -0,0 +1,338 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Waveform ownership across graph replays and asynchronous output copies.""" + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import MossTTSCodecDecoder +from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot + +pytestmark = [pytest.mark.core_model] + + +class CausalCodec(nn.Module): + downsample_rate = 1 + + def __init__(self, channels): + super().__init__() + self.weight = nn.Parameter(torch.zeros(())) + self.channels = channels + self.config = SimpleNamespace(codebook_size=1024) + + def initialize_decoder_state_pool(self, capacity, scratch): + self.state = torch.zeros(capacity + scratch, device=self.weight.device) + + def reset_decoder_state_slots(self, slots): + self.state[slots] = 0 + + def decode_streaming_batch(self, codes, lengths, slots, valid_rows): + audio = codes.sum(0).float().cumsum(-1) + self.state[slots, None] + self.state[slots] = audio[:, -1] + channels = torch.arange(1, self.channels + 1, device=audio.device) + return SimpleNamespace(audio=audio[:, None] * channels[None, :, None], audio_lengths=lengths) + + +class ReplayGraph: + """A real CUDA graph with deliberately reused output storage.""" + + batch_sizes = [1] + + def __init__(self, codec): + self.codes = torch.ones((2, 1, 3), device="cuda", dtype=torch.long) + self.slots = torch.zeros(1, device="cuda", dtype=torch.long) + self.lengths = torch.full((1,), 3, device="cuda", dtype=torch.long) + self.valid = torch.ones(1, device="cuda", dtype=torch.bool) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + codec.decode_streaming_batch(self.codes, self.lengths, self.slots, self.valid) + torch.cuda.current_stream().wait_stream(stream) + self.graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(self.graph): + self.output = codec.decode_streaming_batch(self.codes, self.lengths, self.slots, self.valid) + codec.state.zero_() + + def decode(self, codes, slots, *, allow_frame_padding): + if codes.shape[1:] != (1, 3): + return None + self.codes.copy_(codes) + self.slots.copy_(slots) + self.graph.replay() + return self.output.audio, self.output.audio_lengths, 1 + + +def make_decoder(*, v2=True, channels=2, device="cpu", model_type="moss_tts_local", streaming=True): + config = SimpleNamespace( + model_config=SimpleNamespace( + hf_config=SimpleNamespace(n_vq=2, model_type=model_type), + async_chunk=streaming, + use_v2_model_runner=v2, + ), + scheduler_config=SimpleNamespace(max_num_seqs=8), + ) + decoder = MossTTSCodecDecoder(vllm_config=config) + decoder._codec = CausalCodec(channels).to(device) + decoder._n_channels = channels + decoder._stream_max_step_frames = 3 + return decoder + + +@pytest.mark.cpu +@pytest.mark.parametrize("ramp,frames", [(None, [1, 15]), ([1, 4, 8, 15], [1, 4, 8, 15]), ([1, 4, 4, 15], [1, 4, 15])]) +def test_ramp_chunks_are_captured_at_exact_frame_counts(ramp, frames): + config = SimpleNamespace( + model_config=SimpleNamespace( + hf_config=SimpleNamespace(n_vq=2, model_type="moss_tts_local"), + async_chunk=True, + use_v2_model_runner=True, + stage_connector_config={ + "extra": { + "initial_codec_chunk_frames": 1, + "codec_chunk_frames": 15, + "codec_chunk_ramp": ramp, + } + }, + ), + scheduler_config=SimpleNamespace(max_num_seqs=8), + ) + decoder = MossTTSCodecDecoder(vllm_config=config) + assert decoder._streaming_graph_frame_sizes == frames + assert decoder._stream_max_step_frames == 15 + + +@pytest.mark.cpu +@pytest.mark.parametrize( + "v2,streaming,model_type", + [ + (False, True, "moss_tts_local"), + (True, False, "moss_tts_local"), + (True, True, "moss_tts_realtime"), + (True, True, "moss_tts_local"), + ], +) +def test_gpu_output_requires_local_streaming_mrv2_and_cuda(v2, streaming, model_type): + decoder = make_decoder(v2=v2, streaming=streaming, model_type=model_type) + assert decoder._gpu_stream_output == (v2 and streaming and model_type == "moss_tts_local") + session = decoder._ensure_stream_session() + slot = session.acquire() + assert session.step({slot: torch.ones(2, 3)})[slot].device.type == "cpu" + + +@pytest.mark.cuda +@pytest.mark.parametrize("v2", [False, True]) +@pytest.mark.parametrize("channels", [1, 2]) +@pytest.mark.parametrize("graph", [False, True]) +@torch.no_grad() +def test_replays_terminal_tails_and_slot_reuse_preserve_pending_waveforms(v2, channels, graph): + from vllm_omni.outputs import OmniModelRunnerOutput + from vllm_omni.worker_v2.omni_generation_model_runner import ( + OmniGenerationAsyncOutput, + OmniGenerationModelRunner, + _contains_cuda_tensor, + ) + + decoder = make_decoder(v2=v2, channels=channels, device="cuda") + # Real Local-v1.5 weights replace the class's 24 kHz default at load time. + decoder._sr_tensor = torch.tensor(48_000, dtype=torch.int32) + session = decoder._ensure_stream_session() + if graph: + session._cudagraph_wrapper = ReplayGraph(decoder._codec) + copy_stream = torch.cuda.Stream() + pending = [] + # Do not materialize any output until later forwards have replayed the + # graph, reset terminal/cancelled state, and reused those state slots. + for round_index in range(4): + value = round_index + 1 + codes = torch.cat( + [ + torch.full((14,), value, device="cuda", dtype=torch.long), + torch.full((6,), value + 10, device="cuda", dtype=torch.long), + ] + ) + info = [{"request_id": rid, "meta": {"finished": rid == "tail"}} for rid in ["tail", "empty", "live"]] + output = decoder.forward(codes, runtime_additional_information=info, seq_token_counts=[14, 0, 6]) + # The runner may overwrite its input buffer before codec completion. + codes.fill_(-100) + payload = output.multimodal_outputs + assert _contains_cuda_tensor(payload) == v2 + assert isinstance(payload, PackedOutputSnapshot) == v2 + if v2: + # Exactly one waveform slab, regardless of request/chunk count. + assert sum(slab.is_cuda for _, slab in payload._groups) == 1 + result = OmniModelRunnerOutput( + req_ids=["tail", "empty", "live"], + req_id_to_index={"tail": 0, "empty": 1, "live": 2}, + sampled_token_ids=[[], [], []], + ) + pending.append( + OmniGenerationAsyncOutput( + model_runner_output=result, + multimodal_outputs=payload, + num_reqs=3, + main_stream=torch.cuda.current_stream(), + copy_stream=copy_stream, + ) + ) + else: + pending.append(OmniGenerationModelRunner._build_pooler_output(output, 3)) + decoder.on_requests_finished(["live"]) + assert decoder._stream_req_slots == {} + assert not session._leased_slots + + for round_index, item in enumerate(pending): + rows = item.get_output().multimodal_outputs if v2 else item + for index, (frames, value) in enumerate([(7, round_index + 1), (0, 0), (3, round_index + 11)]): + expected = 2 * value * torch.arange(1, frames + 1, dtype=torch.float32) + if channels > 1: + expected = expected[None] * torch.arange(1, channels + 1)[:, None] + torch.testing.assert_close(rows[index]["model_outputs"], expected, rtol=0, atol=0) + assert rows[index]["sr"].item() == 48_000 + + +@pytest.mark.cuda +@torch.no_grad() +def test_owned_input_copy_clamps_entire_ragged_batch(): + from vllm_omni.worker_v2.omni_ar_model_runner import _async_copy_mm + + decoder = make_decoder(device="cuda", channels=2) + codes = torch.tensor([-1, 1024, 1, -10, 2, 1025, 5000, -2], device="cuda") + output = decoder.forward( + codes, + runtime_additional_information=[{"request_id": name} for name in ("a", "empty", "b")], + seq_token_counts=[6, 0, 2], + ) + codes.fill_(-100) + cpu = _async_copy_mm(output.multimodal_outputs, 0, pin_memory=True) + torch.accelerator.synchronize() + expected = torch.tensor([[0.0, 1025.0, 2049.0], [0.0, 2050.0, 4098.0]]) + torch.testing.assert_close(cpu["model_outputs"][0], expected, rtol=0, atol=0) + assert cpu["model_outputs"][1].numel() == 0 + torch.testing.assert_close(cpu["model_outputs"][2], torch.tensor([[1023.0], [2046.0]]), rtol=0, atol=0) + decoder.on_requests_finished(["a", "b"]) + + +@pytest.mark.cuda +@torch.no_grad() +def test_pinned_metadata_reuse_waits_for_dma_and_preserves_mixed_terminals(): + decoder = make_decoder(device="cuda") + session = decoder._ensure_stream_session() + # Force reuse every call, with an outstanding transfer on the first call. + # This delay tests host-buffer ownership, never performance or overlap. + session._metadata_ring = session._metadata_ring[:1] + slots = [session.acquire() for _ in range(3)] + pending = [] + for step in range(8): + order = slots[step % 3 :] + slots[: step % 3] + terminal = {order[1]} + torch.cuda._sleep(2_000_000) + ids, tails = session._stage_slot_ids(order, terminal) + pending.append((ids.clone(), tails.clone(), order, list(terminal))) + session.close() + for ids, tails, expected_ids, expected_tails in pending: + assert ids.cpu().tolist() == expected_ids + assert tails.cpu().tolist() == expected_tails + + +@pytest.mark.cuda +@torch.no_grad() +def test_mixed_terminal_reset_does_not_reset_live_slot(): + decoder = make_decoder(device="cuda", channels=1) + session = decoder._ensure_stream_session() + first, second = session.acquire(), session.acquire() + codes = torch.ones(2, 3, device="cuda") + initial = session.step({first: codes, second: 2 * codes}, terminal_slots={first}) + following = session.step({second: codes, first: codes}) + torch.testing.assert_close(initial[first].cpu(), torch.tensor([[2.0, 4.0, 6.0]]), rtol=0, atol=0) + torch.testing.assert_close(following[first].cpu(), torch.tensor([[2.0, 4.0, 6.0]]), rtol=0, atol=0) + torch.testing.assert_close(following[second].cpu(), torch.tensor([[14.0, 16.0, 18.0]]), rtol=0, atol=0) + session.close() + + +@pytest.mark.cuda +@pytest.mark.parametrize("explicit_copy_stream", [False, True]) +@torch.no_grad() +def test_output_copy_waits_for_private_codec_stream(explicit_copy_stream): + from vllm_omni.worker_v2.omni_ar_model_runner import _async_copy_mm + + decoder = make_decoder(device="cuda", channels=1) + decoder._codec_stream = torch.cuda.Stream() + with torch.cuda.stream(decoder._codec_stream): + torch.cuda._sleep(20_000_000) + codes = torch.ones(6, device="cuda", dtype=torch.long) + output = decoder.forward( + codes, + runtime_additional_information=[{"request_id": "live"}], + seq_token_counts=[6], + ) + codes.fill_(-100) + payload = output.multimodal_outputs + assert isinstance(payload, PackedOutputSnapshot) + assert payload.producer_event is not None + # No explicit consumer wait here: the shared copier must honor readiness. + copy_stream = torch.cuda.Stream() + with torch.cuda.stream(copy_stream): + cpu = _async_copy_mm(payload, 0, copy_stream=copy_stream if explicit_copy_stream else None, pin_memory=True) + decoder.on_requests_finished(["live"]) + copy_stream.synchronize() + torch.testing.assert_close(cpu["model_outputs"][0], torch.tensor([2.0, 4.0, 6.0]), rtol=0, atol=0) + + +@pytest.mark.cuda +@pytest.mark.parametrize("channels", [1, 2]) +@torch.no_grad() +def test_short_terminal_row_crops_each_channel_into_output_slab(channels): + decoder = make_decoder(device="cuda", channels=channels) + session = decoder._ensure_stream_session() + short, live = session.acquire(), session.acquire() + codes = {short: torch.ones(2, 2, device="cuda"), live: 3 * torch.ones(2, 5, device="cuda")} + slab = torch.empty(channels * 7, device="cuda") + buffers = { + short: slab[: channels * 2].view(channels, 2), + live: slab[channels * 2 :].view(channels, 5), + } + result = session.step(codes, terminal_slots={short}, output_buffers=buffers) + for slot, frames, value in [(short, 2, 1), (live, 5, 3)]: + expected = 2 * value * torch.arange(1, frames + 1, dtype=torch.float32)[None] + expected = expected * torch.arange(1, channels + 1)[:, None] + assert result[slot].data_ptr() == buffers[slot].data_ptr() + torch.testing.assert_close(result[slot].cpu(), expected, rtol=0, atol=0) + session.close() + + +@pytest.mark.cuda +@pytest.mark.parametrize("channels", [1, 2]) +@pytest.mark.parametrize("layout", ["adjacent", "strided", "gapped"]) +@torch.no_grad() +def test_batched_output_copy_preserves_rows_and_slab_ownership(channels, layout): + decoder = make_decoder(device="cuda", channels=channels) + session = decoder._ensure_stream_session() + slots = [session.acquire() for _ in range(4)] + slab = torch.full((6 if layout == "gapped" else 4, channels, 6), -1.0, device="cuda") + if layout == "gapped": + buffers = {slot: slab[row] for row, slot in zip([0, 1, 3, 4], slots, strict=True)} + frames = 6 + elif layout == "adjacent": + buffers = {slot: slab.reshape(4, channels * 6)[row].view(channels, 6) for row, slot in enumerate(slots)} + frames = 6 + else: + # Interleaved with other requests' output: the short views must not + # overwrite the untouched part of another row/channel. + buffers = {slot: slab[row, :, :3] for row, slot in enumerate(slots)} + frames = 3 + codes = {slot: torch.full((2, frames), row + 1, device="cuda") for row, slot in enumerate(slots)} + result = session.step(codes, output_buffers=buffers) + for row, slot in enumerate(slots): + expected = 2 * (row + 1) * torch.arange(1, frames + 1, dtype=torch.float32)[None] + expected = expected * torch.arange(1, channels + 1)[:, None] + assert result[slot].data_ptr() == buffers[slot].data_ptr() + torch.testing.assert_close(result[slot].cpu(), expected, rtol=0, atol=0) + if layout == "strided": + assert torch.all(slab[:, :, 3:] == -1) + elif layout == "gapped": + assert torch.all(slab[[2, 5]] == -1) + session.close() diff --git a/tests/model_executor/models/moss_tts/test_codec_ring_headroom.py b/tests/model_executor/models/moss_tts/test_codec_ring_headroom.py new file mode 100644 index 000000000..aa3fb1797 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_codec_ring_headroom.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Chunked streaming decode must match whole-sequence decode. + +``RingKVCache.complete`` writes a whole chunk before attending. With a ring of +exactly ``context`` entries, a chunk of T tokens evicts up to T keys that its +own queries should still see; for T > context it evicts part of the chunk +itself. ``initialize_decoder_state_pool(chunk_frames=...)`` adds one chunk of +headroom per layer so streaming reproduces the non-streaming decoder. +""" + +import pytest +import torch + +from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import MossAudioTokenizerModel +from vllm_omni.model_executor.models.moss_tts.configuration_moss_audio_tokenizer_v2 import ( + MossAudioTokenizerConfig, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +TRANSFORMER = { + "module_type": "Transformer", + "d_model": 16, + "num_heads": 2, + "num_layers": 1, + "dim_feedforward": 32, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10000, + "gating": "none", + "layer_scale": 1.0, + "conv_layout": True, +} + + +def tiny_codec() -> MossAudioTokenizerModel: + # 64 Hz mono input, 4-sample code frames (16 code frames per second). + # Decoder: transformer @16 Hz (context 0.5 s = 8 tokens) -> x2 -> transformer + # @32 Hz (context 0.25 s = 8 tokens) -> x2 -> 64 Hz. A 6-frame chunk is 12 + # tokens at the second transformer, more than its 8-token ring. + config = MossAudioTokenizerConfig( + sampling_rate=64, + downsample_rate=4, + number_channels=1, + enable_channel_interleave=False, + encoder_kwargs=[ + {"module_type": "PatchedPretransform", "patch_size": 4}, + {**TRANSFORMER, "input_dimension": 4, "output_dimension": 8, "context_duration": 0.5}, + ], + decoder_kwargs=[ + {**TRANSFORMER, "input_dimension": 8, "output_dimension": 16, "context_duration": 0.5}, + {"module_type": "PatchedPretransform", "patch_size": 2}, + {**TRANSFORMER, "input_dimension": 8, "output_dimension": 2, "context_duration": 0.25}, + {"module_type": "PatchedPretransform", "patch_size": 2}, + ], + quantizer_kwargs={ + "input_dim": 8, + "rvq_dim": 8, + "output_dim": 8, + "num_quantizers": 2, + "codebook_size": 16, + "codebook_dim": 4, + "quantizer_type": "rlfq", + }, + ) + torch.manual_seed(0) + model = MossAudioTokenizerModel(config).eval() + with torch.no_grad(): + for parameter in model.parameters(): + parameter.copy_(torch.randn_like(parameter) * 0.5) + model.quantizer.build_decode_lut(2, dtype=torch.float32) + return model + + +def stream_decode(model, codes, chunk_frames, *, headroom): + n_q, _, total = codes.shape + model.initialize_decoder_state_pool(1, 0, chunk_frames=chunk_frames if headroom else 0) + try: + slot = torch.zeros(1, dtype=torch.long) + valid = torch.ones(1, dtype=torch.bool) + parts = [] + for start in range(0, total, chunk_frames): + chunk = codes[:, :, start : start + chunk_frames] + lengths = torch.full((1,), chunk.shape[-1], dtype=torch.long) + parts.append(model.decode_streaming_batch(chunk, lengths, slot, valid).audio) + return torch.cat(parts, dim=-1) + finally: + model.close_decoder_state_pool() + + +@torch.no_grad() +def test_ring_headroom_sizes_each_layer_for_one_chunk(): + model = tiny_codec() + transformers = [m for m in model.decoder if hasattr(m, "tokens_per_input_frame")] + assert [m.tokens_per_input_frame for m in transformers] == [1, 2] + assert [m.transformer.layers[0].self_attn.context for m in transformers] == [8, 8] + model.initialize_decoder_state_pool(1, 0, chunk_frames=6) + try: + rings = [m.transformer.layers[0].self_attn._streaming_state.kv_cache.capacity for m in transformers] + assert rings == [8 + 6, 8 + 12] + finally: + model.close_decoder_state_pool() + model.initialize_decoder_state_pool(1, 0) + try: + rings = [m.transformer.layers[0].self_attn._streaming_state.kv_cache.capacity for m in transformers] + assert rings == [8, 8] + finally: + model.close_decoder_state_pool() + + +@torch.no_grad() +def test_chunked_streaming_matches_whole_sequence_only_with_headroom(): + model = tiny_codec() + torch.manual_seed(1) + codes = torch.randint(0, 16, (2, 1, 18)) + reference = model._decode_frame(codes).audio + with_headroom = stream_decode(model, codes, 6, headroom=True) + legacy = stream_decode(model, codes, 6, headroom=False) + assert with_headroom.shape == reference.shape == legacy.shape + torch.testing.assert_close(with_headroom, reference, rtol=1e-5, atol=1e-5) + # The legacy context-sized ring lets the 12-token chunk evict keys inside + # its own window, so streaming diverges from whole-sequence decoding. + assert not torch.allclose(legacy, reference, rtol=1e-3, atol=1e-3) + + +@torch.no_grad() +def test_padded_terminal_tail_matches_exact_tail_with_headroom(): + model = tiny_codec() + torch.manual_seed(2) + codes = torch.randint(0, 16, (2, 1, 14)) + history, tail = codes[:, :, :12], codes[:, :, 12:] + padded_tail = torch.nn.functional.pad(tail, (0, 4)) # 2 real frames padded to a 6-frame step + outputs = {} + for name, last in [("exact", tail), ("padded", padded_tail)]: + model.initialize_decoder_state_pool(1, 0, chunk_frames=6) + try: + slot, valid = torch.zeros(1, dtype=torch.long), torch.ones(1, dtype=torch.bool) + for chunk in (history[:, :, :6], history[:, :, 6:]): + model.decode_streaming_batch(chunk, torch.tensor([6]), slot, valid) + audio = model.decode_streaming_batch(last, torch.tensor([last.shape[-1]]), slot, valid).audio + outputs[name] = audio[..., : 2 * model.downsample_rate] + finally: + model.close_decoder_state_pool() + torch.testing.assert_close(outputs["padded"], outputs["exact"], rtol=1e-5, atol=1e-5) diff --git a/tests/model_executor/models/moss_tts/test_first_chunk_fast_path.py b/tests/model_executor/models/moss_tts/test_first_chunk_fast_path.py new file mode 100644 index 000000000..85e53e80b --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_first_chunk_fast_path.py @@ -0,0 +1,547 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The first-chunk fast path emits audio out of band; the main path resumes the stream.""" + +import queue +import threading +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from tests.model_executor.models.moss_tts.test_streaming_terminal_batch import session as cpu_session +from vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path import MossFirstChunkFastPath, _SlotHandoff +from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import MossTTSCodecDecoder +from vllm_omni.worker_v2.first_audio_sender import engine_output_queue_sink + +pytestmark = pytest.mark.core_model + +N_VQ, SAMPLES, CHANNELS = 2, 4, 2 + + +class _Session: + def __init__(self, capacity): + self.free = list(range(capacity)) + self.reset_waits = [] + + def acquire(self, *, oldest=False): + return self.free.pop(0 if oldest else -1) if self.free else None + + def order_after_reset(self, slot, stream): + self.reset_waits.append(slot) + + def release(self, slot, **kwargs): + self.free.append(slot) + + +class _Wrapper: + batch_sizes = [1, 2, 4, 8] + + def __init__(self, device): + self.state = torch.zeros(16, device=device) + self.calls = [] + self.gate = threading.Event() + self.gate.set() + + def decode(self, codes, slots): + self.gate.wait() + self.calls.append(int(codes.shape[1])) + base = codes.sum(0)[:, 0].float() + 100 * slots.float() + channel = torch.arange(CHANNELS, device=codes.device).float() + audio = base[:, None, None] + channel[None, :, None] + torch.zeros(1, 1, SAMPLES, device=codes.device) + self.state[slots] = 1 + return audio.to(torch.bfloat16), None, int(codes.shape[1]) + + +@pytest.fixture +def cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + return torch.device("cuda") + + +def _fast(device, capacity=4): + session, wrapper = _Session(capacity), _Wrapper(device) + fast = MossFirstChunkFastPath( + session, + wrapper, + n_vq=N_VQ, + frames=1, + codebook_size=1024, + samples_per_frame=SAMPLES, + n_channels=CHANNELS, + sample_rate=torch.tensor(24000, dtype=torch.int32), + device=device, + ) + return fast, session, wrapper + + +class _CallbackSink: + def __init__(self, callback): + self.callback = callback + + def prepare(self, request_ids): + callback = self.callback + + class Delivery: + routes = {rid: (3 if rid == "r0" else 0) for rid in request_ids} + + def __call__(self, ids, rows, sr): + for rid, row in zip(ids, rows, strict=True): + callback(self.routes[rid], rid, {"model_outputs": row, "sr": sr}) + + def fail(self, ids): + pass + + return Delivery() + + +@pytest.mark.cuda +def test_emits_each_first_chunk_and_hands_slot_to_main_path(cuda): + fast, session, wrapper = _fast(cuda) + emitted = {} + done = threading.Event() + + def sink(client, req_id, payload): + emitted[req_id] = (client, payload["model_outputs"], int(payload["sr"])) + if len(emitted) == 2: + done.set() + + slots: dict[str, int] = {} + assert not fast.submit("r0", "k0", [1, 2], slots) # unbound: main path owns it + fast.bind(_CallbackSink(sink)) + wrapper.gate.clear() + assert fast.submit("r0", "k0", torch.tensor([1, 2]), slots) + assert fast.submit("r1", "k1", [5, 7], slots) + assert not fast.submit("r0", "k0", [1, 2], slots) # already has a slot + assert not fast.submit("r2", "k2", [1, 2, 3, 4], slots) # not a 1-frame chunk + assert slots == {"k0": 0, "k1": 1} and session.free == [2, 3] + wrapper.gate.set() + assert done.wait(10) + fast.close() + + assert wrapper.calls in ([1, 1], [2]) + for req_id, client, code_sum, slot in (("r0", 3, 3, 0), ("r1", 0, 12, 1)): + got_client, wav, sr = emitted[req_id] + assert got_client == client and sr == 24000 and wav.device.type == "cpu" + expected = (code_sum + 100 * slot + torch.arange(CHANNELS).float())[:, None].expand(CHANNELS, SAMPLES) + torch.testing.assert_close(wav, expected) + assert sorted(session.reset_waits) == [0, 1] + + # The main path consumes the marker once and is ordered after the decode. + assert fast.take_decoded("k0") and not fast.take_decoded("k0") + fast.order_after(0) + fast.order_after(0) # no pending handoff: no-op + assert wrapper.state[:2].tolist() == [1, 1] + + +@pytest.mark.cuda +def test_order_after_blocks_until_decode_finishes(cuda): + fast, _, wrapper = _fast(cuda) + fast.bind(_CallbackSink(lambda *args: None)) + wrapper.gate.clear() + assert fast.submit("r0", "k0", [1, 2], {}) + finished = threading.Event() + + def wait_for_slot(): + fast.order_after(0) + finished.set() + + waiter = threading.Thread(target=wait_for_slot) + waiter.start() + assert not finished.wait(0.2) + wrapper.gate.set() + assert finished.wait(10) + waiter.join() + fast.close() + + +@pytest.mark.cuda +def test_no_free_slot_leaves_chunk_to_main_path(cuda): + fast, _, _ = _fast(cuda, capacity=0) + fast.bind(_CallbackSink(lambda *args: None)) + slots: dict[str, int] = {} + assert not fast.submit("r0", "k0", [1, 2], slots) + assert slots == {} and not fast.take_decoded("k0") + fast.close() + + +class _FakeFast: + frames = 1 + + def __init__(self, decoded, unsent=None): + self.decoded = set(decoded) + self.unsent = dict(unsent or {}) + self.ordered = [] + self.forgotten = [] + + def take_decoded(self, key): + if key in self.decoded: + self.decoded.remove(key) + return True + return False + + def get_request_slot(self, key, slots): + return slots.get(key) + + def order_after(self, slot): + self.ordered.append(slot) + + def take_unsent_audio(self, key): + return self.unsent.pop(key, None) + + def forget(self, key): + self.forgotten.append(key) + + def gate(self): + self.gated = getattr(self, "gated", 0) + 1 + + +def _decoder(s, fast): + decoder = object.__new__(MossTTSCodecDecoder) + nn.Module.__init__(decoder) + decoder._stream_max_step_frames = 15 + decoder._stream_state_capacity = 8 + decoder._stream_req_slots = {} + decoder._ensure_stream_session = lambda: s + decoder._first_chunk_fast_path = fast + decoder._stream_session = s + decoder._codec_stream = None + return decoder + + +@pytest.mark.cpu +def test_main_path_skips_fast_decoded_chunk_and_resumes_state(): + s = cpu_session() + reference = cpu_session() + fast = _FakeFast({"a"}) + decoder = _decoder(s, fast) + # The fast path leased slot 0 and decoded frame 0 of "a" into it. + slot = s.acquire(oldest=True) + decoder._stream_req_slots["a"] = slot + s.step({slot: torch.full((2, 1), 3.0)}) + ref_slot = reference.acquire() + reference.step({ref_slot: torch.full((2, 1), 3.0)}) + + out = decoder._decode_streaming_batch([(0, "a", torch.full((2, 1), 3.0), False), (1, "b", torch.ones(2, 1), False)]) + assert 0 not in out and out[1].tolist() == [[2.0]] + assert fast.ordered == [slot] + + out = decoder._decode_streaming_batch([(0, "a", torch.ones(2, 3), True)]) + expected = reference.step({ref_slot: torch.ones(2, 3)})[ref_slot] + torch.testing.assert_close(out[0], expected) + assert "a" not in decoder._stream_req_slots and fast.ordered == [slot, slot] + + +@pytest.mark.cpu +def test_abort_orders_release_after_fast_decode(): + s = cpu_session() + fast = _FakeFast({"a"}) + decoder = _decoder(s, fast) + slot = decoder._stream_req_slots["a"] = s.acquire(oldest=True) + decoder.on_requests_finished({"a"}) + assert fast.ordered == [slot] and fast.forgotten == ["a"] and decoder._stream_req_slots == {} + assert slot in s._free_stream_slots + + +@pytest.mark.cpu +def test_main_path_returns_unsent_first_chunk_without_decoding_twice(): + s = cpu_session() + first = torch.full((1, 1), 7.0) + fast = _FakeFast({"a"}, {"a": first}) + decoder = _decoder(s, fast) + slot = s.acquire(oldest=True) + decoder._stream_req_slots["a"] = slot + s.step({slot: torch.full((2, 1), 3.0)}) + + out = decoder._decode_streaming_batch([(0, "a", torch.full((2, 1), 3.0), False)]) + torch.testing.assert_close(out[0], first) + assert fast.unsent == {} and fast.ordered == [slot] + + # A payload containing a later frame must concatenate it after the + # preserved first samples, while advancing state only for the later frame. + first = torch.full((1, 1), 7.0) + fast.decoded.add("b") + fast.unsent["b"] = first + slot_b = s.acquire(oldest=True) + decoder._stream_req_slots["b"] = slot_b + s.step({slot_b: torch.full((2, 1), 3.0)}) + out = decoder._decode_streaming_batch([(0, "b", torch.full((2, 2), 3.0), False)]) + torch.testing.assert_close(out[0][..., :1], first) + assert out[0].shape[-1] == 2 + + +@pytest.mark.cpu +def test_oldest_acquire_reuses_least_recently_released_slot(): + s = cpu_session() + a, b = s.acquire(), s.acquire() + s.release(a) + s.release(b) + assert s.acquire() == b + assert s.acquire(oldest=True) == 7 # never-used slots were freed before any release + assert s._free_stream_slots[-1] == a + + +@pytest.mark.cuda +def test_failed_decode_reports_error_without_replaying_state(cuda): + fast, _, wrapper = _fast(cuda) + + def fail(codes, slots): + raise RuntimeError("boom") + + wrapper.decode = fail + fast.bind(_CallbackSink(lambda *args: pytest.fail("no output expected"))) + slots: dict[str, int] = {} + assert fast.submit("r0", "k0", [1, 2], slots) + with pytest.raises(RuntimeError, match="first-chunk decode failed"): + fast.order_after(slots["k0"]) + assert fast.take_decoded("k0") + fast.close() + + +@pytest.mark.cuda +def test_failed_output_handoff_retains_decoded_samples(cuda): + fast, _, _ = _fast(cuda) + + def reject(*args): + raise RuntimeError("output queue closed") + + fast.bind(_CallbackSink(reject)) + slots: dict[str, int] = {} + assert fast.submit("r0", "k0", [1, 2], slots) + fast.order_after(slots["k0"]) + assert fast.take_decoded("k0") + wav = fast.take_unsent_audio("k0") + assert wav is not None and wav.shape == (CHANNELS, SAMPLES) + assert fast.take_unsent_audio("k0") is None + fast.close() + + +@pytest.mark.cuda +def test_gate_orders_main_stream_after_inflight_fast_decode(cuda): + fast, _, _ = _fast(cuda) + side = torch.cuda.Stream() + with torch.cuda.stream(side): + torch.cuda._sleep(200_000_000) + pending = torch.cuda.Event() + pending.record(side) + fast._inflight = pending + fast.gate() # disabled: no wait + unordered = torch.cuda.Event() + unordered.record() + fast._gate_main = True + fast.gate() + ordered = torch.cuda.Event() + ordered.record() + unordered.synchronize() + assert not pending.query() and not ordered.query() + ordered.synchronize() + assert pending.query() + + +def _cpu_fast(monkeypatch, *, scheduler=None): + # Replace only CUDA allocations, retaining the real lifecycle and worker. + import vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path as module + + monkeypatch.setattr(torch.Tensor, "pin_memory", lambda tensor: tensor) + monkeypatch.setattr(torch.cuda, "Stream", lambda **kwargs: None) + monkeypatch.setattr(torch.cuda, "Event", lambda: None) + monkeypatch.setattr(module, "current_omni_platform", SimpleNamespace(set_device=lambda device: None)) + session = _Session(2) + fast = MossFirstChunkFastPath( + session, + SimpleNamespace(batch_sizes=[1, 2]), + n_vq=2, + frames=1, + codebook_size=1024, + samples_per_frame=4, + n_channels=1, + sample_rate=torch.tensor(24000), + device=torch.device("cpu:0"), + handoff_timeout_s=0.01, + ) + outputs: queue.Queue = queue.Queue() + scheduler = scheduler or SimpleNamespace(requests={"r": SimpleNamespace(client_index=3)}) + sink = engine_output_queue_sink(outputs, scheduler, upstream_first_audio=False) + return fast, session, sink, outputs + + +def _cancel_during_admission(fast, session, scheduler, sink): + entered, release = threading.Event(), threading.Event() + + class PausedSink: + def prepare(self, ids): + delivery = sink.prepare(ids) + if delivery.routes: + entered.set() + assert release.wait(5) + return delivery + + fast.bind(PausedSink()) + decoder = _decoder(session, fast) + errors = [] + finished = threading.Event() + finishing = threading.Event() + + def submit(): + try: + assert fast.submit("r", "r", [1, 2], decoder._stream_req_slots) + except BaseException as error: + errors.append(error) + + def cancel(): + try: + finishing.set() + decoder.on_requests_finished({"r"}) + finished.set() + except BaseException as error: + errors.append(error) + + producer = threading.Thread(target=submit) + cleaner = threading.Thread(target=cancel) + producer.start() + try: + assert entered.wait(5) + assert decoder._stream_req_slots == {} + # Scheduler cancellation retires the route before the runner hook. + scheduler.requests.pop("r") + cleaner.start() + assert finishing.wait(5) + assert not finished.wait(0.1) + release.set() + producer.join(5) + cleaner.join(5) + assert not producer.is_alive() and not cleaner.is_alive() + assert not errors, errors + assert finished.is_set() + assert decoder._stream_req_slots == {} + assert fast._decoded == {} and fast._handoffs == {} + assert sorted(session.free) == [0, 1] + assert not fast.submit("r", "r", [1, 2], decoder._stream_req_slots) + + # A fresh request can use and return the same capacity immediately. + scheduler.requests["reuse"] = SimpleNamespace(client_index=3) + release.set() + assert fast.submit("reuse", "reuse", [1, 2], decoder._stream_req_slots) + scheduler.requests.pop("reuse") + decoder.on_requests_finished({"reuse"}) + assert decoder._stream_req_slots == {} and sorted(session.free) == [0, 1] + finally: + release.set() + producer.join(5) + if cleaner.ident is not None: + cleaner.join(5) + fast.close() + + +@pytest.mark.cpu +def test_cancel_waits_for_first_chunk_admission_and_reclaims_slot(monkeypatch): + scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=3)}) + fast, session, sink, _ = _cpu_fast(monkeypatch, scheduler=scheduler) + + # Keep the actual queue/worker/handoff lifecycle, omitting CUDA decode. + def decode(jobs): + for job in jobs: + job.handoff.done.set() + + monkeypatch.setattr(fast, "_decode", decode) + fast._handoff_timeout_s = 5.0 + _cancel_during_admission(fast, session, scheduler, sink) + + +@pytest.mark.cuda +def test_cancel_during_first_chunk_admission_orders_cuda_handoff_and_reuse(cuda): + fast, session, _ = _fast(cuda, capacity=2) + scheduler = SimpleNamespace(requests={"r": SimpleNamespace(client_index=3)}) + sink = engine_output_queue_sink(queue.Queue(), scheduler, upstream_first_audio=False) + _cancel_during_admission(fast, session, scheduler, sink) + + +@pytest.mark.cpu +def test_slot_timeout_keeps_handoff_and_prevents_slot_reuse(monkeypatch): + fast, _, _, _ = _cpu_fast(monkeypatch) + handoff = _SlotHandoff() + fast._handoffs[0] = handoff + for _ in range(2): + with pytest.raises(TimeoutError, match="slot 0"): + fast.order_after(0) + assert fast._handoffs[0] is handoff + handoff.done.set() + fast.order_after(0) + assert fast._handoffs == {} + + +@pytest.mark.cpu +def test_closed_decoder_rejects_jobs_and_cannot_restart(monkeypatch): + fast, session, sink, _ = _cpu_fast(monkeypatch) + fast.bind(sink) + fast.close() + slots: dict[str, int] = {} + assert not fast.submit("r", "external", [1, 2], slots) + assert slots == {} and session.free == [0, 1] + with pytest.raises(RuntimeError, match="closed"): + fast.bind(sink) + + +@pytest.mark.cpu +def test_missing_route_does_not_claim_slot(monkeypatch): + fast, session, sink, outputs = _cpu_fast(monkeypatch, scheduler=SimpleNamespace(requests={})) + fast.bind(sink) + try: + assert not fast.submit("gone", "external", [1, 2], {}) + assert session.free == [0, 1] and outputs.empty() + finally: + fast.close() + + +@pytest.mark.cpu +def test_decode_failure_wakes_current_and_queued_handoffs(monkeypatch): + from vllm.v1.engine import FinishReason + + fast, _, sink, outputs = _cpu_fast(monkeypatch) + entered, release = threading.Event(), threading.Event() + + def fail(jobs): + entered.set() + assert release.wait(2) + raise RuntimeError("partial replay failed") + + fast._decode = fail + fast.bind(sink) + slots: dict[str, int] = {} + try: + assert fast.submit("r", "a", [1, 2], slots) + assert entered.wait(2) + assert fast.submit("r", "b", [3, 4], slots) + release.set() + fast._thread.join(2) + assert not fast._thread.is_alive() + for slot in slots.values(): + with pytest.raises(RuntimeError, match="decode failed"): + fast.order_after(slot) + assert all(fast.take_decoded(key) for key in slots) + assert not fast.submit("r", "c", [1, 2], slots) + assert outputs.qsize() == 2 + while not outputs.empty(): + assert outputs.get_nowait()[1].outputs[0].finish_reason == FinishReason.ERROR + finally: + release.set() + fast.close() + + +@pytest.mark.cpu +def test_platform_failure_wakes_pending_handoff(monkeypatch): + import vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path as module + + fast, _, _, _ = _cpu_fast(monkeypatch) + fast._handoffs[0] = _SlotHandoff() + + def fail(device): + raise RuntimeError("device unavailable") + + monkeypatch.setattr(module, "current_omni_platform", SimpleNamespace(set_device=fail)) + fast._run() + assert fast._closed + with pytest.raises(RuntimeError, match="decode failed"): + fast.order_after(0) diff --git a/tests/model_executor/models/moss_tts/test_local_code_rows_snapshot.py b/tests/model_executor/models/moss_tts/test_local_code_rows_snapshot.py new file mode 100644 index 000000000..b9109c0e4 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_local_code_rows_snapshot.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Per-request code rows keep the runner payload layout with one host copy.""" + +import pytest +import torch + +from vllm_omni.model_executor.models.moss_tts.local_model_state import _CodeRowsSnapshot +from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_rows_map_to_batch_positions_and_copy_once(): + codes = torch.arange(6, dtype=torch.long).reshape(2, 3) + snapshot = _CodeRowsSnapshot(codes, [2, 0], 4) + assert isinstance(snapshot, PackedOutputSnapshot) and snapshot.producer_event is None + rows = snapshot["codes"]["audio"] + assert [tuple(row.shape) for row in rows] == [(1, 3), (0, 3), (1, 3), (0, 3)] + assert torch.equal(rows[2], codes[:1]) and torch.equal(rows[0], codes[1:]) + + copies = [] + + def copy(tensor): + copies.append(tensor) + return tensor.clone() + + host = snapshot.copy_to_cpu(copy)["codes"]["audio"] + assert len(copies) == 1 and copies[0] is codes + codes.zero_() + assert host[2].tolist() == [[0, 1, 2]] and host[0].tolist() == [[3, 4, 5]] + assert host[1].numel() == 0 and host[1].dtype == torch.long diff --git a/tests/model_executor/models/moss_tts/test_local_compiled_sampler.py b/tests/model_executor/models/moss_tts/test_local_compiled_sampler.py new file mode 100644 index 000000000..6c926c134 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_local_compiled_sampler.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import pytest +import torch + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cuda, + pytest.mark.tts, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), +] + + +def test_explicit_generator_bypasses_compiled_audio_sampler(): + from transformers import GPT2Config + + from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import ( + MossTTSLocalDepthTransformer, + ) + + model = MossTTSLocalDepthTransformer(GPT2Config(n_embd=80, n_head=1, n_inner=160)).cuda().bfloat16().eval() + heads = torch.nn.ModuleList([torch.nn.Linear(80, 32) for _ in range(2)]).cuda().bfloat16() + embeddings = torch.nn.ModuleList([torch.nn.Embedding(32, 80) for _ in range(2)]).cuda().bfloat16() + binary = torch.nn.Linear(80, 2).cuda().bfloat16() + hidden = torch.randn(1, 80, device="cuda", dtype=torch.bfloat16) + generator = torch.Generator(device="cuda").manual_seed(17) + expected = model.generate_frame(hidden, heads, embeddings, binary, n_vq=2, generator=generator) + + def forbidden(*args, **kwargs): + raise AssertionError("compiled path selected") + + model._compiled_audio_sampler = forbidden + generator.manual_seed(17) + actual = model.generate_frame(hidden, heads, embeddings, binary, n_vq=2, generator=generator) + assert all(torch.equal(a, b) for a, b in zip(expected, actual)) + with pytest.raises(AssertionError, match="compiled path selected"): + model.generate_frame(hidden, heads, embeddings, binary, n_vq=2) diff --git a/tests/model_executor/models/moss_tts/test_local_determined_tokens.py b/tests/model_executor/models/moss_tts/test_local_determined_tokens.py new file mode 100644 index 000000000..2d529cf26 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_local_determined_tokens.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Determined text tokens retain the upstream sampling/counting contract.""" + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from vllm.sampling_params import SamplingParams +from vllm.v1.worker.gpu.sample.sampler import Sampler +from vllm.v1.worker.gpu.states import RequestState + +from tests.model_executor.models.moss_tts.test_local_model_state import _batch, _state +from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState +from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner +from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner + +pytestmark = pytest.mark.core_model + + +@pytest.mark.cuda +@pytest.mark.parametrize("seed", [None, 17]) +@pytest.mark.parametrize("temperature", [0.0, 1.7]) +def test_matches_real_sampler_for_mixed_prefill_decode_stop_and_reordering(seed, temperature): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + device = torch.device("cuda") + reqs = RequestState(5, 32, 16, 0, 8, device) + sampler = Sampler(5, 8, device, reqs) + for i, length in enumerate([4, 1, 3]): + reqs.add_request(str(i), length, [2] * length, 0, 8) + slot = reqs.req_id_to_index[str(i)] + sampler.add_request(slot, length, SamplingParams(temperature=temperature, top_k=3, top_p=0.8, seed=seed)) + reqs.apply_staged_writes() + sampler.apply_staged_writes() + state = _state(MossLocalModelState, device) + state._direct_tokens = True + state.model._batch_state = None + batch = _batch(device, [2, 4, 3], [1, 2, 1]) + batch.seq_lens = torch.tensor([5, 2, 2], device=device, dtype=torch.int32) + batch.cu_num_logits_np = np.arange(4, dtype=np.int32) + batch.cu_num_logits = torch.tensor(batch.cu_num_logits_np, device=device) + batch.expanded_idx_mapping = batch.idx_mapping + batch.expanded_local_pos = torch.zeros(3, dtype=torch.int32, device=device) + batch.logits_indices = batch.query_start_loc[1:].long() - 1 + batch.positions = torch.arange(batch.num_tokens, device=device) + state.model._batch_should_continue = torch.tensor([False, True, True], device=device) + with torch.inference_mode(): + logits = state.model.compute_logits(torch.ones(3, 4, device=device)) + expected = sampler(logits, batch) + actual = state.sample_determined_tokens(batch, sampler) + assert actual is not None + for field in ["sampled_token_ids", "num_sampled", "num_rejected"]: + torch.testing.assert_close(getattr(actual, field), getattr(expected, field), rtol=0, atol=0) + assert actual.num_sampled.tolist() == [1, 0, 1] + snapshot = actual.sampled_token_ids.clone() + state.model._batch_should_continue.logical_not_() + torch.testing.assert_close(actual.sampled_token_ids, snapshot, rtol=0, atol=0) + + +@pytest.mark.cuda +@pytest.mark.parametrize( + "constraint", + [ + {"min_tokens": 2, "stop_token_ids": [1]}, + {"allowed_token_ids": [1, 2]}, + {"logit_bias": {1: -10.0}}, + ], +) +def test_real_sampler_constraints_disable_determined_tokens(constraint): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + device = torch.device("cuda") + reqs = RequestState(5, 32, 16, 0, 8, device) + sampler = Sampler(5, 8, device, reqs) + reqs.add_request("constrained", 1, [2], 0, 8) + slot = reqs.req_id_to_index["constrained"] + sampler.add_request(slot, 1, SamplingParams(**constraint)) + state = _state(MossLocalModelState, device) + state._direct_tokens = True + # Use the actual upstream request registration, including its combined + # min-token/allowed-ID/bias flag, rather than mocking that flag ourselves. + assert state.sample_determined_tokens(SimpleNamespace(idx_mapping_np=np.array([slot])), sampler) is None + + +@pytest.mark.cpu +@pytest.mark.parametrize("feature", ["logprobs", "bias", "penalty", "bad_words", "thinking", "trace", "mask", "nans"]) +def test_distribution_features_fall_back_before_using_gpu(feature): + sampler = object.__new__(Sampler) + sampler.compute_nans = feature == "nans" + sampler.return_sampling_mask = feature == "mask" + sampler.trace_replay_state = object() if feature == "trace" else None + sampler.get_logprobs_dims = lambda rows: (1, 0) if feature == "logprobs" else None + sampler.logit_bias_state = SimpleNamespace(use_logit_bias=np.array([feature == "bias"])) + sampler.penalties_state = SimpleNamespace(use_penalty=np.array([feature == "penalty"])) + sampler.bad_words_state = SimpleNamespace(num_bad_words=SimpleNamespace(np=np.array([feature == "bad_words"]))) + sampler.thinking_budget_state = SimpleNamespace(enabled=True, use_thinking_budget=np.array([feature == "thinking"])) + state = object.__new__(MossLocalModelState) + state._direct_tokens = True + assert state.sample_determined_tokens(SimpleNamespace(idx_mapping_np=np.array([0])), sampler) is None + + +@pytest.mark.cpu +@pytest.mark.parametrize("feature", ["grammar", "sharding", "pipeline", "speculative", "empty", "declined"]) +def test_runner_retains_upstream_path_for_unsupported_execution(mocker, feature): + hook = mocker.Mock(return_value=None) + + class State: + def sample_determined_tokens(self, batch, sampler): + return hook(batch, sampler) + + runner = object.__new__(OmniARModelRunner) + runner.model_state = State() + runner.sampler = object() + runner.batch_sharder = object() if feature == "sharding" else None + runner.pp_handler = object() if feature == "pipeline" else None + batch = SimpleNamespace(num_draft_tokens=int(feature == "speculative"), num_reqs=0 if feature == "empty" else 1) + grammar = object() if feature == "grammar" else None + fallback = mocker.patch.object(OmniGPUModelRunner, "sample", return_value="upstream") + hidden = object() + assert runner.sample(hidden, batch, grammar) == "upstream" + fallback.assert_called_once_with(hidden, batch, grammar) + assert hook.call_count == int(feature == "declined") diff --git a/tests/model_executor/models/moss_tts/test_local_eager_mtp.py b/tests/model_executor/models/moss_tts/test_local_eager_mtp.py new file mode 100644 index 000000000..cf3893cbd --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_local_eager_mtp.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Eager MTP emits the same frames one step earlier with identical backbone inputs.""" + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from tests.model_executor.models.moss_tts.test_local_model_state import _admit, _batch, _state, _step +from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState + +pytestmark = pytest.mark.core_model + + +@pytest.fixture(params=[pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=pytest.mark.cuda)]) +def device(request): + if request.param == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA required") + return torch.device(request.param) + + +def _run(state, device, schedule, stop_step): + computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32) + prompt[3], prompt[0] = 4, 2 + req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed) + embeds = [] + frames: dict[int, list[torch.Tensor]] = {3: [], 0: []} + for index, (slots, counts) in enumerate(schedule): + state.model.force_stop = index == stop_step + embed, payload, _ = _step(state, _batch(device, slots, counts), req_states) + embeds.append(embed) + rows = payload.get("codes", {}).get("audio", []) + for slot, row in zip(slots, rows): + if row.numel() and bool(row.ne(state.model.audio_pad_token_id).any()): + frames[slot].append((index, row.clone())) + for slot, count in zip(slots, counts): + computed[slot] += count + return embeds, frames + + +def test_eager_frames_match_canonical_one_step_earlier(device): + schedule = [([3, 0], [2, 2]), ([3, 0], [2, 1]), ([3, 0], [1, 1]), ([0, 3], [1, 1]), ([3, 0], [1, 1])] + canonical, eager = (_state(MossLocalModelState, device) for _ in range(2)) + eager._local_eager_mtp = True + for state in (canonical, eager): + _admit(state, 3, "long", 17) + _admit(state, 0, "short", 17) + # Stop both streams at the canonical step that would draw the 4th step's frames. + c_embeds, c_frames = _run(canonical, device, schedule, stop_step=4) + e_embeds, e_frames = _run(eager, device, schedule, stop_step=3) + + for c, e in zip(c_embeds, e_embeds): + torch.testing.assert_close(c, e, rtol=0, atol=0) + for slot in (3, 0): + assert [f for _, f in c_frames[slot]] and len(c_frames[slot]) == len(e_frames[slot]) + for (ci, cf), (ei, ef) in zip(c_frames[slot], e_frames[slot]): + assert ei == ci - 1 + torch.testing.assert_close(cf, ef, rtol=0, atol=0) + + +def test_stopped_stream_does_not_emit_again(device): + state = _state(MossLocalModelState, device) + state._local_eager_mtp = True + _admit(state, 0, "short", 17) + computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32) + prompt[0] = 2 + req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed) + emitted = [] + for index, count in enumerate([2, 1, 1, 1]): + state.model.force_stop = index == 1 + _, payload, _ = _step(state, _batch(device, [0], [count]), req_states) + row = payload["codes"]["audio"][0] + emitted.append(bool(row.ne(state.model.audio_pad_token_id).any())) + computed[0] += count + state.model.force_stop = False + # Frame at the prefill step, stop at step 1, then asynchronous overrun steps. + assert emitted == [True, False, False, False] + + +def test_unmarked_decode_row_ignores_stale_slot_state(device): + canonical, eager = (_state(MossLocalModelState, device) for _ in range(2)) + eager._local_eager_mtp = True + for state in (canonical, eager): + _admit(state, 0, "short", 17) + # A decode row whose completing prefill was not observed in this state. + eager._keep_pool[0] = False + eager._eager_emb_pool[0].fill_(3.0) + computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32) + prompt[0], computed[0] = 2, 2 + req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed) + for state in (canonical, eager): + state.intermediate_buffer.buffers[0]["audio_state"] = {"is_stopping": False} + c_embed, c_payload, _ = _step(canonical, _batch(device, [0], [1]), req_states) + e_embed, e_payload, _ = _step(eager, _batch(device, [0], [1]), req_states) + torch.testing.assert_close(c_embed, e_embed, rtol=0, atol=0) + torch.testing.assert_close(c_payload["codes"]["audio"][0], e_payload["codes"]["audio"][0], rtol=0, atol=0) + + +@pytest.mark.cpu +def test_first_prefill_through_runner_preserves_parent_eager_contract(monkeypatch, mocker): + """The Local config must not enable the parent's incompatible eager hook.""" + from vllm_omni.worker_v2 import omni_ar_model_runner as runner_module + from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState + from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState + from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner + + device = torch.device("cpu") + template = _state(MossLocalModelState, device) + model = template.model + model.config = SimpleNamespace(mrv2_eager_mtp=True) + + def initialize_parent(state, *args): + state.__dict__.update(template.__dict__) + state._eager_mtp, state._eager_rows = False, None + state._eager_state = EagerMTPState(state) + + monkeypatch.setattr(OmniModelState, "__init__", initialize_parent) + state = MossLocalModelState(SimpleNamespace(), model, None, device) + _admit(state, 0, "req", 17) + batch = _batch(device, [0], [2]) + batch.req_ids = ["req"] + computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32) + prompt[0] = 2 + _step(state, batch, SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed)) + hook = mocker.spy(state, "run_eager_mtp") + + runner = OmniARModelRunner.__new__(OmniARModelRunner) + runner.execute_model_state = SimpleNamespace( + input_batch=batch, + hidden_states=torch.zeros(2, 4), + finished_req_ids=set(), + ec_connector_output=None, + routed_experts=None, + ) + runner._kv_extracted_req_ids = runner._last_aux_output = runner._last_multimodal_outputs = None + runner.is_last_pp_rank, runner.pp_handler, runner.check_ep_fault = True, None, False + runner.model_config = SimpleNamespace(async_chunk=False) + runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(engine_output_type="text")) + runner.model_state, runner.model = state, model + runner.req_states = SimpleNamespace( + all_token_ids=SimpleNamespace(gpu=torch.zeros(5, 2)), + num_computed_tokens=SimpleNamespace(gpu=torch.from_numpy(computed)), + prompt_len=SimpleNamespace(np=prompt, gpu=torch.from_numpy(prompt)), + ) + runner.main_stream = runner.output_copy_stream = mocker.Mock() + runner.eplb = runner._finalize_native_data_plane_output = runner._reserve_native_data_plane_outputs = mocker.Mock() + runner.sampler = None + runner.sample = mocker.Mock( + return_value=( + SimpleNamespace(sampled_token_ids=torch.tensor([[2]])), + torch.ones(1), + torch.zeros(1), + ) + ) + runner.prompt_logprobs_worker = SimpleNamespace(compute_prompt_logprobs=mocker.Mock(return_value={})) + runner.postprocess_sampled = mocker.Mock() + runner.kv_connector = SimpleNamespace(post_forward=mocker.Mock(return_value=None)) + mock_out = SimpleNamespace(copy_event=None) + monkeypatch.setattr(runner_module, "OmniAsyncOutput", mocker.Mock(return_value=mock_out)) + + assert runner.sample_tokens(None) is mock_out + hook.assert_called_once() + assert state._eager_mtp is False and state._eager_rows is None + assert state._local_eager_mtp is True diff --git a/tests/model_executor/models/moss_tts/test_local_model_state.py b/tests/model_executor/models/moss_tts/test_local_model_state.py new file mode 100644 index 000000000..4956d5c89 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_local_model_state.py @@ -0,0 +1,286 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Compare slot lifecycle with canonical Local hooks, including output ownership.""" + +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +import torch +from vllm.v1.worker.gpu.model_states.default import DefaultModelState + +from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState +from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import MossTTSLocalTalkerForGeneration +from vllm_omni.worker_v2.model_states import init_omni_model_state +from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState +from vllm_omni.worker_v2.model_states.intermediate_buffer import OmniIntermediateBuffer +from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState + +pytestmark = pytest.mark.core_model + + +@pytest.fixture(params=[pytest.param("cpu", marks=pytest.mark.cpu), pytest.param("cuda", marks=pytest.mark.cuda)]) +def device(request): + if request.param == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA required") + return torch.device(request.param) + + +def _model(device): + model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration) + torch.nn.Module.__init__(model) + dtype = torch.bfloat16 if device.type == "cuda" else torch.float32 + model.model = torch.nn.Module() + weight = torch.arange(32, dtype=dtype, device=device).reshape(8, 4) / 16 + model.model.embed_tokens = torch.nn.Embedding(8, 4, _weight=weight) + model.hidden_size, model.n_vq = 4, 2 + model.audio_pad_token_id, model.text_vocab_size = 8, 8 + model.audio_assistant_slot_token_id, model.im_end_token_id = 2, 3 + model.talker_mtp_output_key = ("audio_codes", "current") + model.talker_mtp_graph_safe = False + model.gpu_resident_buffer_keys = {("hidden_states", "last"), ("audio_codes", "current")} + model.audio_lm_heads = model.audio_embeddings = model.local_text_lm_head = None + model._audio_embed = lambda codes: codes[:, :1].expand(-1, 4).to(dtype) / 8 + model.force_stop = False + + def frame(hidden, *_args, generator=None, temperature=None, top_k=None, top_p=None, **_kwargs): + assert (temperature, top_k, top_p) == (1.7, 25, 0.8) + noise = torch.randint(0, 8, (hidden.shape[0], 2), device=device, generator=generator) + codes = (hidden[:, :2].mul(16).long() + noise) % 8 + return torch.full((hidden.shape[0],), not model.force_stop, device=device), codes + + model.local_transformer = SimpleNamespace(generate_frame=frame) + return model + + +def _state(cls, device): + state = object.__new__(cls) + state.model = _model(device) + state.device, state.dtype = device, state.model.model.embed_tokens.weight.dtype + state.has_preprocess = state.has_postprocess = state.have_multimodal_outputs = True + state.scheduler_config = SimpleNamespace(max_num_seqs=5) + state.vllm_config = SimpleNamespace() + state.intermediate_buffer = OmniIntermediateBuffer(5) + state._static_inputs_embeds = torch.zeros(16, 4, device=device, dtype=state.dtype) + state._mtp_input_ids = torch.zeros(5, device=device, dtype=torch.long) + state._mtp_offsets = torch.zeros(5, device=device, dtype=torch.long) + for name in ("_mtp_input_embeds", "_mtp_hidden", "_mtp_text_step"): + setattr(state, name, torch.zeros(5, 4, device=device, dtype=state.dtype)) + state._mtp_runner = state._mtp_sample_uniforms = None + state._mtp_generators = {} + state._stream_pos = {} + state._stream_decode_event = None + state._eager_state = EagerMTPState(state) + if cls is MossLocalModelState: + state._init_slot_buffers(5, 4, device, state.dtype) + return state + + +def _admit(state, slot, name, seed): + params = SimpleNamespace(extra_args={"tts_local_seed": seed} if seed is not None else {}, seed=None) + with patch.object(DefaultModelState, "add_request", return_value=None): + state.add_request(slot, SimpleNamespace(req_id=name, mm_features=[], sampling_params=params)) + state.intermediate_buffer.buffers[slot]["codes"] = {"ref": torch.tensor([[1, 2], [3, 4], [5, 6], [2, 3]])} + + +def _batch(device, slots, counts): + starts = np.array([0, *np.cumsum(counts)], dtype=np.int32) + return SimpleNamespace( + idx_mapping_np=np.array(slots), + idx_mapping=torch.tensor(slots, device=device), + num_reqs=len(slots), + num_tokens=int(starts[-1]), + num_scheduled_tokens=np.array(counts), + query_start_loc_np=starts, + query_start_loc=torch.tensor(starts, device=device), + # Upstream MRV2 uses int32 token buffers, unlike the MTP graph inputs. + input_ids=torch.full((int(starts[-1]),), 2, dtype=torch.int32, device=device), + ) + + +def _step(state, batch, req_states, dispatcher=None): + inputs = {"input_ids": batch.input_ids, "inputs_embeds": state._static_inputs_embeds[: batch.num_tokens]} + with ( + patch("vllm.forward_context.set_forward_context", return_value=nullcontext()), + patch( + "vllm_omni.model_executor.models.moss_tts.local_model_state.set_forward_context", return_value=nullcontext() + ), + torch.inference_mode(), + ): + state.run_preprocess(batch, inputs, req_states, dispatcher) + hidden = inputs["inputs_embeds"] + 0.25 + state.run_postprocess(hidden, batch) + _, payload = state.postprocess_model_output(hidden, batch, req_states) + last = batch.query_start_loc[1:] - 1 + tokens = state.model.compute_logits(hidden.index_select(0, last)).argmax(-1) + result = inputs["inputs_embeds"].clone(), payload, tokens + hidden.fill_(-100) # next graph replay can overwrite its output immediately + return result + + +@pytest.mark.parametrize("seed", [None, 17]) +@pytest.mark.parametrize("batch_prefill", [False, True]) +def test_slot_matches_canonical_mixed_prefill_reorder_stop_and_reuse(device, seed, batch_prefill): + reference, candidate = (_state(cls, device) for cls in (OmniModelState, MossLocalModelState)) + candidate._batch_prefill = batch_prefill + for state in (reference, candidate): + _admit(state, 3, "long", seed) + _admit(state, 0, "short", seed) + computed, prompt = np.zeros(5, dtype=np.int32), np.zeros(5, dtype=np.int32) + prompt[3], prompt[0] = 4, 2 + req_states = SimpleNamespace(prompt_len=prompt, num_computed_tokens=computed) + saved = [] + for index, (slots, counts) in enumerate( + [ + ([3, 0], [2, 2]), + ([0, 3], [1, 1]), + ([3, 0], [1, 1]), + ([3, 0], [1, 1]), + ([0, 3], [1, 1]), + ([3, 0], [1, 1]), + ] + ): + if index == 4: + # Cancellation/removal followed by reuse must not retain hidden, + # reference offset, control, or the old request's RNG stream. + for state in (reference, candidate): + state.remove_request("short") + assert "short" not in state._mtp_generators + _admit(state, 0, "replacement", seed) + computed[0], prompt[0] = 0, 1 + outputs = [] + for state in (reference, candidate): + state.model.force_stop = index == 3 + torch.manual_seed(91 + index) + outputs.append(_step(state, _batch(device, slots, counts), req_states)) + a, b = outputs + torch.testing.assert_close(a[0], b[0], rtol=0, atol=0) + torch.testing.assert_close(a[2], b[2], rtol=0, atol=0) + if index == 3: + assert b[2].tolist() == [3, 3] + assert a[1].keys() == b[1].keys() + for ar, br in zip(a[1].get("codes", {}).get("audio", []), b[1].get("codes", {}).get("audio", [])): + torch.testing.assert_close(ar, br, rtol=0, atol=0) + saved.append((br, br.clone())) + for slot, count in zip(slots, counts): + computed[slot] += count + old = reference.intermediate_buffer.buffers[slot] + new = candidate.intermediate_buffer.buffers[slot] + assert old["ref_offset"] == new["ref_offset"] + torch.testing.assert_close(candidate._hidden_pool[slot], old["hidden_states"]["last"], rtol=0, atol=0) + assert "hidden_states" not in new and "audio_codes" not in new + candidate._codes_pool.fill_(8) + candidate._hidden_pool.zero_() + for current, snapshot in saved: + torch.testing.assert_close(current, snapshot, rtol=0, atol=0) + + +def test_slot_honors_external_stopping_control(device): + states = [_state(cls, device) for cls in (OmniModelState, MossLocalModelState)] + for state in states: + _admit(state, 0, "stop", None) + req = SimpleNamespace(prompt_len=np.ones(5), num_computed_tokens=np.zeros(5)) + _step(state, _batch(device, [0], [1]), req) + state.intermediate_buffer.buffers[0]["audio_state"]["is_stopping"] = True + req.num_computed_tokens[:] = 1 + output = _step(state, _batch(device, [0], [1]), req) + assert output[2].tolist() == [3] + assert output[1]["codes"]["audio"][0].tolist() == [[8, 8]] + + +@pytest.mark.cuda +def test_slot_mtp_graph_padding_reorder_and_owned_output(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + states = [_state(cls, torch.device("cuda")) for cls in (OmniModelState, MossLocalModelState)] + for state in states: + + def frame(hidden, *_args, **_kwargs): + return hidden[:, 0] < 3, hidden[:, :2].mul(16).long() % 8 + + state.model.local_transformer.generate_frame = frame + for slot in [0, 3, 4]: + _admit(state, slot, f"req{slot}", None) + candidate = states[1] + buffers = [ + candidate._mtp_input_ids[:4], + candidate._mtp_input_embeds[:4], + candidate._mtp_hidden[:4], + candidate._mtp_text_step[:4], + ] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream), torch.inference_mode(): + for _ in range(3): + candidate.model.mtp(*buffers) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph), torch.inference_mode(): + graph_output = candidate.model.mtp(*buffers) + + def replay(*args, **kwargs): + assert all(a.data_ptr() == b.data_ptr() for a, b in zip(args, buffers)) + assert args[0].shape[0] == 4 + graph.replay() + return graph_output + + candidate._mtp_runner = replay + candidate._is_mtp_graph_runner = lambda: True + from vllm.config import CUDAGraphMode + + def dispatcher(bsz): + return SimpleNamespace(num_tokens=4, cg_mode=CUDAGraphMode.FULL) + + req = SimpleNamespace(prompt_len=np.ones(5, dtype=np.int32), num_computed_tokens=np.zeros(5, dtype=np.int32)) + retained = [] + for slots in [[0, 3, 4], [3, 0], [4, 0, 3], [3, 4]]: + out = [ + _step( + state, + _batch(torch.device("cuda"), slots, [1] * len(slots)), + req, + dispatcher if state is candidate else None, + ) + for state in states + ] + torch.testing.assert_close(out[0][0], out[1][0], rtol=0, atol=0) + torch.testing.assert_close(out[0][2], out[1][2], rtol=0, atol=0) + for a, b in zip(out[0][1].get("codes", {}).get("audio", []), out[1][1].get("codes", {}).get("audio", [])): + torch.testing.assert_close(a, b, rtol=0, atol=0) + retained.append((b, b.clone())) + req.num_computed_tokens[slots] += 1 + with torch.inference_mode(): + for tensor in graph_output: + tensor.fill_(7) + for tensor, copy in retained: + torch.testing.assert_close(tensor, copy, rtol=0, atol=0) + + +@pytest.mark.cpu +def test_factory_requires_explicit_class_capability_and_valid_state(monkeypatch): + monkeypatch.setattr(OmniModelState, "__init__", lambda *args: None) + dynamic = MagicMock(has_preprocess=True, create_omni_model_state=None) + assert isinstance(init_omni_model_state(None, dynamic, None, torch.device("cpu")), OmniModelState) + dynamic.create_mrv2_model_state.assert_not_called() + + class Invalid: + has_preprocess = True + + def create_mrv2_model_state(self, *args): + return object() + + with pytest.raises(TypeError, match="must return a ModelState"): + init_omni_model_state(None, Invalid(), None, torch.device("cpu")) + + +@pytest.mark.cpu +@pytest.mark.parametrize("enabled", [False, True]) +def test_local_factory_default_and_option(monkeypatch, enabled): + monkeypatch.setattr(OmniModelState, "__init__", lambda *args: None) + monkeypatch.setattr(MossLocalModelState, "__init__", lambda *args: None) + model = _model(torch.device("cpu")) + model.config = SimpleNamespace(mrv2_gpu_slot_state=enabled) + state = init_omni_model_state(None, model, None, torch.device("cpu")) + assert type(state) is (MossLocalModelState if enabled else OmniModelState) diff --git a/tests/model_executor/models/moss_tts/test_local_prefill_batch.py b/tests/model_executor/models/moss_tts/test_local_prefill_batch.py new file mode 100644 index 000000000..8b6856f4a --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_local_prefill_batch.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Reference batching preserves coordinates and asynchronous source ownership.""" + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from tests.model_executor.models.moss_tts.test_local_model_state import _batch, _state +from vllm_omni.model_executor.models.moss_tts.local_model_state import MossLocalModelState +from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import MossTTSLocalTalkerForGeneration + +pytestmark = pytest.mark.core_model + + +@pytest.mark.cpu +@pytest.mark.parametrize("flat", [False, True]) +def test_prefill_reuses_text_embeddings_and_noncontiguous_reference_offsets(mocker, flat): + state = _state(MossLocalModelState, torch.device("cpu")) + state._batch_prefill = True + codes = torch.tensor([[1, 2], [3, 4], [5, 6], [2, 3]]) + for slot, offset in [(3, 1), (0, 2)]: + state.intermediate_buffer.buffers[slot] = { + "req_id": str(slot), + "codes": {"ref": codes.flatten() if flat else codes}, + "ref_offset": offset, + } + batch = _batch(torch.device("cpu"), [3, 0], [2, 1]) + req = SimpleNamespace(prompt_len=np.full(5, 8), num_computed_tokens=np.zeros(5)) + embeds = state._static_inputs_embeds[:3] + expected = state.model.embed_input_ids(batch.input_ids) + state.model._audio_embed(codes[[1, 2, 2]]) + embed = mocker.spy(state.model.model.embed_tokens, "forward") + with torch.inference_mode(): + state.run_preprocess(batch, {"input_ids": batch.input_ids, "inputs_embeds": embeds}, req) + torch.testing.assert_close(embeds, expected, rtol=0, atol=0) + assert embed.call_count == 1 + assert state.intermediate_buffer.buffers[3]["ref_offset"] == 3 + assert state.intermediate_buffer.buffers[0]["ref_offset"] == 3 + + +@pytest.mark.cuda +def test_pending_reference_uploads_survive_staging_reuse_and_overflow(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + state = _state(MossLocalModelState, torch.device("cuda")) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + outputs = [] + with torch.cuda.stream(stream), torch.inference_mode(): + # Keep the first H2D pending while CPU prepares all subsequent batches. + torch.cuda._sleep(100_000_000) + for value in range(1, 8): + embeds = torch.zeros(4, 4, device="cuda", dtype=state.dtype) + codes = torch.full((2, 2), value) + state._apply_reference_batch(embeds, [(1, codes)]) + outputs.append((embeds, value)) + codes.fill_(-1) + stream.synchronize() + assert len(state._prefill_staging) <= 3 + for output, value in outputs: + expected = torch.zeros_like(output) + expected[1:3] = value / 8 + torch.testing.assert_close(output, expected, rtol=0, atol=0) + + +@pytest.mark.cuda +def test_batched_real_audio_embedding_matches_scalar_bfloat16_reduction(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + state = _state(MossLocalModelState, torch.device("cuda")) + model = state.model + model.n_vq = 12 + model.audio_vocab_size = model.audio_pad_token_id = 1024 + model._stacked_audio_emb_w = torch.randn(12, 1024, 3072, device="cuda", dtype=torch.bfloat16) + model._audio_embed = MossTTSLocalTalkerForGeneration._audio_embed.__get__(model) + state._static_inputs_embeds = torch.zeros(512, 3072, device="cuda", dtype=torch.bfloat16) + references = [(start, torch.randint(0, 1025, (count, 12))) for start, count in [(0, 169), (201, 11), (220, 257)]] + with torch.inference_mode(): + actual = torch.randn_like(state._static_inputs_embeds) + expected = actual.clone() + for start, codes in references: + expected[start : start + len(codes)] += model._audio_embed(codes.to("cuda")) + state._apply_reference_batch(actual, references) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/tests/model_executor/models/moss_tts/test_local_short_attention.py b/tests/model_executor/models/moss_tts/test_local_short_attention.py new file mode 100644 index 000000000..212d32f36 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_local_short_attention.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import pytest +import torch +import torch.nn.functional as F + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cuda, + pytest.mark.tts, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), +] + + +@pytest.mark.parametrize("batch,heads,dim", [(1, 4, 64), (16, 32, 80), (128, 32, 80), (3, 8, 128)]) +def test_all_depth_positions_with_strided_cache_and_dirty_future(batch, heads, dim): + from vllm_omni.model_executor.models.moss_tts.local_short_attention import local_short_attention + + torch.manual_seed(43) + query = torch.randn(batch, heads, 1, dim, device="cuda", dtype=torch.bfloat16) + key = torch.randn(batch, heads, 12, dim, device="cuda", dtype=torch.bfloat16) + value = torch.randn_like(key) + for position in range(12): + # Non-contiguous prefixes retain the stride of the full 12-slot cache. + k, v = key[:, :, : position + 1], value[:, :, : position + 1] + before_k, before_v = key.clone(), value.clone() + expected = F.scaled_dot_product_attention(query, k, v) + actual = local_short_attention(query, k, v) + assert torch.equal(key, before_k) and torch.equal(value, before_v) + if position == 0: + assert torch.equal(actual, v) + torch.testing.assert_close(actual, expected, atol=0.015625, rtol=0.01) + # Future cache contents must never affect an earlier position. + key[:, :, position + 1 :] += 100 + value[:, :, position + 1 :] -= 100 + assert torch.equal(local_short_attention(query, k, v), actual) + key.copy_(before_k) + value.copy_(before_v) + + +def test_compile_and_graph_replay_reads_fresh_inputs(): + from vllm_omni.model_executor.models.moss_tts.local_short_attention import local_short_attention + + query = torch.randn(16, 32, 1, 96, device="cuda", dtype=torch.bfloat16) + key = torch.randn(16, 32, 12, 96, device="cuda", dtype=torch.bfloat16) + value = torch.randn_like(key) + compiled = torch.compile(local_short_attention, fullgraph=True) + for _ in range(3): + compiled(query, key, value) + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + result = compiled(query, key, value) + for _ in range(3): + query.normal_() + key.normal_() + value.normal_() + graph.replay() + torch.testing.assert_close(result, F.scaled_dot_product_attention(query, key, value), atol=0.015625, rtol=0.01) + + +def test_depth_block_cache_overwrite_and_compiled_positions(monkeypatch): + from transformers import GPT2Config + + from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import ( + MossTTSLocalDepthTransformer, + ) + + monkeypatch.setenv("VLLM_OMNI_MOSS_LOCAL_SHORT_ATTN", "1") + model = MossTTSLocalDepthTransformer(GPT2Config(n_embd=320, n_head=4, n_inner=640)).cuda().bfloat16().eval() + attn = model.h[0].attn + attn.prepare_rope_cache(12, torch.device("cuda"), torch.bfloat16) + cache = tuple(torch.empty(4, 4, 12, 80, device="cuda", dtype=torch.bfloat16) for _ in range(2)) + forward = torch.compile(model._forward_prefix, fullgraph=True, dynamic=True) + with torch.inference_mode(): + for _ in range(2): + inputs = torch.randn(4, 12, 320, device="cuda", dtype=torch.bfloat16) + # The second frame overwrites the previous frame's K/V before use. + for position in range(12): + expected = model._forward_prefix(inputs[:, : position + 1])[:, -1:] + actual = forward(inputs[:, position : position + 1], cache, position) + torch.testing.assert_close(actual, expected, atol=0.03125, rtol=0.02) diff --git a/tests/model_executor/models/moss_tts/test_reference_encoder.py b/tests/model_executor/models/moss_tts/test_reference_encoder.py index cfd8fb5c2..06efaaf25 100644 --- a/tests/model_executor/models/moss_tts/test_reference_encoder.py +++ b/tests/model_executor/models/moss_tts/test_reference_encoder.py @@ -598,3 +598,236 @@ def test_reference_resampler_is_cached_and_matches_functional_resample(): assert info.misses == 1 and info.hits == 1 finally: _reference_resampler.cache_clear() + + +# --------------------------------------------------------------------------- # +# Inline data: URI fast path # +# --------------------------------------------------------------------------- # + + +class _DigestAudio(_FakeAudio): + """Resolver whose key for data: URIs is the SHA-1 digest, like serving.""" + + def __init__(self, *, digest_keys=True): + super().__init__() + self._digest_keys = digest_keys + + async def resolve(self, ref_str): + wav_list, sr, _ = await super().resolve(ref_str) + if not self._digest_keys: + return wav_list, sr, "rk:" + ref_str + import hashlib + + return wav_list, sr, hashlib.sha1(ref_str.encode("utf-8")).hexdigest() + + def artifact_key(self, cache_key): + return None + + +async def test_inline_reference_skips_repeat_resolution(make_encoder): + proc, audio = _FakeProcessor(), _DigestAudio() + ref = "data:audio/wav;base64,QUJD" + audio.register(ref, 5) + enc = make_encoder(proc) + + first, key1 = await enc.encode(ref, resolve_ref_audio=audio.resolve, get_artifact_key=audio.artifact_key) + first.fill_(0) # callers own their copy + second, key2 = await enc.encode(ref, resolve_ref_audio=audio.resolve, get_artifact_key=audio.artifact_key) + + assert audio.resolve_calls == [ref] + assert key1 == key2 + assert torch.equal(second, torch.full((3, _N_VQ), 5, dtype=torch.long)) + + +async def test_inline_index_requires_matching_resolve_key(make_encoder): + proc, audio = _FakeProcessor(), _DigestAudio(digest_keys=False) + ref = "data:audio/wav;base64,REVG" + audio.register(ref, 6) + enc = make_encoder(proc) + + for _ in range(2): + await enc.encode(ref, resolve_ref_audio=audio.resolve, get_artifact_key=audio.artifact_key) + assert audio.resolve_calls == [ref, ref] + + +async def test_inline_reference_shared_across_encoders(make_encoder, tmp_path_factory, monkeypatch): + shared_dir = tmp_path_factory.mktemp("ref") + monkeypatch.setenv("VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR", str(shared_dir)) + ref = "data:audio/wav;base64,R0hJ" + first_audio, second_audio = _DigestAudio(), _DigestAudio() + first_audio.register(ref, 9) + second_audio.register(ref, 9) + first, second = make_encoder(_FakeProcessor()), make_encoder(_FakeProcessor()) + + a, key_a = await first.encode(ref, resolve_ref_audio=first_audio.resolve, get_artifact_key=first_audio.artifact_key) + b, key_b = await second.encode( + ref, resolve_ref_audio=second_audio.resolve, get_artifact_key=second_audio.artifact_key + ) + + assert first_audio.resolve_calls == [ref] and second_audio.resolve_calls == [] + assert key_a == key_b and torch.equal(a, b) + + +# --------------------------------------------------------------------------- # +# Reference-encoder CUDA graphs: selection and processor-equivalent prep # +# --------------------------------------------------------------------------- # + + +class _GraphsStub: + def __init__(self, result): + self.result = result + self.calls: list[list[torch.Tensor]] = [] + + def encode(self, wavs): + self.calls.append(wavs) + return self.result + + +class _LoudnessProcessor(_FakeProcessor): + @staticmethod + def loudness_normalize(wav): + return wav * 2.0 + + +async def test_encode_prepared_prefers_graphs_and_falls_back_to_processor(make_encoder): + proc = _LoudnessProcessor() + enc = make_encoder(proc) + graphs_codes = [torch.full((3, _N_VQ), 7, dtype=torch.long)] + enc._graphs_enabled, enc._graphs = True, _GraphsStub(graphs_codes) + prepared = [torch.ones(1, 100)] + assert enc._encode_prepared(prepared) is graphs_codes and proc.attempt_sizes == [] + enc._graphs = _GraphsStub([None]) # no graph fits this clip + out = enc._encode_prepared(prepared) + assert proc.attempt_sizes == [1] and out[0].tolist() == [[1] * _N_VQ] * 3 + + +async def test_processor_prepare_matches_processor_channel_and_loudness_handling(make_encoder): + enc = make_encoder(_LoudnessProcessor()) + mono = enc._processor_prepare(torch.ones(100)) + assert mono.shape == (2, 100) and torch.equal(mono, torch.full((2, 100), 2.0)) + many = enc._processor_prepare(torch.arange(3 * 4, dtype=torch.float32).reshape(3, 4)) + assert torch.equal(many, torch.arange(8, dtype=torch.float32).reshape(2, 4) * 2.0) + + +async def test_processor_prepare_resamples_to_the_tokenizer_rate(make_encoder): + torchaudio = pytest.importorskip("torchaudio") + proc = _LoudnessProcessor(sampling_rate=48000) + enc = make_encoder(proc) # works at 24 kHz, the Local-v1.5 reference rate + wav = torch.randn(1, 2400) + expected = torchaudio.functional.resample(wav.repeat(2, 1), 24000, 48000) * 2.0 + # The encoder resamples the mono clip before duplicating it. Some CPU conv + # kernels round a one-row batch differently from a two-row one, so the two + # orders agree to float32 precision, not bit for bit. + torch.testing.assert_close(enc._processor_prepare(wav), expected, rtol=0, atol=1e-5) + + +@pytest.mark.parametrize("role", ["", "host", "client"]) +async def test_prepare_captures_graphs_or_starts_the_shared_host(make_encoder, monkeypatch, tmp_path, role): + from types import SimpleNamespace + + from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod + + enc = make_encoder(_FakeProcessor()) + enc._shares_encoder, enc._shared_codes_dir = True, str(tmp_path) + monkeypatch.setattr(re_mod, "shared_encoder_role", lambda: role) + captured, hosts, ready = [], [], [] + monkeypatch.setattr(enc, "_reference_graphs", lambda: captured.append(True)) + monkeypatch.setattr(enc, "_host_workers", lambda: []) + helper = SimpleNamespace(wait_until_ready=lambda: ready.append(True)) + + def make_host(*args, **kwargs): + hosts.append(args) + return helper + + monkeypatch.setattr(re_mod, "SharedReferenceEncoderHost", make_host) + monkeypatch.setattr(re_mod, "SharedReferenceEncoderClient", lambda *args: helper) + enc.prepare() + enc.prepare() + # Every API process waits for the shared host before startup completes. + assert (len(captured), len(hosts)) == {"": (2, 0), "host": (0, 1), "client": (0, 0)}[role] + assert len(ready) == (2 if role else 0) + + +async def test_clips_reach_the_encoder_at_the_tokenizer_rate(make_encoder, monkeypatch): + torchaudio = pytest.importorskip("torchaudio") + proc = _LoudnessProcessor(sampling_rate=48000) + enc = make_encoder(proc) # works at 24 kHz + sent = [] + + def record_prepared(prepared): + sent.extend(prepared) + return [None] * len(prepared) + + monkeypatch.setattr(enc, "_encode_shared_or_local", record_prepared) + wav = [float(v) for v in torch.randn(2400)] + enc._encode_batch_sync([(wav, _SR)]) + # Resampled where the request arrived (a shared encoder receives it this way); still one channel. + expected = torchaudio.functional.resample(torch.tensor(wav)[None], 24000, 48000) + torch.testing.assert_close(sent[0], expected, rtol=0, atol=0) + # The eager path is told the clip is already at the tokenizer rate. + calls = [] + + def record_encoding(wavs, sampling_rate, n_vq=None): + calls.append(sampling_rate) + return [None] + + monkeypatch.setattr(proc, "encode_audios_from_wav", record_encoding) + enc._graphs_enabled = False + enc._encode_prepared(sent) + assert calls == [48000] + + +async def test_batcher_runs_up_to_max_inflight_batches_at_once(): + inside = threading.Barrier(2, timeout=5) + seen: list[int] = [] + + def encode_batch(payload): + inside.wait() # two batches are encoding at the same time + seen.append(len(payload)) + return [torch.zeros(1) for _ in payload] + + batcher = _RefEncodeBatcher(encode_batch, window_ms=0.0, max_batch=8, max_inflight=2) + try: + first = asyncio.ensure_future(batcher.submit([1.0], _SR)) + await asyncio.sleep(0.05) # the first batch is in flight before the second clip arrives + await asyncio.gather(first, batcher.submit([2.0], _SR)) + finally: + await batcher.aclose() + assert seen == [1, 1] + + +def test_reference_stream_priority(monkeypatch): + from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod + + monkeypatch.delenv(re_mod._REF_ENCODE_PRIORITY_ENV, raising=False) + assert re_mod.reference_stream_priority() == 0 + monkeypatch.setenv(re_mod._REF_ENCODE_PRIORITY_ENV, "high") + expected = torch.cuda.Stream.priority_range()[1] if torch.cuda.is_available() else 0 + assert re_mod.reference_stream_priority() == expected + + +async def test_device_prep_resamples_on_the_tokenizer_device(make_encoder, monkeypatch): + torchaudio = pytest.importorskip("torchaudio") + from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod + + proc = _LoudnessProcessor(sampling_rate=48000) + monkeypatch.setattr(proc, "audio_tokenizer", torch.nn.Linear(1, 1), raising=False) # CPU placement + enc = make_encoder(proc) + enc._gpu_prep = True + wav = torch.randn(1, 2400) + clip = enc._encoder_input(wav) + assert clip.shape == (1, 2400) # left at the working rate for the encoder's device + expected = torchaudio.functional.resample(wav, 24000, 48000).repeat(2, 1) * 2.0 + torch.testing.assert_close(enc._finish_prepare(clip), expected, rtol=1e-5, atol=1e-6) + # The eager fallback is told the clip's own rate. + calls = [] + + def record_encoding(wavs, sampling_rate, n_vq=None): + calls.append(sampling_rate) + return [None] + + monkeypatch.setattr(proc, "encode_audios_from_wav", record_encoding) + enc._graphs_enabled = False + enc._encode_prepared([clip]) + assert calls == [24000] + assert re_mod._REF_ENCODE_GPU_PREP_ENV == "VLLM_OMNI_MOSS_REF_GPU_PREP" diff --git a/tests/model_executor/models/moss_tts/test_reference_encoder_graphs.py b/tests/model_executor/models/moss_tts/test_reference_encoder_graphs.py new file mode 100644 index 000000000..74969005e --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_reference_encoder_graphs.py @@ -0,0 +1,320 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Reference-encoder CUDA graphs reproduce the tokenizer's eager encode.""" + +from contextlib import nullcontext + +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +from vllm_omni.model_executor.models.moss_tts.reference_encoder_graphs import ( + MossReferenceEncoderGraphs, + _cached_reference_rope, + _reference_rope_factors, + _validate_unmasked_reference_encoder, + _windowed_attention, + install_windowed_attention, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + +FRAME, N_VQ = 8, 3 + + +class _Stage(nn.Module): + """Causal per-frame stage: frames never mix, like a causal encoder layer.""" + + def __init__(self, dim_in, dim_out): + super().__init__() + self.proj = nn.Linear(dim_in, dim_out) + + def forward(self, x, lengths): + return torch.tanh(self.proj(x)), lengths + + +class _Quantizer(nn.Module): + def forward(self, hidden, lengths, n_q): + # hidden (B, frames, D): codes (n_q, B, frames) from per-frame argmax slices. + codes = torch.stack([hidden[..., q::n_q].argmax(-1) for q in range(n_q)]) + return hidden, codes, lengths + + +class _Tokenizer(nn.Module): + number_channels = 2 + downsample_rate = FRAME + sampling_rate = 64 + + def __init__(self): + super().__init__() + self.encoder = nn.ModuleList([_Stage(FRAME * 2, 16), _Stage(16, 12)]) + self.quantizer = _Quantizer() + + def _flatten_channels_for_codec(self, audio, lengths): + batch, channels, samples = audio.shape + frames = samples // FRAME + values = audio.reshape(batch, channels, frames, FRAME).permute(0, 2, 1, 3).reshape(batch, frames, -1) + return values, lengths // FRAME + + def _codec_inference_autocast(self): + return nullcontext() + + def eager(self, wavs): + out = [] + for wav in wavs: + values, lengths = self._flatten_channels_for_codec(wav[None], torch.tensor([wav.shape[-1]])) + for module in self.encoder: + values, lengths = module(values, lengths) + _, codes, _ = self.quantizer(values, lengths, N_VQ) + out.append(codes[:, 0, : wav.shape[-1] // FRAME].transpose(0, 1).long()) + return out + + +@pytest.fixture +def cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + return torch.device("cuda", torch.accelerator.current_device_index()) + + +def test_cached_rope_reuses_longest_table_without_modifying_inputs(cuda): + from types import SimpleNamespace + + module = SimpleNamespace(rope=SimpleNamespace(max_period=10000.0)) + cache: dict[tuple[int, float, torch.device], tuple[torch.Tensor, torch.Tensor]] = {} + torch.manual_seed(17) + # Prime once; all subsequent sequences and batches use an immutable prefix. + q = torch.randn(1, 3, 128, 32, device=cuda, dtype=torch.bfloat16) + k = torch.randn_like(q) + _cached_reference_rope(module, q, k, cache) + stored = next(iter(cache.values())) + for batch, length in [(1, 17), (4, 64), (2, 128)]: + q = torch.randn(batch, 3, length, 32, device=cuda, dtype=torch.bfloat16) + k = torch.randn_like(q) + before = (q.clone(), k.clone()) + got = _cached_reference_rope(module, q, k, cache) + cos, sin = _reference_rope_factors(length, 32, 10000.0, cuda) + torch.testing.assert_close(stored[0][:length], cos, rtol=0, atol=0) + torch.testing.assert_close(stored[1][:length], sin, rtol=0, atol=0) + for value, source, saved in zip(got, (q, k), before): + real, imag = source.float().reshape(batch, 3, length, 16, 2).unbind(-1) + expected = torch.stack((real * cos - imag * sin, real * sin + imag * cos), -1).to(source.dtype) + torch.testing.assert_close(value, expected.reshape_as(source), rtol=0, atol=0) + torch.testing.assert_close(source, saved, rtol=0, atol=0) + assert len(cache) == 1 + assert next(iter(cache.values())) is stored + + +def test_cached_rope_distinguishes_period_and_head_width(cuda): + from types import SimpleNamespace + + cache: dict[tuple[int, float, torch.device], tuple[torch.Tensor, torch.Tensor]] = {} + for width, period in [(32, 10000.0), (64, 10000.0), (32, 1000.0)]: + module = SimpleNamespace(rope=SimpleNamespace(max_period=period)) + q = torch.randn(1, 2, 32, width, device=cuda, dtype=torch.bfloat16) + _cached_reference_rope(module, q, q, cache) + assert len(cache) == 3 + + +@pytest.mark.parametrize("batched_transfer", [False, True]) +def test_graph_codes_match_eager_per_clip_and_fall_back_outside_buckets(cuda, batched_transfer): + torch.manual_seed(0) + tokenizer = _Tokenizer().to(cuda) + # Buckets of 2 s and 4 s at 64 Hz: 128 and 256 samples. + graphs = MossReferenceEncoderGraphs( + tokenizer, n_vq=N_VQ, batch_sizes=(1, 2, 4), bucket_seconds=(2.0, 4.0), batched_transfer=batched_transfer + ) + graphs.capture() + assert graphs.captured == [(1, 128), (1, 256), (2, 128), (2, 256), (4, 128), (4, 256)] + assert graphs.max_batch == 4 + + for count, frames in [(1, [5]), (3, [16, 3, 9]), (4, [32, 1, 20, 17]), (2, [16, 16])]: + wavs = [torch.randn(2, n * FRAME) for n in frames] + codes = graphs.encode(wavs) + expected = tokenizer.eager([w.to(cuda) for w in wavs]) + assert len(codes) == count + for got, want in zip(codes, expected): + assert got.device.type == "cpu" and got.dtype == torch.long + torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0) + + # Longer than every bucket: left to the caller; the rest still use graphs. + long_clip, short_clip = torch.randn(2, 33 * FRAME), torch.randn(2, 4 * FRAME) + codes = graphs.encode([long_clip, short_clip]) + assert codes[0] is None + torch.testing.assert_close(codes[1], tokenizer.eager([short_clip.to(cuda)])[0].cpu(), rtol=0, atol=0) + + # Wider than every batch: split, longest first, and returned in input order. + wavs = [torch.randn(2, n * FRAME) for n in (3, 30, 7, 12, 1)] + for got, want in zip(graphs.encode(wavs), tokenizer.eager([w.to(cuda) for w in wavs])): + torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0) + + # Result ownership includes the degenerate single-frame shape. Reusing + # the pinned output slab must not overwrite a previous caller's codes. + previous = graphs.encode([torch.randn(2, FRAME)])[0] + snapshot = previous.clone() + graphs.encode([torch.randn(2, 16 * FRAME) for _ in range(4)]) + torch.testing.assert_close(previous, snapshot, rtol=0, atol=0) + + +def test_batch_audio_budget_limits_captured_graphs(cuda): + tokenizer = _Tokenizer().to(cuda) + graphs = MossReferenceEncoderGraphs( + tokenizer, n_vq=N_VQ, batch_sizes=(1, 2, 4), bucket_seconds=(2.0, 4.0), max_batch_seconds=8.0 + ) + graphs.capture() + # 4 x 4 s exceeds the 8 s budget; a single clip is always captured. + assert graphs.captured == [(1, 128), (1, 256), (2, 128), (2, 256), (4, 128)] + wavs = [torch.randn(2, 30 * FRAME) for _ in range(3)] + for got, want in zip(graphs.encode(wavs), tokenizer.eager([w.to(cuda) for w in wavs])): + torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0) + + +def test_finer_singleton_buckets_preserve_multiclip_batching(cuda): + tokenizer = _Tokenizer().to(cuda) + graphs = MossReferenceEncoderGraphs( + tokenizer, + n_vq=N_VQ, + batch_sizes=(1, 2, 4), + bucket_seconds=(2.0, 4.0), + singleton_bucket_seconds=(1.0, 3.0), + ) + graphs.capture() + assert (1, 64) in graphs.captured and (1, 192) in graphs.captured + assert (2, 64) not in graphs.captured and (4, 192) not in graphs.captured + replay = graphs._replay + shapes = [] + + def record(entry, wavs): + shapes.append(tuple(entry.audio.shape)) + return replay(entry, wavs) + + graphs._replay = record + one = torch.randn(2, 20 * FRAME) + out = graphs.encode([one]) + assert shapes == [(1, 2, 192)] + torch.testing.assert_close(out[0], tokenizer.eager([one.to(cuda)])[0].cpu(), rtol=0, atol=0) + shapes.clear() + wavs = [one, one[:, : 8 * FRAME], one[:, : 12 * FRAME]] + out = graphs.encode(wavs) + assert shapes == [(4, 2, 256)] + for got, want in zip(out, tokenizer.eager([w.to(cuda) for w in wavs])): + torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0) + + +def test_compiled_capture_matches_eager(cuda): + torch.manual_seed(2) + tokenizer = _Tokenizer().to(cuda) + graphs = MossReferenceEncoderGraphs( + tokenizer, n_vq=N_VQ, batch_sizes=(1, 4), bucket_seconds=(2.0, 4.0), compile_core=True + ) + graphs.capture() + assert graphs.captured == [(1, 128), (1, 256), (4, 128), (4, 256)] + wavs = [torch.randn(2, n * FRAME) for n in (5, 30, 12)] + for got, want in zip(graphs.encode(wavs), tokenizer.eager([w.to(cuda) for w in wavs])): + torch.testing.assert_close(got, want.cpu(), rtol=0, atol=0) + + +class _Attention(nn.Module): + """The attention surface the tokenizer's MHA exposes to the windowed kernel.""" + + def __init__(self, embed_dim=128, heads=2, context=5): + super().__init__() + self.embed_dim, self.num_heads, self.context, self.causal = embed_dim, heads, context, True + self.in_proj = nn.Linear(embed_dim, 3 * embed_dim, bias=False) + + def _project_qkv(self, x): + b, t, _ = x.shape + return self.in_proj(x).reshape(b, t, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4).unbind(0) + + def _apply_dense_rope(self, q, k): + return q, k + + def _forward_non_streaming_sdpa(self, x, input_lengths): + # Reference semantics: causal, distance < context, padded rows zeroed. + b, t, _ = x.shape + q, k, v = self._project_qkv(x) + pos = torch.arange(t, device=x.device) + delta = pos.view(-1, 1) - pos.view(1, -1) + mask = (delta >= 0) & (delta < self.context) + out = F.scaled_dot_product_attention(q, k, v, mask[None, None]) + valid = (pos.view(1, t) < input_lengths.view(-1, 1)).view(b, 1, t, 1) + out = torch.where(valid, out, torch.zeros((), device=out.device, dtype=out.dtype)) + return out.transpose(1, 2).reshape(b, t, self.embed_dim) + + +@pytest.mark.parametrize("fa_version", [2, 3]) +def test_windowed_attention_matches_masked_sdpa(cuda, fa_version): + if fa_version == 3 and torch.cuda.get_device_capability(cuda)[0] != 9: + pytest.skip("FA3 requires Hopper") + torch.manual_seed(1) + model = nn.Module() + model.encoder = nn.ModuleList([_Attention().to(cuda, torch.bfloat16)]) + attention = model.encoder[0] + x = torch.randn(3, 40, 128, device=cuda, dtype=torch.bfloat16) + lengths = torch.tensor([40, 17, 3], device=cuda) + expected = attention._forward_non_streaming_sdpa(x, lengths) + assert install_windowed_attention(model, fa_version=fa_version) == 1 + actual = attention._forward_non_streaming_sdpa(x, lengths) + torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) + assert actual[1, 17:].count_nonzero() == 0 and actual[2, 3:].count_nonzero() == 0 + # Non-bf16 inputs keep the tokenizer's own SDPA path. + x32 = x.float() + attention.float() + torch.testing.assert_close( + attention._forward_non_streaming_sdpa(x32, lengths), expected.float(), rtol=3e-2, atol=3e-2 + ) + + +@pytest.mark.parametrize("fa_version", [2, 3]) +def test_unmasked_causal_queries_preserve_valid_prefix_and_ignore_future_padding(cuda, fa_version): + if fa_version == 3 and torch.cuda.get_device_capability(cuda)[0] != 9: + pytest.skip("FA3 requires Hopper") + torch.manual_seed(44) + attention = _Attention().to(cuda, torch.bfloat16) + x = torch.randn(4, 40, 128, device=cuda, dtype=torch.bfloat16) + lengths = torch.tensor([40, 17, 1, 0], device=cuda) + kwargs = dict(sdpa=attention._forward_non_streaming_sdpa, fa_version=fa_version) + with torch.no_grad(): + masked = _windowed_attention(attention, x, lengths, **kwargs) + unmasked = _windowed_attention(attention, x, lengths, skip_padded_query_mask=True, **kwargs) + poisoned = x.clone() + for row, length in enumerate([40, 17, 1, 0]): + poisoned[row, length:] = torch.randn_like(poisoned[row, length:]) * 30 + changed = _windowed_attention(attention, poisoned, lengths, skip_padded_query_mask=True, **kwargs) + for row, length in enumerate([40, 17, 1, 0]): + torch.testing.assert_close(unmasked[row, :length], masked[row, :length], rtol=0, atol=0) + torch.testing.assert_close(changed[row, :length], masked[row, :length], rtol=0, atol=0) + assert masked[row, length:].count_nonzero() == 0 + assert unmasked[1, 17:].count_nonzero() > 0 + # The FP32 fallback retains its original masking contract. + attention.float() + got = _windowed_attention(attention, x.float(), lengths, skip_padded_query_mask=True, **kwargs) + expected = attention._forward_non_streaming_sdpa(x.float(), lengths) + torch.testing.assert_close(got, expected, rtol=0, atol=0) + + +def test_unmasked_queries_reject_unknown_encoder_before_patching(): + tokenizer = _Tokenizer() + with pytest.raises(ValueError, match="patch-downsample"): + install_windowed_attention(tokenizer, skip_padded_query_mask=True) + + +@pytest.mark.parametrize("bad_part", ["upsample", "noncausal", "quantizer", "convolution"]) +def test_unmasked_queries_reject_structures_that_can_read_invalid_rows(bad_part): + from types import SimpleNamespace + + patch = type("MossAudioTokenizerPatchedPretransform", (), {})() + patch.is_downsample = bad_part != "upsample" + projected = type("MossAudioTokenizerProjectedTransformer", (), {})() + projected.transformer = SimpleNamespace( + layers=[SimpleNamespace(self_attn=SimpleNamespace(causal=bad_part != "noncausal"))] + ) + quantizer_cls = type( + "MossAudioTokenizerResidualVQ" if bad_part != "quantizer" else "UnknownQuantizer", (nn.Module,), {} + ) + quantizer = quantizer_cls() + quantizer.proj = nn.Conv1d(2, 2, 3 if bad_part == "convolution" else 1) + tokenizer = SimpleNamespace(encoder=[patch, projected], quantizer=quantizer) + with pytest.raises(ValueError): + _validate_unmasked_reference_encoder(tokenizer) diff --git a/tests/model_executor/models/moss_tts/test_shared_reference_encoder.py b/tests/model_executor/models/moss_tts/test_shared_reference_encoder.py new file mode 100644 index 000000000..b8bd50184 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_shared_reference_encoder.py @@ -0,0 +1,339 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""API processes share one reference encoder through a Unix socket.""" + +import threading + +import pytest +import torch + +from vllm_omni.model_executor.models.moss_tts import shared_reference_encoder as shared + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _codes(wav): + # Deterministic per clip: frames from the clip length, values from its first sample. + return torch.full((wav.shape[-1] // 10, 4), int(wav.reshape(-1)[0]), dtype=torch.long) + + +class _Recorder: + def __init__(self, gate=None): + self.batches = [] + self.gate = gate + + def __call__(self, wavs): + if self.gate is not None: + self.gate.wait(5) + self.batches.append(len(wavs)) + return [RuntimeError("bad clip") if int(w.reshape(-1)[0]) < 0 else _codes(w) for w in wavs] + + +@pytest.fixture +def host_dir(tmp_path_factory, monkeypatch): + monkeypatch.setattr(shared, "_host_lock_handle", None) + # AF_UNIX limits the full socket path, including pytest directory names. + return str(tmp_path_factory.mktemp("ref")) + + +def test_only_one_process_per_directory_hosts(host_dir): + assert shared.elect_host(host_dir) is True + assert shared.elect_host(host_dir) is False # a second open file description cannot take the lock + + +def test_close_releases_queued_and_active_callers(host_dir): + entered, release = threading.Event(), threading.Event() + + def encode(wavs): + entered.set() + assert release.wait(5) + return [_codes(wav) for wav in wavs] + + host = shared.SharedReferenceEncoderHost(host_dir, [(encode, None)]) + results: dict[str, list[torch.Tensor | BaseException]] = {} + threads = [ + threading.Thread(target=lambda key=key: results.setdefault(key, host.encode([torch.ones(1, 20)])), daemon=True) + for key in ("active", "queued") + ] + try: + threads[0].start() + assert entered.wait(2) + threads[1].start() + with host._ready: + assert host._ready.wait_for(lambda: len(host._pending) == 1, timeout=2) + host.close() + for thread in threads: + thread.join(1) + assert not thread.is_alive(), "closed encoder left a request waiting" + assert all(isinstance(results[key][0], RuntimeError) for key in results) + finally: + release.set() + host.close() + for thread in threads: + if thread.ident is not None: + thread.join(1) + + +def test_worker_result_count_mismatch_fails_entire_batch(host_dir): + host = shared.SharedReferenceEncoderHost(host_dir, [(lambda wavs: [_codes(wavs[0])], None)]) + results = [] + caller = threading.Thread(target=lambda: results.extend(host.encode([torch.ones(1, 20)] * 2)), daemon=True) + try: + caller.start() + caller.join(2) + assert not caller.is_alive(), "missing encoder result left a request waiting" + assert len(results) == 2 + assert all(isinstance(result, RuntimeError) for result in results) + finally: + host.close() + caller.join(1) + + +def test_client_round_trip_and_per_clip_errors(host_dir): + host = shared.SharedReferenceEncoderHost(host_dir, [(_Recorder(), None)]) + try: + client = shared.SharedReferenceEncoderClient(host_dir) + wavs = [torch.full((1, 40), 3.0), torch.full((1, 70), -1.0), torch.full((1, 20), 5.0)] + out = client.encode(wavs) + assert out[0].dtype == torch.long and torch.equal(out[0], _codes(wavs[0])) + assert isinstance(out[1], RuntimeError) and "bad clip" in str(out[1]) + assert torch.equal(out[2], _codes(wavs[2])) + # The host process's own clips skip the socket. + assert torch.equal(host.encode([wavs[2]])[0], _codes(wavs[2])) + finally: + host.close() + + +def test_host_merges_clips_from_concurrent_clients(host_dir): + gate = threading.Event() + recorder = _Recorder(gate) + host = shared.SharedReferenceEncoderHost(host_dir, [(recorder, None)]) + try: + results = {} + + def send(name, value): + results[name] = shared.SharedReferenceEncoderClient(host_dir).encode([torch.full((1, 30), value)] * 2) + + # Encoding is held until both clients' clips are queued. + warm = threading.Thread(target=lambda: host.encode([torch.full((1, 10), 1.0)])) + warm.start() + threads = [threading.Thread(target=send, args=(n, v)) for n, v in (("a", 2.0), ("b", 4.0))] + for thread in threads: + thread.start() + threading.Timer(0.5, gate.set).start() + for thread in (warm, *threads): + thread.join(10) + # Both clients' clips were encoded together, not one request at a time. + assert sum(recorder.batches) == 5 and max(recorder.batches) >= 4 + assert all(int(c[0, 0]) == 2 for c in results["a"]) and all(int(c[0, 0]) == 4 for c in results["b"]) + finally: + host.close() + + +def test_client_without_host_raises(host_dir, monkeypatch): + monkeypatch.setattr(shared, "_CONNECT_TIMEOUT_S", 0.3) + with pytest.raises((FileNotFoundError, ConnectionRefusedError)): + shared.SharedReferenceEncoderClient(host_dir).encode([torch.zeros(1, 10)]) + + +def test_host_window_collects_across_workers_without_reserving_singletons(host_dir): + recorder = _Recorder() + host = shared.SharedReferenceEncoderHost(host_dir, [(recorder, None)] * 3, batch_window_ms=500) + try: + barrier = threading.Barrier(3) + results = [] + + def send(value): + barrier.wait() + results.extend(host.encode([torch.full((1, 20), value)])) + + threads = [threading.Thread(target=send, args=(v,)) for v in (1.0, 2.0)] + for thread in threads: + thread.start() + barrier.wait() + for thread in threads: + thread.join(5) + assert not thread.is_alive() + assert recorder.batches == [2] + assert sorted(int(x[0, 0]) for x in results) == [1, 2] + finally: + host.close() + + +def test_reference_encoder_client_falls_back_to_local_encode(host_dir, monkeypatch): + from tests.model_executor.models.moss_tts.test_reference_encoder import _N_VQ, _SR, _FakeProcessor + from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod + from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache + + monkeypatch.setenv(re_mod._SHARED_CODES_DIR_ENV, host_dir) + monkeypatch.setattr(re_mod, "shared_encoder_role", lambda: "client") + monkeypatch.setattr(shared, "_CONNECT_TIMEOUT_S", 0.3) + proc = _FakeProcessor() + enc = re_mod.MossReferenceEncoder( + proc, variant="local", n_vq=_N_VQ, sr_target=_SR, speaker_cache=SpeakerEmbeddingCache(max_bytes=1 << 20) + ) + placed = [] + monkeypatch.setattr(enc, "_place_tokenizer_for_local_encode", lambda: placed.append(True)) + out = enc._encode_batch_sync([([7.0] * 50, _SR)]) + assert placed == [True] and proc.attempt_sizes == [1] and int(out[0][0, 0]) == 7 + + +def test_workers_encode_concurrently_after_their_warmups(host_dir): + inside, release = threading.Barrier(3, timeout=5), threading.Event() + warmed = [] + + def worker(name): + def encode(wavs): + inside.wait() # both workers are inside an encode at the same time + release.wait(5) + return [_codes(w) for w in wavs] + + return encode, lambda: warmed.append(name) + + host = shared.SharedReferenceEncoderHost(host_dir, [worker("a"), worker("b")]) + try: + out = [] + + def submit(value): + out.append(host.encode([torch.full((1, 20), value)])) + + first = threading.Thread(target=submit, args=(1.0,)) + first.start() + # The first worker is busy by now, so the second clip goes to the other one. + threading.Event().wait(0.1) + second = threading.Thread(target=submit, args=(2.0,)) + second.start() + inside.wait() + release.set() + for thread in (first, second): + thread.join(5) + assert sorted(warmed) == ["a", "b"] and sorted(int(o[0][0, 0]) for o in out) == [1, 2] + finally: + host.close() + + +def test_client_requests_in_flight_use_their_own_connections(host_dir): + entered, release = threading.Semaphore(0), threading.Event() + + def worker(): + def encode(wavs): + entered.release() + release.wait(5) + return [_codes(w) for w in wavs] + + return encode, None + + host = shared.SharedReferenceEncoderHost(host_dir, [worker(), worker()]) + try: + client = shared.SharedReferenceEncoderClient(host_dir) + out = [] + threads = [ + threading.Thread(target=lambda v=v: out.append(client.encode([torch.full((1, 20), v)]))) for v in (1.0, 2.0) + ] + for thread in threads: + thread.start() + # A worker holds this request, still unanswered, before the next is + # sent. Sent together, one free worker may take both as one batch. + assert entered.acquire(timeout=5) + release.set() # both requests reached the host at once + for thread in threads: + thread.join(5) + assert sorted(int(o[0][0, 0]) for o in out) == [1, 2] + assert len(client._idle) == 2 # both connections are kept for reuse + finally: + host.close() + + +def test_readiness_waits_for_every_worker_warmup(host_dir, monkeypatch): + entered, release = threading.Event(), threading.Event() + recorder = _Recorder() + warmed = [] + + def last_warmup(): + entered.set() + assert release.wait(5) + warmed.append("last") + + host = shared.SharedReferenceEncoderHost( + host_dir, [(recorder, lambda: warmed.append("first")), (recorder, last_warmup)] + ) + client = shared.SharedReferenceEncoderClient(host_dir) + try: + assert entered.wait(5) + monkeypatch.setattr(shared, "_STARTUP_TIMEOUT_S", 0.03) + with pytest.raises(shared.SharedReferenceEncoderStartupError): + host.wait_until_ready() + with pytest.raises(shared.SharedReferenceEncoderStartupError): + client.wait_until_ready() + assert warmed == ["first"] and recorder.batches == [] and not client._idle + release.set() + monkeypatch.setattr(shared, "_STARTUP_TIMEOUT_S", 5) + host.wait_until_ready() + client.wait_until_ready() + assert warmed == ["first", "last"] and recorder.batches == [] + assert len(client._idle) == 1 + assert torch.equal(client.encode([torch.ones(1, 20)])[0], _codes(torch.ones(1, 20))) + finally: + release.set() + host.close() + + +def test_slow_startup_does_not_use_encode_timeout(host_dir, monkeypatch): + release = threading.Event() + monkeypatch.setattr(shared, "_REQUEST_TIMEOUT_S", 0.05) + monkeypatch.setattr(shared, "_STARTUP_TIMEOUT_S", 5, raising=False) + host = shared.SharedReferenceEncoderHost(host_dir, [(_Recorder(), lambda: release.wait(5))]) + timer = threading.Timer(0.2, release.set) + timer.start() + try: + # The first socket waits for readiness before starting its much shorter + # encode deadline. No API warmup call is required to protect this path. + client = shared.SharedReferenceEncoderClient(host_dir) + assert torch.equal(client.encode([torch.ones(1, 20)])[0], _codes(torch.ones(1, 20))) + finally: + release.set() + timer.join(5) + host.close() + + +def test_failed_host_warmup_is_not_ready(host_dir): + def fail(): + raise RuntimeError("warmup failed") + + recorder = _Recorder() + host = shared.SharedReferenceEncoderHost(host_dir, [(recorder, fail)]) + try: + with pytest.raises(shared.SharedReferenceEncoderStartupError): + host.wait_until_ready() + with pytest.raises(shared.SharedReferenceEncoderStartupError): + shared.SharedReferenceEncoderClient(host_dir).wait_until_ready() + with pytest.raises(shared.SharedReferenceEncoderStartupError): + host.encode([torch.ones(1, 20)]) + assert recorder.batches == [] + finally: + host.close() + + +def test_startup_failure_does_not_trigger_local_capture(host_dir, monkeypatch): + from unittest.mock import Mock + + from tests.model_executor.models.moss_tts.test_reference_encoder import _N_VQ, _SR, _FakeProcessor + from vllm_omni.model_executor.models.moss_tts import reference_encoder as re_mod + from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache + + monkeypatch.setenv(re_mod._SHARED_CODES_DIR_ENV, host_dir) + monkeypatch.setattr(re_mod, "shared_encoder_role", lambda: "client") + enc = re_mod.MossReferenceEncoder( + _FakeProcessor(), + variant="local", + n_vq=_N_VQ, + sr_target=_SR, + speaker_cache=SpeakerEmbeddingCache(max_bytes=1 << 20), + ) + enc._shared_client = Mock() + enc._shared_client.encode.side_effect = shared.SharedReferenceEncoderStartupError("still compiling") + local = Mock() + monkeypatch.setattr(enc, "_place_tokenizer_for_local_encode", local) + with pytest.raises(shared.SharedReferenceEncoderStartupError): + enc._encode_shared_or_local([torch.ones(1, 20)]) + local.assert_not_called() diff --git a/tests/model_executor/models/moss_tts/test_slot_attention.py b/tests/model_executor/models/moss_tts/test_slot_attention.py new file mode 100644 index 000000000..3ee822371 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_slot_attention.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import copy + +import pytest +import torch +import torch.nn.functional as F + +from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import ( + MossAudioTokenizerMultiheadAttention, + MossAudioTokenizerRotaryEmbedding, + RingKVCache, + StreamingExecutionContext, +) +from vllm_omni.model_executor.models.moss_tts.slot_attention import slot_ring_attention, slot_ring_attention_rows + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + + +@pytest.mark.parametrize( + "frames,capacity", [(1, 125), (2, 250), (15, 125), (30, 250), (32, 400), (60, 400), (480, 400)] +) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("execution", ["graph", "compiled"]) +def test_empty_tile_skip_matches_static_ring_across_reuse(frames, capacity, dtype, execution): + torch.manual_seed(45) + batch, heads, states, dim = 3, 2, 5, 64 + projected = torch.randn(batch, frames, 3, heads, dim, device="cuda", dtype=dtype) + q, k, v = projected.permute(2, 0, 3, 1, 4).unbind(0) + cache = torch.randn(2, states, heads, capacity, dim, device="cuda", dtype=dtype) + expected_cache = cache.clone() + end = torch.zeros(states, 2, device="cuda", dtype=torch.long)[:, 0] + expected_end = end.clone() + slots = torch.tensor([0, 4, 2], device="cuda") + lengths = torch.full((batch, 2), frames, device="cuda", dtype=torch.int32)[:, 0] + + def candidate(): + return slot_ring_attention(q, k, v, cache, end, slots, lengths, 17, skip_empty_tiles=True) + + graph = None + if execution == "compiled": + run = torch.compile(candidate, fullgraph=True) + else: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + candidate() + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + actual = candidate() + + for step, start in enumerate([0, 1, capacity - 1, capacity + 9, 0, 3 * capacity]): + end.fill_(start) + expected_end.fill_(start) + cache.copy_(expected_cache) + projected.normal_() + slots.copy_(torch.roll(torch.arange(states, device="cuda"), step)[:batch]) + lengths.copy_(torch.tensor([frames, max(1, frames - 3), 0], device="cuda", dtype=torch.int32)) + expected = slot_ring_attention(q, k, v, expected_cache, expected_end, slots, lengths, 17) + if graph is not None: + graph.replay() + else: + actual = run() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(cache, expected_cache, rtol=0, atol=0) + torch.testing.assert_close(end, expected_end, rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("frames,capacity", [(1, 125), (15, 140), (480, 400), (480, 880)]) +def test_all_padding_preserves_state_and_replay_can_become_active(dtype, frames, capacity): + q, k, v = [torch.randn(2, 2, frames, 64, device="cuda", dtype=dtype) for _ in range(3)] + cache = torch.randn(2, 4, 2, capacity, 64, device="cuda", dtype=dtype) + end = torch.tensor([0, capacity + 3, 9, 0], device="cuda") + slots = torch.tensor([1, 3], device="cuda") + lengths = torch.zeros(2, device="cuda", dtype=torch.int32) + before_cache, before_end = cache.clone(), end.clone() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + slot_ring_attention(q, k, v, cache, end, slots, lengths, capacity) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = slot_ring_attention(q, k, v, cache, end, slots, lengths, capacity) + for active in [False, True, False, True]: + lengths.fill_(frames if active else 0) + expected = _ring_reference(q, k, v, before_cache, before_end, slots, lengths, capacity) + graph.replay() + torch.testing.assert_close(cache, before_cache, rtol=0, atol=0) + torch.testing.assert_close(end, before_end, rtol=0, atol=0) + torch.testing.assert_close(out, expected, rtol=0.025, atol=0.025) + if not active: + assert out.count_nonzero() == 0 + + +def test_dense_ring_retains_last_tokens_when_chunk_exceeds_capacity(): + torch.manual_seed(19) + capacity, frames = 400, 480 + state = RingKVCache(3, 2, 64, capacity, device=torch.device("cuda"), dtype=torch.float32) + slots = torch.tensor([1, 0], device="cuda") + context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=torch.ones(2, device="cuda", dtype=torch.bool)) + for start in [0, frames]: + k = torch.randn(2, 2, frames, 64, device="cuda") + v = torch.randn_like(k) + state.complete(k, v, execution_context=context) + retained = torch.arange(frames - capacity, frames, device="cuda") + destinations = (start + retained) % capacity + torch.testing.assert_close(state.cache[0, slots][:, :, destinations], k[:, :, retained], rtol=0, atol=0) + torch.testing.assert_close(state.cache[1, slots][:, :, destinations], v[:, :, retained], rtol=0, atol=0) + torch.testing.assert_close(state.end_offset[slots], torch.full_like(slots, start + frames), rtol=0, atol=0) + + +def _ring_reference(q, k, v, cache, end, slots, lengths, context): + """Independent chronological ring oracle with FP32 attention arithmetic.""" + capacity = cache.shape[-2] + output = torch.zeros_like(q) + for row, (slot, length) in enumerate(zip(slots.tolist(), lengths.tolist())): + if length == 0: + continue + start = int(end[slot]) + for token in range(max(0, length - capacity), length): + cache[0, slot, :, (start + token) % capacity] = k[row, :, token] + cache[1, slot, :, (start + token) % capacity] = v[row, :, token] + stop = start + length + physical_positions = torch.full((capacity,), -1, device=q.device, dtype=torch.long) + available = torch.arange(max(0, stop - capacity), stop, device=q.device) + physical_positions[available % capacity] = available + queries = torch.arange(start, stop, device=q.device) + distance = queries[:, None] - physical_positions[None] + mask = (physical_positions[None] >= 0) & (distance >= 0) + if context > 0: + mask &= distance < context + output[row, :, :length] = F.scaled_dot_product_attention( + q[row, :, :length].float(), cache[0, slot].float(), cache[1, slot].float(), mask[None] + ).to(q.dtype) + end[slot] = stop + return output + + +@pytest.mark.parametrize("frames,capacity,head_dim", [(1, 125, 64), (15, 125, 64), (480, 400, 64), (33, 65, 32)]) +@pytest.mark.parametrize("execution", ["eager", "graph", "compiled"]) +def test_slot_attention_preserves_ring_state_and_padding(frames, capacity, head_dim, execution): + torch.manual_seed(29) + batch, heads, state_capacity = 3, 2, 5 + projected = torch.randn(batch, frames, 3, heads, head_dim, device="cuda", dtype=torch.bfloat16) + q, k, v = projected.permute(2, 0, 3, 1, 4).unbind(0) + cache = torch.randn(2, state_capacity, heads, capacity, head_dim, device="cuda", dtype=q.dtype) + # Strided metadata is allowed by the kernel contract. + end = torch.zeros(state_capacity, 2, device="cuda", dtype=torch.long)[:, 0] + slots = torch.arange(batch * 2, device="cuda")[::2] + lengths = torch.full((batch, 2), frames, device="cuda", dtype=torch.int32)[:, 0] + context = capacity + attention = slot_ring_attention + if execution == "compiled": + # Functionalization must preserve mutations of the persistent ring + # and its offsets between calls, including slot reuse. + attention = torch.compile(slot_ring_attention, backend="inductor", fullgraph=True) + if execution == "graph": + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + slot_ring_attention(q, k, v, cache, end, slots, lengths, context) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = slot_ring_attention(q, k, v, cache, end, slots, lengths, context) + end.zero_() + expected_cache, expected_end = cache.clone(), end.clone() + for step in range(8): + slots.copy_(torch.roll(torch.arange(state_capacity, device="cuda"), step)[:batch]) + lengths.copy_(torch.tensor([frames, max(1, frames - 3), 0], device="cuda", dtype=torch.int32)) + projected.normal_() + if step == 4: + end[1:3].zero_() + expected_end[1:3].zero_() + expected = _ring_reference(q, k, v, expected_cache, expected_end, slots, lengths, context) + if execution == "graph": + graph.replay() + else: + actual = attention(q, k, v, cache, end, slots, lengths, context) + torch.testing.assert_close(cache, expected_cache, rtol=0, atol=0) + torch.testing.assert_close(end, expected_end, rtol=0, atol=0) + torch.testing.assert_close(actual, expected, rtol=0.025, atol=0.025) + padding = torch.arange(frames, device="cuda")[None] >= lengths[:, None] + assert actual.permute(0, 2, 1, 3)[padding].count_nonzero() == 0 + + +@pytest.mark.parametrize("frames,context", [(1, 9), (3, 9), (15, 125), (480, 400)]) +def test_slot_attention_integrates_with_current_mha(frames, context): + torch.manual_seed(9) + reference = MossAudioTokenizerMultiheadAttention( + 128, + 2, + causal=True, + context=context, + rope=MossAudioTokenizerRotaryEmbedding(), + device="cuda", + dtype=torch.bfloat16, + ) + candidate = copy.deepcopy(reference) + candidate._slot_attention = slot_ring_attention + for model in (reference, candidate): + model._streaming_state = model._init_streaming_state(4) + for step in range(8): + slots = torch.tensor([1, 0, 3] if step % 2 else [0, 1, 3], device="cuda") + valid = torch.tensor([True, True, False], device="cuda") + execution_context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=valid) + x = torch.randn(3, frames, 128, device="cuda", dtype=torch.bfloat16) + with torch.inference_mode(): + expected = reference(x, x, x, execution_context=execution_context) + actual = candidate(x, x, x, execution_context=execution_context) + torch.testing.assert_close(actual, expected, atol=0.01, rtol=0.02) + cs, rs = candidate._streaming_state, reference._streaming_state + torch.testing.assert_close(cs.offset, rs.offset, rtol=0, atol=0) + torch.testing.assert_close(cs.kv_cache.end_offset, rs.kv_cache.end_offset, rtol=0, atol=0) + torch.testing.assert_close(cs.kv_cache.cache[:, :3], rs.kv_cache.cache[:, :3], rtol=0, atol=0) + if step == 4: + for model in (reference, candidate): + model._streaming_state.reset_slots(torch.tensor([0], device="cuda")) + + +@pytest.mark.parametrize( + "frames,capacity,head_dim", [(1, 125, 64), (15, 140, 64), (40, 125, 64), (480, 400, 64), (33, 65, 32)] +) +@pytest.mark.parametrize("execution", ["eager", "graph", "compiled"]) +def test_fused_rows_attention_matches_three_launch_op_exactly(frames, capacity, head_dim, execution): + torch.manual_seed(31) + batch, heads, state_capacity = 3, 2, 6 + projected = torch.randn(batch, frames, 3, heads, head_dim, device="cuda", dtype=torch.bfloat16) + q, k, v = projected.permute(2, 0, 3, 1, 4).unbind(0) + cache = torch.randn(2, state_capacity, heads, capacity, head_dim, device="cuda", dtype=q.dtype) + reference_cache = cache.clone() + end = torch.zeros(state_capacity, device="cuda", dtype=torch.long) + slots = torch.zeros(batch, device="cuda", dtype=torch.long) + lengths = torch.zeros(batch, device="cuda", dtype=torch.int32) + starts = torch.zeros(batch, device="cuda", dtype=torch.long) + attention = slot_ring_attention_rows + if execution == "compiled": + attention = torch.compile(slot_ring_attention_rows, backend="inductor", fullgraph=True) + if execution == "graph": + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + slot_ring_attention_rows(q, k, v, cache, starts, slots, lengths, capacity) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = slot_ring_attention_rows(q, k, v, cache, starts, slots, lengths, capacity) + cache.copy_(reference_cache) + for step in range(8): + slots.copy_(torch.roll(torch.arange(state_capacity, device="cuda"), step)[:batch]) + lengths.copy_(torch.tensor([frames, max(1, frames - 3), 0], device="cuda", dtype=torch.int32)) + starts.copy_(end[slots]) + projected.normal_() + expected = slot_ring_attention(q, k, v, reference_cache, end, slots, lengths, capacity) + if execution == "graph": + graph.replay() + else: + actual = attention(q, k, v, cache, starts, slots, lengths, capacity) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + torch.testing.assert_close(cache, reference_cache, rtol=0, atol=0) + + +@pytest.mark.parametrize("frames,context", [(1, 9), (15, 125), (480, 400)]) +def test_fused_rows_attention_integrates_with_mha_offsets(frames, context): + torch.manual_seed(11) + reference = MossAudioTokenizerMultiheadAttention( + 128, + 2, + causal=True, + context=context, + rope=MossAudioTokenizerRotaryEmbedding(), + device="cuda", + dtype=torch.bfloat16, + ) + candidate = copy.deepcopy(reference) + reference._slot_attention = slot_ring_attention + candidate._slot_attention_rows = slot_ring_attention_rows + for model in (reference, candidate): + model._streaming_state = model._init_streaming_state(4) + for step in range(8): + slots = torch.tensor([1, 0, 3] if step % 2 else [0, 1, 3], device="cuda") + valid = torch.tensor([True, True, False], device="cuda") + execution_context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=valid) + x = torch.randn(3, frames, 128, device="cuda", dtype=torch.bfloat16) + with torch.inference_mode(): + expected = reference(x, x, x, execution_context=execution_context) + actual = candidate(x, x, x, execution_context=execution_context) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + cs, rs = candidate._streaming_state, reference._streaming_state + torch.testing.assert_close(cs.offset, rs.offset, rtol=0, atol=0) + torch.testing.assert_close(cs.kv_cache.end_offset, rs.kv_cache.end_offset, rtol=0, atol=0) + torch.testing.assert_close(cs.kv_cache.cache, rs.kv_cache.cache, rtol=0, atol=0) + if step == 4: + for model in (reference, candidate): + model._streaming_state.reset_slots(torch.tensor([0], device="cuda")) diff --git a/tests/model_executor/models/moss_tts/test_streaming_decoder_wrapper_concurrency.py b/tests/model_executor/models/moss_tts/test_streaming_decoder_wrapper_concurrency.py new file mode 100644 index 000000000..79c4351a4 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_streaming_decoder_wrapper_concurrency.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""A second streaming-decoder graph set can replay concurrently with the first.""" + +import threading + +import pytest +import torch +from torch import nn + +from vllm_omni.model_executor.models.moss_tts.cuda_graph_streaming_decoder_wrapper import ( + CUDAGraphStreamingDecoderWrapper, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + +CAP, DIM = 8, 64 + + +class _Codec(nn.Module): + def __init__(self, device): + super().__init__() + self.weight = nn.Parameter(torch.randn(DIM, DIM, device=device) / DIM) + self.state = torch.zeros(CAP + 16, DIM, device=device) + + def decode_streaming_tensors(self, codes, lengths, slots, valid_rows): + x = codes.float().sum(0)[..., None].expand(-1, -1, DIM) + self.state[slots][:, None] + y = x @ self.weight + self.state.index_copy_(0, slots, y[:, -1]) + return y.transpose(1, 2).contiguous(), lengths + + def reset_decoder_state_slots(self, slots): + self.state.index_fill_(0, slots.to(self.state.device), 0) + + +def _wrapper(codec, **kwargs): + return CUDAGraphStreamingDecoderWrapper( + codec, + state_capacity=CAP, + num_quantizers=2, + vllm_config=None, + compiled_decode=codec.decode_streaming_tensors, + **kwargs, + ) + + +def test_private_wrapper_captures_on_its_own_stream_and_replays_concurrently(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + device = torch.device("cuda", torch.accelerator.current_device_index()) + codec = _Codec(device) + main = _wrapper(codec, batch_sizes=[4], frame_sizes=[3]) + fast = _wrapper(codec, batch_sizes=[1, 2], frame_sizes=[1], scratch_base=CAP + 4, private_pool=True) + main.warmup(device) + fast.warmup(device) + assert main.is_ready and fast.is_ready + assert main._capture_stream is None and fast._capture_stream is not None + + codes_main = torch.randint(0, 5, (2, 4, 3), device=device) + codes_fast = torch.randint(0, 5, (2, 2, 1), device=device) + errors = [] + + def run(wrapper, codes, slots, iters): + stream = torch.cuda.Stream(device=device) + slot_ids = torch.tensor(slots, device=device) + try: + with torch.cuda.stream(stream): + for _ in range(iters): + codec.reset_decoder_state_slots(slot_ids) + audio = wrapper.decode(codes, slot_ids)[0][: len(slots)].clone() + expected = (codes.float().sum(0)[..., None].expand(-1, -1, DIM) @ codec.weight).transpose(1, 2) + stream.synchronize() + torch.testing.assert_close(audio, expected, rtol=1e-3, atol=1e-3) + except Exception as error: # surfaced below + errors.append(error) + + threads = [ + threading.Thread(target=run, args=(main, codes_main, [0, 1, 2, 3], 200)), + threading.Thread(target=run, args=(fast, codes_fast, [4, 5], 200)), + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=60) + assert not any(thread.is_alive() for thread in threads), "concurrent graph replay did not finish" + assert not errors, errors diff --git a/tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py b/tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py index fff97cd96..5ccd758f3 100644 --- a/tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py +++ b/tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py @@ -54,8 +54,19 @@ def test_terminal_padding_preserves_audio_live_state_and_reused_slots(): assert [resumed[i].item() for i in (0, 1)] == [2, 36] -@pytest.mark.parametrize("graph_capacity, expected_calls", [(None, 5), (4, 5), (8, 3)]) -def test_only_existing_compatible_graph_buckets_are_coalesced(mocker, graph_capacity, expected_calls): +@pytest.mark.parametrize( + "graph_capacity,expected_groups", + [ + # No graph, or graphs that do not cover the state capacity: exact tails. + (None, {3: [0], 7: [1], 8: [2], 1: [3], 15: [4]}), + (4, {3: [0], 7: [1], 8: [2], 1: [3], 15: [4]}), + # Full graph coverage: finished tails (3, 7) join the padded 8-frame + # bucket with the live 8-frame row; the finished 1-frame row keeps the + # exact 1-frame graph and the live 15-frame row stays alone. + (8, {8: [0, 1, 2], 1: [3], 15: [4]}), + ], +) +def test_terminal_tails_coalesce_into_padded_frame_buckets(mocker, graph_capacity, expected_groups): s = session() s._cudagraph_wrapper = ( None @@ -81,9 +92,16 @@ def test_only_existing_compatible_graph_buckets_are_coalesced(mocker, graph_capa for i, (length, done) in enumerate(zip(lengths, finished, strict=True)) ] ) - assert step.call_count == expected_calls + # Slots are leased in request order, so slot == row here. + groups = {} + for call in step.call_args_list: + plan = call.args[0] + groups[max(int(codes.shape[1]) for codes in plan.values())] = sorted(plan) + terminal = call.kwargs["terminal_slots"] or set() + assert terminal == {slot for slot in plan if finished[slot]} + assert groups == expected_groups + assert step.call_count == len(expected_groups) + # Output is always cropped to the exact frame count, padded or not. assert [result[i].shape[-1] for i in range(5)] == lengths - if graph_capacity == 8: - assert list(step.call_args_list[0].args[0]) == [0, 1, 2] - assert step.call_args_list[0].kwargs["terminal_slots"] == {0, 1} - assert list(step.call_args_list[1].args[0]) == [3] # Preserve the exact one-frame graph. + # Finished rows released their slots; live rows keep theirs. + assert set(decoder._stream_req_slots) == {"2", "4"} diff --git a/tests/model_executor/stage_input_processors/test_moss_raw_async_chunk.py b/tests/model_executor/stage_input_processors/test_moss_raw_async_chunk.py new file mode 100644 index 000000000..9d9f6b7c1 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_moss_raw_async_chunk.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""MOSS raw codec rows are buffered into first/regular chunks in codebook-major order.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.stage_input_processors.moss_tts import ( + _MOSS_AUDIO_PAD_CODE, + talker2codec_raw_async_chunk, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _manager(first=1, regular=3): + config = {"extra": {"codec_chunk_frames": regular, "initial_codec_chunk_frames": first}} + return SimpleNamespace(connector=SimpleNamespace(config=config)) + + +def _step(manager, frames, finished=False): + request = SimpleNamespace(request_id="r", external_req_id="r") + payload = talker2codec_raw_async_chunk(manager, {"codes": {"audio": frames}}, request, is_finished=finished) + if payload is not None and payload.codes is not None: + manager.put_req_chunk["r"] += 1 + manager.ramp_chunk_count["r"] += 1 + return payload + + +def test_first_regular_and_terminal_chunks(): + manager = _manager() + rows = torch.arange(6 * 4, dtype=torch.int32).reshape(6, 4) + first = _step(manager, rows[:1]) + assert first.meta.codec_chunk_frames == 1 + assert first.codes.audio.dtype == torch.int64 + assert torch.equal(first.codes.audio, rows[0].long()) + + assert _step(manager, rows[1:2]) is None + padded = torch.full((1, 4), _MOSS_AUDIO_PAD_CODE, dtype=torch.int32) + assert _step(manager, padded) is None # PAD rows are not audio + assert _step(manager, rows[2:3]) is None + regular = _step(manager, rows[3:4]) + assert regular.meta.codec_chunk_frames == 3 + assert torch.equal(regular.codes.audio, rows[1:4].long().T.reshape(-1)) + + tail = _step(manager, rows[4:6], finished=True) + assert tail.meta.codec_chunk_frames == 2 + assert torch.equal(tail.codes.audio, rows[4:6].long().T.reshape(-1)) + assert bool(tail.meta.finished) + assert "r" not in manager.code_prompt_token_ids + + +def test_buffered_rows_do_not_alias_producer_storage(): + manager = _manager(first=2, regular=2) + frames = torch.tensor([[7, 8, 9]]) + assert _step(manager, frames) is None + frames.fill_(0) + chunk = _step(manager, torch.tensor([[1, 2, 3]])) + assert torch.equal(chunk.codes.audio, torch.tensor([7, 1, 8, 2, 9, 3])) + + +@pytest.mark.parametrize("total", [1, 3, 5, 10, 13, 28, 31]) +def test_ramp_preserves_every_frame_and_terminal(total): + manager = _manager(first=1, regular=15) + manager.connector.config["extra"]["codec_chunk_ramp"] = [1, 4, 8, 15] + rows = torch.arange(total * 4, dtype=torch.int32).reshape(total, 4) + packets = [] + for index in range(total): + packet = _step(manager, rows[index : index + 1], finished=index == total - 1) + if packet is not None: + packets.append(packet) + rebuilt = torch.cat([packet.codes.audio.reshape(4, -1).T for packet in packets]) + assert torch.equal(rebuilt, rows.long()) + assert all(not bool(packet.meta.finished) for packet in packets[:-1]) + assert bool(packets[-1].meta.finished) + assert "r" not in manager.code_prompt_token_ids + if total == 31: + assert [packet.meta.codec_chunk_frames for packet in packets] == [1, 4, 8, 15, 3] diff --git a/tests/model_executor/stage_input_processors/test_moss_raw_empty_terminal.py b/tests/model_executor/stage_input_processors/test_moss_raw_empty_terminal.py new file mode 100644 index 000000000..a2277fa62 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_moss_raw_empty_terminal.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Empty MOSS terminals must not execute prewarmed or previous codec tokens.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.core.sched.omni_scheduling_coordinator import OmniSchedulingCoordinator +from vllm_omni.data_entry_keys import to_dict +from vllm_omni.distributed.omni_connectors.model_runner.omni_connector_payload_transport import ( + _OmniConnectorPayloadTransportMixin as Transport, +) +from vllm_omni.model_executor.stage_input_processors.moss_tts import ( + _MOSS_AUDIO_PAD_CODE, + talker2codec_raw_async_chunk, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _request(tokens): + return SimpleNamespace( + request_id="r", + external_req_id="r", + prompt_token_ids=list(tokens), + num_prompt_tokens=len(tokens), + _all_token_ids=list(tokens), + _output_token_ids=[], + num_computed_tokens=0, + ) + + +@pytest.mark.parametrize("after_audio", [False, True]) +def test_empty_terminal_clears_native_codec_prompt(after_audio): + manager = SimpleNamespace( + connector=SimpleNamespace(config={"extra": {"initial_codec_chunk_frames": 1, "codec_chunk_frames": 3}}) + ) + source = _request([0] * 88) + previous = [0] * 88 + if after_audio: + packet = talker2codec_raw_async_chunk(manager, {"codes": {"audio": torch.ones(1, 12)}}, source) + manager.put_req_chunk["r"] += 1 + previous = packet.codes.audio.tolist() + terminal = talker2codec_raw_async_chunk( + manager, + {"codes": {"audio": torch.full((1, 12), _MOSS_AUDIO_PAD_CODE)}}, + source, + is_finished=True, + ) + payload = to_dict(terminal) + assert not Transport._payload_is_consumable(payload) + metadata = Transport._extract_scheduling_metadata(payload) + receiver = _request(previous) + coordinator = OmniSchedulingCoordinator(stage_id=1) + coordinator.update_request_metadata({"r": receiver}, {"r": metadata}, model_mode="generation") + assert receiver.prompt_token_ids == [] + assert receiver._all_token_ids == [] + assert receiver.num_prompt_tokens == receiver.num_computed_tokens == 0 + assert "r" in coordinator.input_terminal_req_ids + assert "r" not in manager.code_prompt_token_ids + + +@pytest.mark.parametrize("empty", [[], torch.empty(0, dtype=torch.long)]) +def test_explicit_empty_codes_clear_generation_prompt(empty): + request = _request([17, 18, 19]) + request._output_token_ids = [23] + request.num_computed_tokens = 3 + coordinator = OmniSchedulingCoordinator(stage_id=1) + coordinator.update_request_metadata( + {"r": request}, {"r": {"code_predictor_codes": empty, "input_terminal": True}}, model_mode="generation" + ) + assert request.prompt_token_ids == request._all_token_ids == request._output_token_ids == [] + assert request.num_prompt_tokens == request.num_computed_tokens == 0 + assert "r" in coordinator.input_terminal_req_ids + + +@pytest.mark.parametrize("metadata", [{"input_terminal": True}, {"code_predictor_codes": None}]) +def test_absent_codes_preserve_generation_prompt(metadata): + request = _request([17, 18, 19]) + coordinator = OmniSchedulingCoordinator(stage_id=1) + coordinator.update_request_metadata({"r": request}, {"r": metadata}, model_mode="generation") + assert request.prompt_token_ids == request._all_token_ids == [17, 18, 19] + + +def test_ar_prompt_is_not_replaced_by_empty_audio_codes(): + request = _request([17, 18, 19]) + coordinator = OmniSchedulingCoordinator(stage_id=1) + coordinator.update_request_metadata( + {"r": request}, {"r": {"code_predictor_codes": [], "input_terminal": True}}, model_mode="ar" + ) + assert request.prompt_token_ids == request._all_token_ids == [17, 18, 19] diff --git a/tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py new file mode 100644 index 000000000..4625c6399 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_moss_tts_async_chunk.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Progressive Local codec chunks preserve frames and agree with graph shapes.""" + +from collections import defaultdict +from dataclasses import dataclass, field +from unittest.mock import Mock + +import pytest +import torch +from transformers import PretrainedConfig +from vllm.config import CompilationConfig, ModelConfig, SchedulerConfig, VllmConfig + +from vllm_omni.data_entry_keys import OmniPayloadStruct +from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import MossTTSCodecDecoder +from vllm_omni.model_executor.stage_input_processors.moss_tts import talker2codec_raw_async_chunk + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@dataclass +class _Connector: + config: dict[str, dict[str, object]] + + +@dataclass +class _TransferManager: + connector: _Connector + code_prompt_token_ids: dict[str, list[object]] = field(default_factory=lambda: defaultdict(list)) + request_payload: dict[str, object] = field(default_factory=dict) + put_req_chunk: dict[str, int] = field(default_factory=lambda: defaultdict(int)) + ramp_chunk_count: dict[str, int] = field(default_factory=lambda: defaultdict(int)) + + +@dataclass +class _Request: + external_req_id: str + + +def _manager(ramp: object = None, *, first: int = 1) -> _TransferManager: + extra = {"initial_codec_chunk_frames": first, "codec_chunk_frames": 15, "codec_chunk_ramp": ramp} + return _TransferManager(_Connector({"extra": extra})) + + +def _emit( + manager: _TransferManager, req_id: str, value: int | None, finished: bool = False +) -> OmniPayloadStruct | None: + output = None if value is None else {"codes": {"audio": torch.full((1, 4), value)}} + payload = talker2codec_raw_async_chunk(manager, output, _Request(req_id), finished) + if payload is not None: + manager.put_req_chunk[req_id] += 1 + manager.ramp_chunk_count[req_id] += 1 + return payload + + +@pytest.mark.parametrize( + "ramp,first,expected", + [ + (None, 1, [1, 15, 15, 9]), + ([1, 2, 4, 8, 15], 1, [1, 2, 4, 8, 15, 10]), + ([1, 4, 15], 7, [1, 4, 15, 15, 5]), + ([4, 8, 15], 1, [4, 8, 15, 13]), + ([1, 20, 15], 1, [1, 20, 15, 4]), + ("1,4,15", 1, [1, 4, 15, 15, 5]), + ([1, 0, 15], 1, [1, 15, 15, 9]), + ([1], 1, [1, 15, 15, 9]), + ("invalid", 1, [1, 15, 15, 9]), + ], +) +def test_order_and_final_flush(ramp: object, first: int, expected: list[int]) -> None: + manager = _manager(ramp, first=first) + packets = [] + for index in range(40): + packet = _emit(manager, "a", index, finished=index == 39) + if packet is not None: + packets.append(packet) + assert [p.meta.codec_chunk_frames for p in packets] == expected + actual = torch.cat([p.codes.audio.reshape(4, -1).T for p in packets]) + assert torch.equal(actual, torch.arange(40).unsqueeze(1).expand(-1, 4)) + assert all(not bool(p.meta.finished) for p in packets[:-1]) + assert bool(packets[-1].meta.finished) + assert "a" not in manager.code_prompt_token_ids + + +def test_requests_progress_independently_and_empty_finish() -> None: + manager = _manager([1, 2, 4, 8, 15]) + for name in ("a", "b"): + first = _emit(manager, name, 0) + assert first is not None + assert first.meta.codec_chunk_frames == 1 + assert _emit(manager, "a", 1) is None + final_a = _emit(manager, "a", None, finished=True) + assert final_a is not None + assert final_a.meta.codec_chunk_frames == 1 + assert bool(final_a.meta.finished) + final_b = _emit(manager, "b", None, finished=True) + assert final_b is not None + assert final_b.meta.codec_chunk_frames == 0 + assert bool(final_b.meta.finished) + assert final_b.codes.audio.numel() == 0 + assert not manager.code_prompt_token_ids + + +@pytest.mark.parametrize( + "ramp,first,expected,max_step,effective_first", + [ + (None, 1, [1, 15], 15, 1), + ([1, 2, 4, 8, 15], 1, [1, 2, 4, 8, 15], 15, 1), + ([1, 20, 15], 1, [1, 15, 20], 20, 1), + ([4, 8, 15], 1, [4, 8, 15], 15, 4), + ([1, 4, 15], 7, [1, 4, 15], 15, 1), + ([1, 0, 15], 1, [1, 15], 15, 1), + ], +) +def test_graph_shapes_and_fast_path_follow_sender( + ramp: object, first: int, expected: list[int], max_step: int, effective_first: int +) -> None: + connector = _manager(ramp, first=first).connector.config + connector["extra"]["codec_first_chunk_fast_path"] = 1 + config = Mock( + spec=VllmConfig, + model_config=Mock( + spec=ModelConfig, + hf_config=PretrainedConfig(model_type="moss_tts_local"), + async_chunk=True, + use_v2_model_runner=True, + enforce_eager=False, + stage_connector_config=connector, + ), + scheduler_config=Mock(spec=SchedulerConfig, max_num_seqs=8), + compilation_config=Mock(spec=CompilationConfig, cudagraph_capture_sizes=[1, 2, 4, 8]), + ) + codec = MossTTSCodecDecoder(vllm_config=config) + assert codec._streaming_graph_frame_sizes == expected + assert codec._stream_max_step_frames == max_step + assert codec._first_chunk_fast + assert codec._initial_stream_chunk_frames == effective_first diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index 9006d58c8..52b7d61dd 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.py @@ -39,6 +39,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] class MockConnector: + sender_host: str | None = None + sender_zmq_port: int | None = None + """In-memory connector for testing (mimics OmniConnectorBase).""" def __init__(self, stage_id: int = 0): @@ -187,7 +190,7 @@ def test_init_payload_connector_ownership(role, custom_func, expected): @pytest.mark.parametrize("async_chunk", [False, True]) @pytest.mark.parametrize("rank", [0, 1]) -def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank): +def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank, mocker): host = MixinHost() host.init_omni_connectors(_make_model_config(async_chunk=async_chunk)) # No background threads: deterministic interleaved polling. @@ -196,7 +199,7 @@ def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank): connector = host._omni_connector connector.sender_host = "configured-host" connector.sender_zmq_port = 50051 - connector.get = MagicMock(return_value=None) + get = mocker.patch.object(connector, "get", return_value=None) first = _make_request("r1", "external1") first.payload_sender_info = {"host": "producer-a", "zmq_port": "50101"} second = _make_request("r2", "external2") @@ -209,14 +212,14 @@ def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank): for req_id in ("r2", "r1", "r2"): assert not host._poll_single_request(req_id) if rank == 1: - connector.get.assert_not_called() + get.assert_not_called() else: - assert connector.get.call_args_list == [ + assert get.call_args_list == [ unittest.mock.call("0", "1", "external2_0_0", {"source_host": "producer-b", "source_port": 51101}), unittest.mock.call("0", "1", "external1_0_0", {"source_host": "producer-a", "source_port": 50101}), unittest.mock.call("0", "1", "external2_0_0", {"source_host": "producer-b", "source_port": 51101}), ] - connector.get.return_value = ({"ids": {"output": [1]}, "meta": {"finished": True}}, 1) + get.return_value = ({"ids": {"output": [1]}, "meta": {"finished": True}}, 1) assert host._poll_single_request("r2") assert "r2" not in host._pending_load_reqs assert "r1" in host._pending_load_reqs @@ -225,14 +228,14 @@ def test_split_request_endpoint_retry_completion_and_cleanup(async_chunk, rank): assert not host._pending_load_reqs # Reused internal ID without sender info cannot inherit either endpoint. replacement = _make_request("r1", "external-new") - connector.get.reset_mock(return_value=True) - connector.get.return_value = None + get.reset_mock(return_value=True) + get.return_value = None host.register_chunk_recv(replacement) assert not host._poll_single_request("r1") if rank == 0: - connector.get.assert_called_once_with("0", "1", "external-new_0_0") + get.assert_called_once_with("0", "1", "external-new_0_0") else: - connector.get.assert_not_called() + get.assert_not_called() assert connector.sender_host == "configured-host" assert connector.sender_zmq_port == 50051 assert first.payload_sender_info == {"host": "producer-a", "zmq_port": "50101"} @@ -1437,6 +1440,58 @@ class TestAsyncPayloadLifecycle(unittest.TestCase): host.shutdown_omni_connectors() + def test_first_chunk_hook_sees_only_chunk_zero_before_staging(self): + host = MixinHost() + host.init_omni_connectors( + model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="gen"), + ) + host._omni_connector = MagicMock() + host._stage_id = 1 + host._async_chunk = True + host._model_mode = "gen" + seen = [] + + def hook(req_id, handle, payload): + # The chunk is not yet visible to the scheduler when the hook runs. + seen.append((req_id, handle.external_req_id, payload["codes"]["audio"], "r1" in host._finished_load_reqs)) + + host.set_first_chunk_hook(hook) + host.register_chunk_recv(_make_request("r1", "ext-r1")) + chunk = {"codes": {"audio": [1, 2]}, "meta": {"finished": torch.tensor(False)}} + host._omni_connector.get.return_value = (chunk, 1) + self.assertTrue(host._poll_single_request("r1")) + self.assertEqual(seen, [("r1", "ext-r1", [1, 2], False)]) + self.assertIn("r1", host._finished_load_reqs) + + host.get_omni_connector_output() + host._local_stage_payload_cache.clear() + host._local_request_metadata.clear() + host._omni_connector.get.return_value = ({"codes": {"audio": [3, 4]}, "meta": {}}, 1) + self.assertTrue(host._poll_single_request("r1")) + self.assertEqual(len(seen), 1) + + host.shutdown_omni_connectors() + + def test_first_chunk_hook_skips_terminal_first_chunk(self): + host = MixinHost() + host.init_omni_connectors( + model_config=_make_model_config(stage_id=1, async_chunk=True, worker_type="gen"), + ) + host._omni_connector = MagicMock() + host._stage_id = 1 + host._async_chunk = True + host._model_mode = "gen" + hook = MagicMock() + host.set_first_chunk_hook(hook) + host.register_chunk_recv(_make_request("r1")) + host._omni_connector.get.return_value = ( + {"codes": {"audio": [1, 2]}, "meta": {"finished": torch.tensor(True)}}, + 1, + ) + self.assertTrue(host._poll_single_request("r1")) + hook.assert_not_called() + host.shutdown_omni_connectors() + def test_non_ar_recv_waits_for_scheduler_handoff_before_fetching_next_chunk(self): host = MixinHost() host.init_omni_connectors( @@ -1480,6 +1535,39 @@ class TestAsyncPayloadLifecycle(unittest.TestCase): host.shutdown_omni_connectors() +@pytest.mark.parametrize("cancel_at", ["connector_get", "first_chunk_hook"]) +def test_cancelled_chunk_receive_does_not_recreate_delivery_state(mocker, cancel_at): + host = MixinHost() + host.init_omni_connectors(_make_model_config(async_chunk=True, worker_type="gen")) + host._stage_id = 1 + host._omni_connector = mocker.MagicMock() + host.register_chunk_recv(_make_request("r1")) + seen = [] + payload = {"codes": {"audio": [1, 2]}, "meta": {"finished": torch.tensor(False)}} + + def get(*args, **kwargs): + if cancel_at == "connector_get": + host.cleanup_finished_request("r1") + return payload, 1 + + def hook(req_id, handle, chunk): + seen.append(req_id) + host.cleanup_finished_request(req_id) + + host._omni_connector.get.side_effect = get + host.set_first_chunk_hook(hook) + try: + assert not host._poll_single_request("r1") + assert seen == (["r1"] if cancel_at == "first_chunk_hook" else []) + assert host.get_local_stage_payload("r1") is None + assert "r1" not in host._get_req_chunk + assert "r1" not in host._pending_load_reqs + output = host.get_omni_connector_output() + assert output.chunk_ready_req_ids == set() and output.request_metadata == {} + finally: + host.shutdown_omni_connectors() + + class TestRankAwareKVRouting(unittest.TestCase): def _make_host(self, *, from_tp: int, to_tp: int, local_rank: int) -> MixinHost: host = MixinHost() diff --git a/tests/worker_v2/test_first_audio_sender.py b/tests/worker_v2/test_first_audio_sender.py index 181f01d8b..8734dbf1e 100644 --- a/tests/worker_v2/test_first_audio_sender.py +++ b/tests/worker_v2/test_first_audio_sender.py @@ -235,3 +235,19 @@ def test_sender_copy_failure_reports_all_prepared_routes(): assert client == 3 assert [(o.request_id, o.finish_reason) for o in outputs.outputs] == [("r", FinishReason.ERROR)] assert output_queue.empty() + + +def test_final_codec_delivery_omits_upstream_marker_and_freezes_route(): + from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY + + outputs: queue.Queue = queue.Queue() + scheduler = _scheduler(r=3) + delivery = engine_output_queue_sink(outputs, scheduler, upstream_first_audio=False).prepare(["r"]) + scheduler.requests.clear() + pcm = torch.ones(4) + delivery(["r"], [pcm], torch.tensor(24000)) + client, batch = outputs.get_nowait() + assert client == 3 + assert FIRST_AUDIO_KEY not in batch.outputs[0].multimodal_output + assert batch.outputs[0].multimodal_output["model_outputs"] is pcm + assert batch.outputs[0].finish_reason is None diff --git a/tests/worker_v2/test_omni_ar_model_runner.py b/tests/worker_v2/test_omni_ar_model_runner.py index 0bbebf491..cd0c09877 100644 --- a/tests/worker_v2/test_omni_ar_model_runner.py +++ b/tests/worker_v2/test_omni_ar_model_runner.py @@ -15,6 +15,7 @@ from vllm.v1.outputs import RoutedExpertsTensors from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors import vllm_omni.worker_v2.omni_ar_model_runner as omni_ar_model_runner +from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner, OmniAsyncOutput from vllm_omni.worker_v2.output_snapshot import pack_output_snapshot @@ -144,6 +145,7 @@ def test_async_mm_snapshot_owns_output_until_copy_finishes() -> None: runner.model = SimpleNamespace() runner._async_mm_snapshot_slots, runner._async_mm_snapshot_events = [{}], [None] runner._async_mm_snapshot_pending, runner._async_mm_snapshot_cursor = [False], 0 + runner._last_multimodal_snapshot_slot = None waited: list[object] = [] runner.main_stream = SimpleNamespace(wait_event=waited.append) source = torch.tensor([[7, 8]], dtype=torch.long) @@ -161,6 +163,35 @@ def test_async_mm_snapshot_owns_output_until_copy_finishes() -> None: assert waited == [copy_event] +def test_producer_snapshot_is_not_repacked() -> None: + runner = OmniARModelRunner.__new__(OmniARModelRunner) + runner.model_config = SimpleNamespace(async_chunk=True) + runner.model = SimpleNamespace() + runner._async_mm_snapshot_slots, runner._async_mm_snapshot_events = [{}], [None] + runner._async_mm_snapshot_pending, runner._async_mm_snapshot_cursor = [False], 0 + runner._last_multimodal_snapshot_slot = None + source = torch.tensor([[7, 8]], dtype=torch.float32) + snapshot = pack_output_snapshot({"audio": source}, {}, max_buckets=1) + assert isinstance(snapshot, PackedOutputSnapshot) + + retained = runner._retain_multimodal_outputs(snapshot) + + assert retained is snapshot + assert runner._last_multimodal_snapshot_slot is None + assert runner._async_mm_snapshot_pending == [False] + + +def test_fresh_per_step_outputs_are_not_repacked() -> None: + runner = OmniARModelRunner.__new__(OmniARModelRunner) + runner.model_config = SimpleNamespace(async_chunk=True) + runner.model = SimpleNamespace(mm_outputs_fresh_per_step=True) + runner._last_multimodal_snapshot_slot = None + outputs = {"model_outputs": torch.ones(2, 4)} + + assert runner._retain_multimodal_outputs(outputs) is outputs + assert runner._last_multimodal_snapshot_slot is None + + def test_snapshot_slots_bounded_by_shape_and_packed_grouping_isolation(monkeypatch) -> None: slot: dict[tuple[Any, ...], torch.Tensor] = {} omni_ar_model_runner._copy_mm_to_snapshot_slot(torch.ones(1, 2), slot) diff --git a/tests/worker_v2/test_omni_model_state.py b/tests/worker_v2/test_omni_model_state.py index 1cbef67e8..86764b691 100644 --- a/tests/worker_v2/test_omni_model_state.py +++ b/tests/worker_v2/test_omni_model_state.py @@ -224,6 +224,61 @@ def test_static_decode_embeddings_refresh_from_input_ids(): assert OmniModelState._preprocess_result_needs_writeback(original, original.view_as(original)) is True +def test_moss_local_decode_runs_depth_predictor_and_routes_eos(mocker): + """Exercise MRV2 dispatch through the real Local hook, output and logits.""" + from vllm.sampling_params import SamplingParams + from vllm.v1.worker.gpu.states import RequestState + + from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import MossTTSLocalDepthTransformer + from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import ( + MossTTSLocalTalkerForGeneration, + ) + + model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration) + torch.nn.Module.__init__(model) + model.model = torch.nn.Module() + model.model.embed_tokens = torch.nn.Embedding(8, 4, _weight=torch.zeros(8, 4)) + model.hidden_size = 4 + model.n_vq = 2 + model.audio_pad_token_id = 8 + model.audio_assistant_slot_token_id = 2 + model.im_end_token_id = 3 + model.text_vocab_size = 8 + model.talker_mtp_output_key = ("audio_codes", "current") + model.talker_mtp_graph_safe = False + model.gpu_resident_buffer_keys = set() + model.audio_lm_heads = model.audio_embeddings = model.local_text_lm_head = None + model._audio_embed = lambda codes: codes[:, :1].expand(-1, 4).float() + frame = mocker.Mock(return_value=(torch.tensor([True, False]), torch.tensor([[1, 2], [3, 4]]))) + model.local_transformer = mocker.Mock(spec=MossTTSLocalDepthTransformer, generate_frame=frame) + + state = _make_state(max_num_reqs=2, has_preprocess=True, have_multimodal_outputs=True) + state.model = model + for idx in range(2): + state.intermediate_buffer.buffers[idx] = { + "req_id": f"r{idx}", + "audio_state": {"is_stopping": False}, + "hidden_states": {"last": torch.full((4,), float(idx + 1))}, + } + batch = _DummyInputBatch([1, 0], num_computed_tokens_cpu=[1, 1]) + inputs = {"input_ids": torch.tensor([2, 2]), "inputs_embeds": torch.zeros(2, 4)} + request_state = mocker.Mock(spec=RequestState, prompt_len=np.array([1, 1]), num_computed_tokens=None) + state.run_preprocess(batch, inputs, request_state) + + frame.assert_called_once() + torch.testing.assert_close(frame.call_args.args[0], torch.tensor([[2.0] * 4, [1.0] * 4])) + assert frame.call_args.kwargs["temperature"] == 1.7 + assert frame.call_args.kwargs["top_k"] == 25 + assert frame.call_args.kwargs["top_p"] == 0.8 + torch.testing.assert_close(inputs["inputs_embeds"], torch.tensor([[1.0] * 4, [0.0] * 4])) + _, payload = state.postprocess_model_output(torch.zeros(2, 4), batch, request_state) + assert [codes.tolist() for codes in payload["codes"]["audio"]] == [[[1, 2]], [[8, 8]]] + assert model.compute_logits(torch.zeros(2, 4)).argmax(-1).tolist() == [2, 3] + # An explicit local seed must reach the request-owned MRV2 generator. + params = SamplingParams(extra_args={"tts_local_seed": 17}, seed=99) + assert state._get_mtp_generator("seeded", params, torch.device("cpu")).initial_seed() == 17 + + @pytest.mark.parametrize("owned", [False, True]) def test_batched_postprocess_gpu_snapshot_writeback(owned): # batch=[1, 0]: last-token indices follow the reordered batch, the scalar diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index 06dcb3ef4..965e8d340 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -182,6 +182,24 @@ _MODEL_PROMOTE = ( "VLLM_GEPARD_FIRST_CHUNK_FRAMES", "VLLM_GEPARD_LOOKBACK_FRAMES", "VLLM_OMNI_FISH_KVCACHE_ATTN", + "VLLM_OMNI_MOSS_LOCAL_COMPILE_AUDIO_SAMPLER", + "VLLM_OMNI_MOSS_LOCAL_SHORT_ATTN", + "VLLM_OMNI_MOSS_REF_ATTN", + "VLLM_OMNI_MOSS_REF_BATCHED_TRANSFER", + "VLLM_OMNI_MOSS_REF_BATCH_WINDOW_MS", + "VLLM_OMNI_MOSS_REF_CACHE_ROPE", + "VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR", + "VLLM_OMNI_MOSS_REF_COMPILE", + "VLLM_OMNI_MOSS_REF_ENCODER_WORKERS", + "VLLM_OMNI_MOSS_REF_FA_VERSION", + "VLLM_OMNI_MOSS_REF_GPU_PREP", + "VLLM_OMNI_MOSS_REF_GRAPHS", + "VLLM_OMNI_MOSS_REF_HOST_WINDOW_MS", + "VLLM_OMNI_MOSS_REF_INFLIGHT", + "VLLM_OMNI_MOSS_REF_SHARED_ENCODER", + "VLLM_OMNI_MOSS_REF_SINGLETON_BUCKET_MS", + "VLLM_OMNI_MOSS_REF_SKIP_PADDED_QUERY_MASK", + "VLLM_OMNI_MOSS_REF_STREAM_PRIORITY", "VLLM_OMNI_SANA_WM_STAGE1_TEXT_ENCODER", "VLLM_OMNI_SEEDVR2_LONG_JOB_TTL_SECONDS", "VLLM_OMNI_SEEDVR2_LONG_MAX_UPLOAD_BYTES", diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index cefa2a37a..b103776fa 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -3,7 +3,9 @@ from __future__ import annotations +import math import os +import queue import time from collections import defaultdict from typing import TYPE_CHECKING, Any @@ -58,6 +60,48 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): self._init_omni_io_scheduling_state() self._retains_state_across_chunks = bool(getattr(model_config, "retains_state_across_chunks", False)) self._pending_finish_reqs: list[Request] = [] + connector_config = getattr(model_config, "stage_connector_config", None) + extra = ( + connector_config.get("extra", connector_config) + if isinstance(connector_config, dict) + else getattr(connector_config, "extra", None) + ) or {} + self._generation_min_batch_size = int(extra.get("generation_min_batch_size", 1)) + self._generation_max_wait_s = float(extra.get("generation_max_wait_ms", 0)) / 1000 + self._generation_coalescing_policy = extra.get("generation_coalescing_policy", "fixed") + self._generation_max_regular_batch = int(extra.get("generation_max_regular_batch", 0) or 0) + if self._generation_max_regular_batch < 0: + raise ValueError("generation_max_regular_batch must be nonnegative") + if self._generation_coalescing_policy not in ("fixed", "idle_wait"): + raise ValueError("generation_coalescing_policy must be fixed or idle_wait") + if not 1 <= self._generation_min_batch_size <= self.max_num_running_reqs: + raise ValueError("generation_min_batch_size must be in [1, max_num_seqs]") + if not math.isfinite(self._generation_max_wait_s) or self._generation_max_wait_s < 0: + raise ValueError("generation_max_wait_ms must be finite and nonnegative") + if self._generation_max_wait_s and self._generation_min_batch_size > 1: + parallel = self.vllm_config.parallel_config + if not self._native_data_plane: + # CUDA-oriented profiles can inherit a platform's V1 runner. + # Their connector extras survive that override, but batching + # must revert to the fallback runner's default scheduling. + logger.warning( + "Generation batch waiting and regular-batch limit ignored: " + "requires native MRV2; using default generation scheduling." + ) + self._generation_min_batch_size = 1 + self._generation_max_wait_s = 0.0 + self._generation_max_regular_batch = 0 + elif not ( + self._retains_state_across_chunks + and parallel.tensor_parallel_size == parallel.pipeline_parallel_size == 1 + ): + raise ValueError("Generation batch waiting requires a stateful native MRV2 TP1/PP1 stage") + else: + logger.info( + "Generation input coalescing: target batch=%d, max wait=%.3f ms", + self._generation_min_batch_size, + self._generation_max_wait_s * 1000, + ) self._first_chunk_express = os.environ.get("VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS") == "1" self._last_step_express = False # Streams that already had a chunk scheduled; the rest await their first. @@ -76,6 +120,101 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): self._express_min_slack_s, ) + def _drain_omni_connector_outputs(self): + self._generation_defer_batch = False + outputs = super()._drain_omni_connector_outputs() + max_wait = getattr(self, "_generation_max_wait_s", 0) + target = getattr(self, "_generation_min_batch_size", 1) + coordinator = self.input_coordinator + if not max_wait or target <= 1 or coordinator is None or self._pause_state != PauseState.UNPAUSED: + return outputs + + # Count distinct runnable streams, not notifications, padding rows, or + # chunks whose previous execution still owns their state. Waiting here + # precedes allocation and metadata replacement. Do not schedule a batch + # and then discard it to try to accumulate more input. + ready = set(coordinator.requests_with_ready_chunks) + terminal = set(coordinator.finished_requests) | set(coordinator.input_terminal_req_ids) + deadline = getattr(self, "_generation_batch_deadline", None) + if deadline is None: + deadline = time.monotonic() + max_wait + in_flight = any( + not request.is_finished() and request.num_in_flight_tokens > 0 for request in self.requests.values() + ) + new_outputs = outputs + waited_without_input = False + while True: + for output in new_outputs: + ready.update(output.chunk_ready_req_ids) + terminal.update(output.chunk_finished_req_ids) + terminal.update( + req_id for req_id, metadata in output.request_metadata.items() if metadata.get("input_terminal") + ) + runnable = { + req_id + for req_id in ready | terminal + if (request := self.requests.get(req_id)) is not None + and not request.is_finished() + and request.num_in_flight_tokens == 0 + } + if runnable and waited_without_input: + # Idle waiting must not eat the ready batch's coalescing budget. + # Start that budget at its first arrival, once only. The two + # bounded windows can delay control processing by at most 2W. + deadline = time.monotonic() + max_wait + waited_without_input = False + # A first/final chunk participates in the same bounded window. + # At high concurrency almost every batch contains one; bypassing + # the entire batch in that case defeats coalescing. Low concurrency + # still dispatches immediately, including single-request tails. + if not runnable: + # With no output to retire and all receivers already parked, + # a wakeable inbox wait avoids repeatedly executing empty steps. + # Admission/registration/terminal work must reach the runner + # first; waiting for a not-yet-registered receiver would stall it. + can_wait_idle = ( + getattr(self, "_generation_coalescing_policy", "fixed") == "idle_wait" + and not in_flight + and bool(self.requests) + and not getattr(coordinator, "pending_chunk_registrations", ()) + and not getattr(coordinator, "pending_input_registrations", ()) + and not getattr(self, "_pending_data_plane_terminal_req_ids", ()) + and all(request.status == RequestStatus.WAITING_FOR_CHUNK for request in self.requests.values()) + ) + if not can_wait_idle: + break + waited_without_input = True + elif len(runnable) >= target or len(self.requests) < target: + break + remaining = deadline - time.monotonic() + if remaining <= 0: + break + if in_flight: + # Let EngineCore retire the prior output before waiting for + # more input. Blocking this thread also blocks the in-flight + # fence that makes those streams eligible for the next batch. + # Emit a normal empty scheduler step (including connector and + # cancellation bookkeeping), never discard an allocated step. + self._generation_defer_batch = True + self._generation_batch_deadline = deadline + return outputs + try: + output = self._omni_connector_output_inbox.get(timeout=remaining) + except queue.Empty: + break + outputs.append(output) + new_outputs = [output] + # Drain a burst before re-evaluating, retaining notification order. + while True: + try: + output = self._omni_connector_output_inbox.get_nowait() + except queue.Empty: + break + outputs.append(output) + new_outputs.append(output) + self._generation_batch_deadline = None + return outputs + def _build_generation_scheduler_output( self, new_reqs_data: list[OmniNewRequestData], @@ -256,6 +395,11 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): self._drop_aborted_queued_requests() self._requeue_completed_native_chunks() self._process_pending_omni_inputs(model_mode="generation") + if getattr(self, "_generation_defer_batch", False): + token_budget = 0 + max_regular = getattr(self, "_generation_max_regular_batch", 0) + if max_regular: + execution_batch_size = min(execution_batch_size, max_regular) self._drop_aborted_queued_requests() self._resync_streaming_input_counter() async_chunk_transport = self._async_chunk_transport_enabled() diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 6f9f03d6c..105b79a00 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -331,13 +331,8 @@ class OmniSchedulerMixin(_SchedulerMixinBase): # ready-chunk marker remains until scheduler admission succeeds. replaced_ids.discard(request_id) - def _consume_pending_connector_output(self, model_mode: str) -> None: - """Drain ``self._latest_omni_connector_output`` into the coordinator. - - Called at the top of every ``schedule()`` cycle. Identical between - AR and generation schedulers except for the ``model_mode`` argument - forwarded to ``update_request_metadata``. - """ + def _drain_omni_connector_outputs(self) -> list[OmniConnectorOutput]: + """Collect control messages in order, without changing request state.""" connector_outputs: list[OmniConnectorOutput] = [] inbox = getattr(self, "_omni_connector_output_inbox", None) if inbox is not None: @@ -350,6 +345,11 @@ class OmniSchedulerMixin(_SchedulerMixinBase): self._latest_omni_connector_output = None if connector_output is not None: connector_outputs.append(connector_output) + return connector_outputs + + def _consume_pending_connector_output(self, model_mode: str) -> None: + """Drain input notifications into the coordinator on the scheduler thread.""" + connector_outputs = self._drain_omni_connector_outputs() input_coordinator = getattr(self, "input_coordinator", None) if input_coordinator is None: return diff --git a/vllm_omni/core/sched/omni_scheduling_coordinator.py b/vllm_omni/core/sched/omni_scheduling_coordinator.py index dee3a9b43..f3042b736 100644 --- a/vllm_omni/core/sched/omni_scheduling_coordinator.py +++ b/vllm_omni/core/sched/omni_scheduling_coordinator.py @@ -379,14 +379,17 @@ class OmniSchedulingCoordinator: ) if model_mode != "ar": - new_ids = self._flatten_prompt_token_ids(metadata.get("code_predictor_codes")) + new_codes = metadata.get("code_predictor_codes") + new_ids = self._flatten_prompt_token_ids(new_codes) runtime_seed = None if "left_context_size" in metadata: runtime_seed = { "meta": {"left_context_size": metadata["left_context_size"]}, } request._omni_initial_model_buffer = runtime_seed - if new_ids: + # An explicit empty snapshot clears a previous codec chunk + # or prewarmed placeholder. Missing codes leave it intact. + if new_codes is not None: request.prompt_token_ids = new_ids request.num_prompt_tokens = len(new_ids) request._all_token_ids.clear() diff --git a/vllm_omni/deploy/moss_tts_local.yaml b/vllm_omni/deploy/moss_tts_local.yaml index 77c2a42b6..c351041f5 100644 --- a/vllm_omni/deploy/moss_tts_local.yaml +++ b/vllm_omni/deploy/moss_tts_local.yaml @@ -18,6 +18,9 @@ connectors: connector_get_max_wait: 300 initial_codec_chunk_frames: 1 # 1 frame * 80 ms/frame = 80 ms first audio chunk codec_chunk_frames: 15 # 15 frames * 80 ms/frame = 1.2 s steady chunks + # Optional startup ramp; its first entry overrides initial chunk size. + # Extra graph shapes/calls trade codec cost and latency for less underrun. + # codec_chunk_ramp: [1, 4, 15] codec_left_context_frames: 0 # Chunk ramp (opt-in): gradual chunk-size ladder replacing the # initial_codec_chunk_frames -> codec_chunk_frames cliff. The last diff --git a/vllm_omni/deploy/moss_tts_local_mrv2.yaml b/vllm_omni/deploy/moss_tts_local_mrv2.yaml new file mode 100644 index 000000000..4b25a41c7 --- /dev/null +++ b/vllm_omni/deploy/moss_tts_local_mrv2.yaml @@ -0,0 +1,30 @@ +# Opt-in CUDA MRV2 for MOSS-TTS-Local-Transformer-v1.5. +# Keep the base profile's batching, chunk geometry and V1 default. +base_config: moss_tts_local.yaml +model_runner: v2 + +platforms: + cuda: + stages: + - stage_id: 0 + # Bound prefill work so new prompts do not monopolize streaming decode. + max_num_batched_tokens: 512 + - stage_id: 1 + compilation_config: + # The codec owns CUDA Graph capture/replay. Skip Inductor compilation + # of its stateful decoder; plain CUDA graphs remain enabled. + mode: 0 + # FULL works without Inductor; the default FULL_AND_PIECEWISE would + # normalize to NONE and erase the codec's capture sizes. + cudagraph_mode: FULL + # Platform overrides replace compilation_config as a whole. + cudagraph_capture_sizes: [1, 2, 4, 8, 16, 32, 64] + cudagraph_num_of_warmups: 1 + npu: + model_runner: v1 + xpu: + model_runner: v1 + rocm: + model_runner: v1 + musa: + model_runner: v1 diff --git a/vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency.yaml b/vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency.yaml new file mode 100644 index 000000000..eb155f284 --- /dev/null +++ b/vllm_omni/deploy/moss_tts_local_mrv2_high_concurrency.yaml @@ -0,0 +1,53 @@ +# Opt-in CUDA profile for MOSS-TTS-Local-Transformer-v1.5 on a large-memory GPU. +# Both stages share GPU 0. Validated with one H200 and Seed-TTS EN1088. +base_config: moss_tts_local_mrv2.yaml + +connectors: + shm: + extra: + # Coalesce ready codec streams before submitting a small execution batch. + # This is a dispatch target, not a cap: execution capacity remains 128. + generation_min_batch_size: 32 + generation_max_wait_ms: 12 + # Optional; preserves the first/steady sizes while adding a middle chunk. + # codec_chunk_ramp: [1, 4, 15] + +# Optional codec attention-ring headroom (see recipes/OpenMOSS/MOSS-TTS.md): +# sizes each ring as context + one chunk so chunked streaming matches +# whole-sequence decoding. Requires additional state/graph memory; evaluate +# throughput and audio quality before enabling for your context and capacity. +# Add codec_ring_headroom: 1 under connectors.shm.extra above. + +platforms: + cuda: + stages: + - stage_id: 0 + max_num_seqs: 128 + max_num_batched_tokens: 512 + hf_overrides: + mrv2_gpu_slot_state: true + # Throughput-oriented pair; measure first-packet latency for your load. + mrv2_batch_prefill: true + mrv2_direct_tokens: true + # Mixed prefill/decode fits the same bounded token budget as capture. + compilation_config: + mode: 3 + cudagraph_mode: FULL + max_cudagraph_capture_size: 512 + cudagraph_num_of_warmups: 1 + # Leave room for the codec state pool and its batch-128 CUDA graphs. + kv_cache_memory_bytes: 34359738368 # 32 GiB + - stage_id: 1 + max_num_seqs: 128 + hf_overrides: + codec_attention_backend: triton_slot + compilation_config: + mode: 3 + backend: inductor + # The codec owns capture/replay after Inductor compiles its decoder. + cudagraph_mode: FULL + cudagraph_capture_sizes: [1, 2, 4, 6, 8, 12, 16, 24, 32, 64, 128] + cudagraph_num_of_warmups: 1 + inductor_compile_config: + combo_kernels: false + benchmark_combo_kernel: false diff --git a/vllm_omni/deploy/moss_tts_local_mrv2_low_latency.yaml b/vllm_omni/deploy/moss_tts_local_mrv2_low_latency.yaml new file mode 100644 index 000000000..959238ee2 --- /dev/null +++ b/vllm_omni/deploy/moss_tts_local_mrv2_low_latency.yaml @@ -0,0 +1,66 @@ +# Opt-in CUDA profile for MOSS-TTS-Local-Transformer-v1.5 tuned for first-packet +# latency at client concurrency 128 on one H200 (Seed-TTS EN1088, closed loop). +# Both stages share GPU 0. Use moss_tts_local_mrv2_high_concurrency.yaml for +# throughput-first serving. +# +# Optional launch-level settings (see the full command in recipes/OpenMOSS/MOSS-TTS.md): +# - CUDA MPS for the GPU (stage 0 below lowers its own MPS client priority); +# - VLLM_OMNI_EVENT_DRIVEN_ORCH=1; +# - --api-server-count 4 with VLLM_WORKER_MULTIPROC_METHOD=spawn and +# VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR=/dev/shm/ so frontends share reference codes; +# - VLLM_OMNI_CONNECTOR_RECV_POLL_MS=1 so Stage 1 picks up first chunks sooner. +base_config: moss_tts_local_mrv2.yaml + +connectors: + shm: + extra: + generation_min_batch_size: 32 + generation_max_wait_ms: 12 + # Stage 1 decodes each first chunk from its receive thread on a small + # dedicated graph set and emits it directly, so first chunks no longer + # wait for the engine loop or a regular codec step. + codec_first_chunk_fast_path: 1 + # Regular codec replays wait for an in-flight first-chunk decode. + codec_first_chunk_gate: 1 + # Short regular codec steps bound contention with first-chunk decodes. + generation_max_regular_batch: 16 + # Optional ramp. Both the first-chunk graph and regular graphs follow it. + # codec_chunk_ramp: [1, 4, 15] + +platforms: + cuda: + stages: + - stage_id: 0 + max_num_seqs: 128 + # Larger prefill budget admits bursts of new requests in fewer steps. + max_num_batched_tokens: 2048 + hf_overrides: + mrv2_gpu_slot_state: true + mrv2_batch_prefill: true + mrv2_direct_tokens: true + # Sample each frame right after the forward that produced it. + mrv2_eager_mtp: true + compilation_config: + mode: 3 + cudagraph_mode: FULL + max_cudagraph_capture_size: 2048 + cudagraph_num_of_warmups: 1 + kv_cache_memory_bytes: 34359738368 # 32 GiB + env: + # Under MPS, let codec kernels schedule ahead of the talker's. + CUDA_MPS_CLIENT_PRIORITY: "1" + - stage_id: 1 + max_num_seqs: 128 + # Keep synchronous scheduling; benchmark async scheduling separately. + async_scheduling: false + hf_overrides: + codec_attention_backend: triton_slot + compilation_config: + mode: 3 + backend: inductor + cudagraph_mode: FULL + cudagraph_capture_sizes: [1, 2, 4, 6, 8, 12, 16, 24, 32, 64, 128] + cudagraph_num_of_warmups: 1 + inductor_compile_config: + combo_kernels: false + benchmark_combo_kernel: false diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py index 7d2b16d50..d7744ffb1 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py @@ -1169,7 +1169,12 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): self._payload_finished(payload_data), ) - self._get_req_chunk[req_id] += 1 + with self._lock: + if self._async_chunk and request is not None and req_id not in self._pending_load_reqs: + # A connector get can return after cancellation removed its + # receiver. Do not recreate delivery state for that request. + return False + self._get_req_chunk[req_id] += 1 if self._async_chunk: is_finished = self._payload_finished(payload_data) @@ -1180,8 +1185,16 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): if not self._payload_value_has_content(new_ids) and not is_finished: return False payload_consumable = self._payload_is_consumable(payload_data) + first_chunk_hook = getattr(self, "_first_chunk_hook", None) + if first_chunk_hook is not None and chunk_id == 0 and not is_finished and payload_consumable: + first_chunk_hook(req_id, request, payload_data) with self._lock: + if request is not None and req_id not in self._pending_load_reqs: + # The first-chunk hook may have overlapped cancellation. + # Its slot cleanup is ordered by the model; this payload + # must not republish readiness after receive teardown. + return False if self._model_mode == "ar": # Accumulation, staging, and model-side consume/ack share # this lock. Keeping the transition atomic prevents the @@ -1723,3 +1736,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): sink: Callable[[OmniConnectorOutput], None] | None, ) -> None: self._omni_connector_output_sink = sink + + def set_first_chunk_hook(self, hook: Callable[[str, Any, OmniPayload], bool]) -> None: + """Claim eligible first chunks before publishing them to the model loop.""" + self._first_chunk_hook = hook diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index f9fcd34b9..c84d9ed18 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -80,9 +80,12 @@ def _signal_exit_code(signum: int) -> int: def _bind_first_audio_sink(model_executor: Any, output_queue: Any, scheduler: Any) -> bool: - """Let a TP1 in-process runner that decodes first audio deliver it as its own output.""" + """Bind in-process first audio to the generic engine output sink.""" if not isinstance(model_executor, UniProcExecutor): return False + parallel_config = model_executor.vllm_config.parallel_config + if parallel_config.tensor_parallel_size != 1 or parallel_config.pipeline_parallel_size != 1: + return False worker = getattr(getattr(model_executor, "driver_worker", None), "worker", None) model_runner = getattr(worker, "model_runner", None) model_state = getattr(model_runner, "model_state", None) @@ -90,12 +93,23 @@ def _bind_first_audio_sink(model_executor: Any, output_queue: Any, scheduler: An decodes_audio = getattr(model, "first_frame_decoder", None) is not None or ( getattr(model, "stream_decoder", None) is not None and bool(getattr(model, "stream_first_audio", False)) ) - if not decodes_audio or not hasattr(model_state, "set_first_audio_sink"): - return False from vllm_omni.worker_v2.first_audio_sender import engine_output_queue_sink - assert model_state is not None - model_state.set_first_audio_sink(engine_output_queue_sink(output_queue, scheduler)) + if decodes_audio and hasattr(model_state, "set_first_audio_sink"): + assert model_state is not None + model_state.set_first_audio_sink(engine_output_queue_sink(output_queue, scheduler)) + return True + data_plane = getattr(model_runner, "_omni_data_plane", None) + get_model = getattr(model_runner, "get_model", None) + if data_plane is None or not callable(get_model): + return False + bind = getattr(get_model(), "bind_first_chunk_fast_path", None) + if not callable(bind): + return False + hook = bind(engine_output_queue_sink(output_queue, scheduler, upstream_first_audio=False)) + if hook is None: + return False + data_plane.set_first_chunk_hook(hook) return True diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 79217ca8b..ca86587e9 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -441,7 +441,7 @@ class OmniBase(PDDisaggregationMixin): if self.log_stats: # Emit per-request orchestrator timing (including e2e_total_ms) # before dropping request state. - req_state.metrics.build_and_log_summary() + req_state.metrics.log_timing_summary() except Exception: logger.exception( "[%s] Failed to build/log summary for req=%s", diff --git a/vllm_omni/entrypoints/openai/tts_adapters/base.py b/vllm_omni/entrypoints/openai/tts_adapters/base.py index 6ad7dcdcc..89eb463d9 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/base.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/base.py @@ -81,13 +81,24 @@ def conditioning_cache_salt( This replaces the uploaded data URI in the hash without changing the request. """ h = hashlib.sha256() + ref_audio: str | list[str] | tuple[str, int] | None = request.ref_audio + if ( + isinstance(ref_audio, str) + and len(ref_audio) > 256 + and tts_params is not None + and tts_params.get("ref_audio_cache_key") is not None + ): + # The resolve key below is already a content-aware digest of this + # locator. Hashing a repr of a large inline data URI again dominates + # the per-request serving cost without adding information. + ref_audio = ("ref_audio_cache_key", len(ref_audio)) for part in ( request.input, request.task_type, request.language, request.voice, request.ref_text, - registered_voice if registered_voice is not None else request.ref_audio, + registered_voice if registered_voice is not None else ref_audio, request.instructions, request.x_vector_only_mode, request.speaker_embedding, diff --git a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py index db2833f42..ed7865789 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py @@ -32,6 +32,59 @@ if TYPE_CHECKING: logger = init_logger(__name__) +def _local_continuation_prompt(proc: Any, user_kwargs: dict[str, Any], reference: Any) -> tuple[list[int], Any] | None: + """Build the Local-v1.5 continuation prompt without the generic processor call. + + Mirrors ``_build_continuation_codes`` of the model's processor: every + template segment is encoded separately there, so the constant segments + (keyed by ``language``) are cached and only the request text is + tokenized. Returns ``(text_ids, audio_codes)`` equal to + ``unified[:, 0].tolist()`` / ``unified[:, 1:]``, or ``None`` when the + request or processor does not match that exact layout. + """ + import torch + + if set(user_kwargs) - {"text", "language"}: + return None + if not isinstance(reference, list) or len(reference) != 1 or not isinstance(reference[0], torch.Tensor): + return None + cfg = getattr(proc, "model_config", None) + n_vq = getattr(cfg, "n_vq", None) + if cfg is None or not isinstance(n_vq, int) or n_vq <= 0: + return None + codes = reference[0] + if codes.ndim != 2 or int(codes.shape[1]) != n_vq or codes.shape[0] == 0: + return None + required = ( + "_user_prompt_prefix_ids", + "_encode_text", + "_user_prompt_after_reference_ids", + "_assistant_prompt_prefix_ids", + ) + if not all(callable(getattr(proc, name, None)) for name in required): + return None + language = user_kwargs.get("language") + cache = proc.__dict__.setdefault("_vllm_omni_continuation_segments", {}) + segments = cache.get(language) + if segments is None: + fields = {} if language is None else {"language": language} + head = ( + proc._user_prompt_prefix_ids() + + proc._encode_text("None") + + proc._user_prompt_after_reference_ids(language, fields) + ) + tail = proc._assistant_prompt_prefix_ids() + [int(cfg.audio_start_token_id)] + segments = cache[language] = ([int(t) for t in head], [int(t) for t in tail]) + head, tail = segments + text = user_kwargs.get("text") + prompt_ids = head + [int(t) for t in proc._encode_text("" if text is None else str(text))] + tail + num_audio = int(codes.shape[0]) + text_ids = prompt_ids + [int(cfg.audio_assistant_slot_token_id)] * num_audio + audio_codes = torch.full((len(text_ids), n_vq), int(cfg.audio_pad_token_id), dtype=torch.int64) + audio_codes[len(prompt_ids) :] = codes.to(dtype=torch.int64, device="cpu") + return text_ids, audio_codes + + class _MossTTSAdapterBase(ARTTSAdapter): accumulate_nonstreaming: bool = False @@ -82,6 +135,28 @@ class _MossTTSAdapterBase(ARTTSAdapter): self._moss_ref_encoder = encoder return encoder + async def warmup(self) -> None: + """Set up the reference encoder at startup, not on the first cloning request. + + Its CUDA graph capture takes seconds per API process (about a minute for + a shared encoder's workers). + """ + if self._moss_variant not in ("tts", "local", "ttsd"): + return + from vllm_omni.model_executor.models.moss_tts.reference_encoder import reference_graphs_enabled + from vllm_omni.model_executor.models.moss_tts.shared_reference_encoder import SharedReferenceEncoderStartupError + + if not reference_graphs_enabled(): + return + try: + await asyncio.to_thread(self._get_moss_ref_encoder().prepare) + except SharedReferenceEncoderStartupError: + # Do not advertise readiness and trigger concurrent local captures + # while the shared host is still compiling or has failed startup. + raise + except Exception: # noqa: BLE001 - requests still encode (and capture) on demand + logger.warning("MOSS reference encoder warmup failed", exc_info=True) + async def _encode_moss_references( self, request: "OpenAICreateSpeechRequest", @@ -269,8 +344,16 @@ class _MossTTSAdapterBase(ARTTSAdapter): proc = AutoProcessor.from_pretrained(model_id, trust_remote_code=True) if hasattr(proc, "audio_tokenizer"): 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) + proc._vllm_omni_ref_encoder_device = device + from vllm_omni.model_executor.models.moss_tts.reference_encoder import shared_encoder_role + + if shared_encoder_role() == "client": + # Another API process hosts the shared encoder on the GPU. + proc.audio_tokenizer = proc.audio_tokenizer.eval() + logger.info("MOSS reference-audio encoder: using the shared encoder of another API process") + else: + 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 @@ -405,18 +488,24 @@ class _MossTTSAdapterBase(ARTTSAdapter): # Build the unified-codes prompt: (L, 1+n_vq) where col 0 is text/special # tokens and cols 1..n_vq are the delay-pattern audio code grid (mostly # audio_pad_code outside the reference block). + fast = None if v == "local" and request.ref_text and request.ref_text.strip(): reference = user_kwargs.pop("reference") user_kwargs["text"] = request.ref_text.strip() + " " + (request.input or "") - user_msg = proc.build_user_message(**user_kwargs) - assistant_msg = proc.build_assistant_message(audio_codes_list=reference) - batch = proc(conversations=[[user_msg, assistant_msg]], mode="continuation") + fast = _local_continuation_prompt(proc, user_kwargs, reference) + if fast is None: + user_msg = proc.build_user_message(**user_kwargs) + assistant_msg = proc.build_assistant_message(audio_codes_list=reference) + batch = proc(conversations=[[user_msg, assistant_msg]], mode="continuation") else: user_msg = proc.build_user_message(**user_kwargs) batch = proc(conversations=[[user_msg]], mode="generation") - unified = batch["input_ids"][0] # torch.LongTensor (L, 1+n_vq) - text_ids: list[int] = unified[:, 0].tolist() - audio_codes: torch.Tensor = unified[:, 1:].contiguous().to(torch.int64) + if fast is not None: + text_ids, audio_codes = fast + else: + unified = batch["input_ids"][0] # torch.LongTensor (L, 1+n_vq) + text_ids = unified[:, 0].tolist() + audio_codes = unified[:, 1:].contiguous().to(torch.int64) params = { "prompt_token_ids": text_ids, diff --git a/vllm_omni/metrics/stats.py b/vllm_omni/metrics/stats.py index df5183f45..edb9a87db 100644 --- a/vllm_omni/metrics/stats.py +++ b/vllm_omni/metrics/stats.py @@ -4,6 +4,7 @@ from __future__ import annotations import copy +import logging import time from collections import defaultdict from collections.abc import Callable, Sequence @@ -851,9 +852,62 @@ class OrchestratorAggregator: for cached_key in [k for k in self._replica_cache if k[1] == rid_key]: self._replica_cache.pop(cached_key, None) + def log_timing_summary(self) -> None: + """Emit only the concise per-request ``[OmniTiming]`` lines. + + The serving path uses this unless DEBUG tables are requested; it skips + the table rows that ``build_and_log_summary`` builds for its result. + """ + if not self.log_stats: + return + if logger.isEnabledFor(logging.DEBUG): + self.build_and_log_summary() + return + for rid in sorted(set(self.stage_events.keys()) | {e.request_id for e in self.e2e_events}): + e2e_evt = next((e for e in self.e2e_events if e.request_id == rid), None) + stage_evts = sorted( + self.stage_events.get(rid, []), + key=lambda e: e.stage_id if e.stage_id is not None else -1, + ) + self._log_omni_timing(rid, e2e_evt, stage_evts) + + def _log_omni_timing(self, rid: str, e2e_evt: RequestE2EStats | None, stage_evts: list[StageRequestStats]) -> None: + pt: dict[str, float] = {} + if stage_evts: + pt = stage_evts[-1].pipeline_timings or {} + if pt or e2e_evt: + parts = [f"req={rid}"] + if e2e_evt: + parts.append(f"total={e2e_evt.e2e_total_ms / 1000.0:.2f}s") + if "preprocess_ms" in pt: + parts.append(f"preprocess={pt['preprocess_ms'] / 1000.0:.2f}s") + if e2e_evt: + engine_ms = e2e_evt.e2e_total_ms - pt.get("preprocess_ms", 0.0) + parts.append(f"engine={engine_ms / 1000.0:.2f}s") + stage_parts = [] + for evt in stage_evts: + sid = evt.stage_id if evt.stage_id is not None else "?" + t = evt.stage_gen_time_ms / 1000.0 + stage_parts.append(f"{sid}:{t:.2f}s") + if stage_parts: + parts.append(f"stages=[{','.join(stage_parts)}]") + transfer_parts = [] + for te in self.transfer_events.values(): + if te.request_id == rid: + transfer_parts.append(f"{te.from_stage}->{te.to_stage}={te.tx_time_ms:.2f}ms") + if transfer_parts: + parts.append(f"transfers=[{','.join(transfer_parts)}]") + if "ar2diffusion_ms" in pt: + parts.append(f"ar2diffusion={pt['ar2diffusion_ms']:.2f}ms") + logger.info("[OmniTiming] %s", " ".join(parts)) + def build_and_log_summary(self) -> dict[str, Any]: if not self.log_stats: return {} + # Per-request tables are diagnostic. Formatting them on the serving + # event loop costs more than the request's own bookkeeping, so they + # are emitted at DEBUG; the concise [OmniTiming] line stays at INFO. + log_tables = logger.isEnabledFor(logging.DEBUG) wall_time_ms = max(0.0, (self.last_finish_ts - self.wall_start_ts) * 1000.0) e2e_avg_req = (wall_time_ms / self.e2e_count) if self.e2e_count > 0 else 0.0 e2e_avg_tok = (self.e2e_total_tokens * 1000.0 / wall_time_ms) if wall_time_ms > 0 else 0.0 @@ -895,8 +949,8 @@ class OrchestratorAggregator: v = overall_summary.get(k, None) if v not in (0, 0.0, 0.000, None, ""): overall_fields.append(k) - if overall_fields: - logger.info( + if log_tables and overall_fields: + logger.debug( "\n%s", _format_table("Overall Summary", overall_summary, overall_fields), ) @@ -921,8 +975,8 @@ class OrchestratorAggregator: nonzero_e2e_fields.add(k) value_fields_e2e = sorted(nonzero_e2e_fields) - if value_fields_e2e: - logger.info( + if log_tables and value_fields_e2e: + logger.debug( "\n%s", _format_table( f"RequestE2EStats [request_id={rid}]", @@ -936,34 +990,7 @@ class OrchestratorAggregator: self.stage_events.get(rid, []), key=lambda e: e.stage_id if e.stage_id is not None else -1, ) - pt = {} - if stage_evts: - pt = stage_evts[-1].pipeline_timings or {} - if pt or e2e_evt: - parts = [f"req={rid}"] - if e2e_evt: - parts.append(f"total={e2e_evt.e2e_total_ms / 1000.0:.2f}s") - if "preprocess_ms" in pt: - parts.append(f"preprocess={pt['preprocess_ms'] / 1000.0:.2f}s") - if e2e_evt: - engine_ms = e2e_evt.e2e_total_ms - pt.get("preprocess_ms", 0.0) - parts.append(f"engine={engine_ms / 1000.0:.2f}s") - stage_parts = [] - for evt in stage_evts: - sid = evt.stage_id if evt.stage_id is not None else "?" - t = evt.stage_gen_time_ms / 1000.0 - stage_parts.append(f"{sid}:{t:.2f}s") - if stage_parts: - parts.append(f"stages=[{','.join(stage_parts)}]") - transfer_parts = [] - for te in self.transfer_events.values(): - if te.request_id == rid: - transfer_parts.append(f"{te.from_stage}->{te.to_stage}={te.tx_time_ms:.2f}ms") - if transfer_parts: - parts.append(f"transfers=[{','.join(transfer_parts)}]") - if "ar2diffusion_ms" in pt: - parts.append(f"ar2diffusion={pt['ar2diffusion_ms']:.2f}ms") - logger.info("[OmniTiming] %s", " ".join(parts)) + self._log_omni_timing(rid, e2e_evt, stage_evts) # === Stage table (columns = stage_id) === # if any stage has diffusion_metrics, remove postprocess_time_ms field @@ -987,7 +1014,7 @@ class OrchestratorAggregator: result_stage_table.append({"request_id": rid, "stages": stage_rows}) - if stage_rows: + if log_tables and stage_rows: # filter out all-zero fields for logging all_value_fields = set() for row in stage_rows: @@ -1006,7 +1033,7 @@ class OrchestratorAggregator: value_fields_list.append(field) if value_fields_list: - logger.info( + logger.debug( "\n%s", _format_table( f"StageRequestStats [request_id={rid}]", @@ -1027,7 +1054,7 @@ class OrchestratorAggregator: ] result_trans_table.append({"request_id": rid, "transfers": transfer_rows}) - if transfer_rows: + if log_tables and transfer_rows: # filter out all-zero fields for logging all_value_fields = set() for row in transfer_rows: @@ -1046,7 +1073,7 @@ class OrchestratorAggregator: value_fields_list.append(field) if value_fields_list: - logger.info( + logger.debug( "\n%s", _format_table( f"TransferEdgeStats [request_id={rid}]", diff --git a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py index 4dbedbd19..89f32ee55 100644 --- a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py +++ b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py @@ -539,11 +539,15 @@ class RingKVCache: end_offset = self.end_offset.index_select(0, slots) row_cache = self.cache.index_select(1, slots) - indexes = torch.arange(T, device=end_offset.device, dtype=end_offset.dtype) + # Upsampling can make a chunk longer than the ring. Only the last + # capacity tokens survive; writing all T tokens gives scatter_ + # duplicate destinations and nondeterministic cache contents. + write_length = min(T, self.capacity) + indexes = torch.arange(T - write_length, T, device=end_offset.device, dtype=end_offset.dtype) indexes = (indexes + end_offset.view(-1, 1)) % self.capacity - scatter_indexes = indexes.view(B, 1, T, 1).expand(-1, H, T, D) - row_cache[0].scatter_(2, scatter_indexes, k) - row_cache[1].scatter_(2, scatter_indexes, v) + scatter_indexes = indexes.view(B, 1, write_length, 1).expand(-1, H, write_length, D) + row_cache[0].scatter_(2, scatter_indexes, k[:, :, -write_length:]) + row_cache[1].scatter_(2, scatter_indexes, v[:, :, -write_length:]) # Live and graph-padding rows always map to distinct slots. The # latter map only to scratch state, so this write cannot corrupt a # request even though dense graph operators still execute it. @@ -729,7 +733,13 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule): if self.context is None: capacity = self.weights_per_step if self.weights_per_step else 1024 else: - capacity = self.context + # RingKVCache.complete writes a whole chunk before attending, so a + # ring of exactly ``context`` entries lets a T-token chunk evict up + # to T tokens of history that the chunk's own queries should still + # see (and, for T > context, evict part of the chunk itself). + # ``_ring_headroom`` (max chunk tokens at this layer) keeps every + # in-window key resident; 0 preserves the legacy truncation. + capacity = self.context + int(getattr(self, "_ring_headroom", 0)) kv_cache = RingKVCache( batch_size, @@ -791,29 +801,63 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule): if self.rope: q, k = self.rope(q, k, offset, time_before_heads=False) - k, v, pos_k = self._complete_kv(k, v, execution_context) - pos_k = pos_k[:, None] - - if self.causal: - pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view(-1, 1) - delta = pos_q - pos_k - attn_bias = (pos_k >= 0) & (delta >= 0) - if self.context is not None: - attn_bias = attn_bias & (delta < self.context) - attn_bias = attn_bias[:, None] + slot_attention = getattr(self, "_slot_attention", None) + slot_attention_rows = getattr(self, "_slot_attention_rows", None) + advance_kv_offset = False + slot_path = state is not None and execution_context is not None and self.causal + if slot_path and slot_attention_rows is not None: + # One launch: the gathered MHA offset equals the ring's end_offset + # for every slot, and both advance together below. + valid_lengths = execution_context.valid_rows.to(dtype=torch.int32) * T + x = slot_attention_rows( + q, + k, + v, + state.kv_cache.cache, + offset, + execution_context.state_slot_ids, + valid_lengths, + self.context if self.context is not None else -1, + ) + advance_kv_offset = True + elif slot_attention is not None and slot_path: + # Current streaming batches use one exact T. Padding rows own + # scratch slots and must not advance persistent attention state. + valid_lengths = execution_context.valid_rows.to(dtype=torch.int32) * T + x = slot_attention( + q, + k, + v, + state.kv_cache.cache, + state.kv_cache.end_offset, + execution_context.state_slot_ids, + valid_lengths, + self.context if self.context is not None else -1, + ) else: - attn_bias = None + k, v, pos_k = self._complete_kv(k, v, execution_context) + pos_k = pos_k[:, None] - streaming_attention = getattr(self, "_streaming_attention", None) - if ( - streaming_attention is not None - and attn_bias is not None - and q.dtype == torch.bfloat16 - and q.shape[-1] == 64 - ): - x = streaming_attention(q, k, v, attn_bias) - else: - x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0) + if self.causal: + pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view(-1, 1) + delta = pos_q - pos_k + attn_bias = (pos_k >= 0) & (delta >= 0) + if self.context is not None: + attn_bias = attn_bias & (delta < self.context) + attn_bias = attn_bias[:, None] + else: + attn_bias = None + + streaming_attention = getattr(self, "_streaming_attention", None) + if ( + streaming_attention is not None + and attn_bias is not None + and q.dtype == torch.bfloat16 + and q.shape[-1] == 64 + ): + x = streaming_attention(q, k, v, attn_bias) + else: + x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0) x = x.transpose(1, 2).reshape(B, T, self.embed_dim) x = apply_weights_per_step(self.out_projs, self.weights_per_step_schedule, x, offset_cpu) @@ -824,6 +868,8 @@ class MossAudioTokenizerMultiheadAttention(StreamingModule): else: next_offset = torch.where(execution_context.valid_rows, offset + T, offset) state.offset.index_copy_(0, execution_context.state_slot_ids, next_offset) + if advance_kv_offset: + state.kv_cache.end_offset.index_copy_(0, execution_context.state_slot_ids, next_offset) return x @@ -1624,6 +1670,7 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): decoder_kwargs_list = copy.deepcopy(config.decoder_kwargs) self.decoder = nn.ModuleList() + decoder_input_frame_rate = current_frame_rate for decoder_kwargs_i in decoder_kwargs_list: decoder_kwargs_i = dict(decoder_kwargs_i) if decoder_kwargs_i["module_type"] == "PatchedPretransform": @@ -1638,6 +1685,11 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): context=int(round(current_frame_rate * context_duration)), ) ) + # Tokens this transformer sees per input code frame. Streaming + # ring headroom is sized from it (see initialize_decoder_state_pool). + self.decoder[-1].tokens_per_input_frame = max( + 1, int(round(current_frame_rate / decoder_input_frame_rate)) + ) current_frame_rate *= self.decoder[-1].downsample_ratio expected_output_frame_rate = float(self.sampling_rate * channel_interleave_factor) @@ -1731,10 +1783,29 @@ class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): self._decoder_state_capacity = 0 self._decoder_slot_offsets = None - def initialize_decoder_state_pool(self, state_capacity: int, scratch_capacity: int = 0) -> None: - """Allocate persistent decoder state independently of execution B.""" - if state_capacity <= 0 or scratch_capacity < 0: - raise ValueError(f"Invalid decoder state capacities: state={state_capacity}, scratch={scratch_capacity}.") + def initialize_decoder_state_pool( + self, state_capacity: int, scratch_capacity: int = 0, chunk_frames: int = 0 + ) -> None: + """Allocate persistent decoder state independently of execution B. + + ``chunk_frames`` > 0 sizes every decoder attention ring as + ``context + chunk_frames * tokens_per_input_frame`` so that a chunk of + up to ``chunk_frames`` code frames never evicts keys inside its own + causal window; chunked streaming then matches whole-sequence decoding. + ``chunk_frames`` = 0 keeps the legacy ``context``-sized rings. + """ + if state_capacity <= 0 or scratch_capacity < 0 or chunk_frames < 0: + raise ValueError( + "Invalid decoder state capacities: " + f"state={state_capacity}, scratch={scratch_capacity}, chunk_frames={chunk_frames}." + ) + for module in self.decoder: + if not isinstance(module, MossAudioTokenizerProjectedTransformer): + continue + tokens_per_frame = int(getattr(module, "tokens_per_input_frame", 1)) + for attention in module.modules(): + if isinstance(attention, MossAudioTokenizerMultiheadAttention): + attention._ring_headroom = chunk_frames * tokens_per_frame self._start_streaming(state_capacity + scratch_capacity, decoder_only=True) self._decoder_state_capacity = state_capacity + scratch_capacity diff --git a/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py b/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py index 4d277f44e..26e9f2fde 100644 --- a/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py +++ b/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py @@ -81,15 +81,31 @@ class CUDAGraphStreamingDecoderWrapper: frame_sizes: list[int], num_quantizers: int, vllm_config: VllmConfig, + scratch_base: int | None = None, + private_pool: bool = False, + compiled_decode: nn.Module | None = None, ) -> None: self.codec = codec self.state_capacity = int(state_capacity) + # Graph padding rows use scratch slots [scratch_base, scratch_base+B). + # A wrapper replaying concurrently with another needs its own scratch + # range and graph memory pool. + self.scratch_base = self.state_capacity if scratch_base is None else int(scratch_base) + self._private_pool = bool(private_pool) + # Graphs bake in the cuBLAS workspace of the stream they were captured + # on, and torch.cuda.graph shares one default capture stream. Split-K + # GEMMs synchronize through that workspace, so two graphs replaying + # concurrently from one capture stream can spin on each other forever. + self._capture_stream: torch.cuda.Stream | None = None self.batch_sizes = sorted({int(size) for size in batch_sizes if 0 < int(size) <= state_capacity}) self.frame_sizes = sorted({int(size) for size in frame_sizes if int(size) > 0}) self.num_quantizers = int(num_quantizers) self.graphs: dict[tuple[int, int], _CapturedStreamingDecodeGraph] = {} self._pool = None self._warmed_up = False + if compiled_decode is not None: + self._compiled_decode: nn.Module | None = compiled_decode + return # vLLM owns Inductor compilation; this wrapper remains the sole owner # of CUDA Graph capture/replay because it understands persistent codec # state slots. Disable vLLM's CUDA Graph layer to avoid nested graphs. @@ -98,7 +114,7 @@ class CUDAGraphStreamingDecoderWrapper: compile_config.compilation_config.cudagraph_mode = CUDAGraphMode.NONE compile_config.compilation_config.static_forward_context = {} with set_current_vllm_config(compile_config): - self._compiled_decode: nn.Module | None = _MossStreamingDecodeCompileAdapter( + self._compiled_decode = _MossStreamingDecodeCompileAdapter( codec, vllm_config=compile_config, ) @@ -175,7 +191,7 @@ class CUDAGraphStreamingDecoderWrapper: frame_size, exc_info=True, ) - scratch_slots = self.state_capacity + torch.arange( + scratch_slots = self.scratch_base + torch.arange( batch_size, dtype=torch.long, device=device, @@ -206,10 +222,13 @@ class CUDAGraphStreamingDecoderWrapper: device=device, ) lengths = torch.zeros(batch_size, dtype=torch.long, device=device) - scratch_slots = self.state_capacity + torch.arange(batch_size, dtype=torch.long, device=device) + scratch_slots = self.scratch_base + torch.arange(batch_size, dtype=torch.long, device=device) valid_rows = torch.zeros(batch_size, dtype=torch.bool, device=device) - stream = torch.cuda.Stream() + if self._private_pool and self._capture_stream is None: + self._capture_stream = torch.cuda.Stream(device=device) + # Warm up on the capture stream so its workspace exists before capture. + stream = self._capture_stream or torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): for _ in range(3): @@ -219,9 +238,11 @@ class CUDAGraphStreamingDecoderWrapper: self.codec.reset_decoder_state_slots(scratch_slots) if self._pool is None: - self._pool = current_platform.get_global_graph_pool() + self._pool = ( + torch.cuda.graph_pool_handle() if self._private_pool else current_platform.get_global_graph_pool() + ) graph = CUDAGraph() - with torch.cuda.graph(graph, pool=self._pool, capture_error_mode="thread_local"): + with torch.cuda.graph(graph, pool=self._pool, stream=self._capture_stream, capture_error_mode="thread_local"): audio, audio_lengths = decode(codes, lengths, scratch_slots, valid_rows) self.graphs[(batch_size, frame_size)] = _CapturedStreamingDecodeGraph( @@ -272,7 +293,7 @@ class CUDAGraphStreamingDecoderWrapper: entry.static_lengths.zero_() entry.static_lengths[:actual_batch_size].fill_(int(frame_size)) entry.static_state_slot_ids.copy_( - self.state_capacity + torch.arange(batch_size, dtype=torch.long, device=entry.static_state_slot_ids.device) + self.scratch_base + torch.arange(batch_size, dtype=torch.long, device=entry.static_state_slot_ids.device) ) entry.static_state_slot_ids[:actual_batch_size].copy_(state_slot_ids, non_blocking=True) entry.static_valid_rows.zero_() diff --git a/vllm_omni/model_executor/models/moss_tts/first_chunk_fast_path.py b/vllm_omni/model_executor/models/moss_tts/first_chunk_fast_path.py new file mode 100644 index 000000000..34361e547 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/first_chunk_fast_path.py @@ -0,0 +1,333 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Decode a stream's first codec chunk as soon as Stage 1 receives it. + +The Stage-1 engine loop only sees a new chunk between steps, and a step waits +for its codec graph to finish. A first chunk that arrives mid-step therefore +waits for the rest of that step, then for its own replay. This path takes the +first chunk from the receive thread, replays a small dedicated graph on its +own stream while the main step runs, and emits the audio straight to the +engine output queue. The main path later receives the same chunk, finds it +already decoded, and continues the stream from the state written here. +""" + +from __future__ import annotations + +import math +import queue +import threading +import time +from dataclasses import dataclass +from typing import Any + +import torch +from vllm.logger import init_logger + +from vllm_omni.platforms import current_omni_platform +from vllm_omni.worker_v2.first_audio_sender import FirstAudioSink, _PreparedDelivery + +logger = init_logger(__name__) + + +class _SlotHandoff: + """Orders the main path's first use of a slot after the fast decode.""" + + __slots__ = ("done", "event", "error") + + def __init__(self) -> None: + self.done = threading.Event() + self.event: torch.cuda.Event | None = None + self.error: Exception | None = None + + +@dataclass +class _Job: + request_id: str + request_key: str + delivery: _PreparedDelivery + slot: int + codes: torch.Tensor # (n_vq, T) int64 on host + handoff: _SlotHandoff + + +class MossFirstChunkFastPath: + def __init__( + self, + session: Any, + wrapper: Any, + *, + n_vq: int, + frames: int, + codebook_size: int, + samples_per_frame: int, + n_channels: int, + sample_rate: torch.Tensor, + device: torch.device, + gate_main: bool = False, + handoff_timeout_s: float = 30.0, + ) -> None: + if not math.isfinite(handoff_timeout_s) or handoff_timeout_s <= 0: + raise ValueError("handoff_timeout_s must be positive and finite") + self._handoff_timeout_s = float(handoff_timeout_s) + self._session = session + # When set, the main codec stream waits for an in-flight fast decode + # before its next replay, so the small first-chunk graph runs mostly + # uncontended instead of interleaving with a large regular replay. + self._gate_main = bool(gate_main) + self._inflight: torch.cuda.Event | None = None + self._wrapper = wrapper + self._n_vq = int(n_vq) + self.frames = int(frames) + self._codebook_size = int(codebook_size) + self._samples = int(frames) * int(samples_per_frame) + self._n_channels = int(n_channels) + self._sr = sample_rate + if device.index is None: + device = torch.device(device.type, torch.accelerator.current_device_index()) + self._device = device + self._max_batch = max(wrapper.batch_sizes) + self._stream = torch.cuda.Stream(device=device, priority=-1) + self._host_codes = torch.zeros((self._n_vq, self._max_batch, self.frames), dtype=torch.long).pin_memory() + self._host_slots = torch.zeros(self._max_batch, dtype=torch.long).pin_memory() + audio_shape = (self._max_batch, self._n_channels, self._samples) + self._host_audio = torch.empty(audio_shape, dtype=torch.float32).pin_memory() + self._copied = torch.cuda.Event() + self._sink: FirstAudioSink | None = None + # Request key -> slot for chunks decoded here and not yet seen by the main path. + self._decoded: dict[str, int] = {} + # A failed output-queue handoff must be replayed from the already + # decoded samples. Re-decoding would advance the streaming state twice. + self._unsent_audio: dict[str, torch.Tensor] = {} + self._handoffs: dict[int, _SlotHandoff] = {} + self._lock = threading.Lock() + self._jobs: queue.SimpleQueue[_Job | None] = queue.SimpleQueue() + self._thread: threading.Thread | None = None + self._closed = False + self._worker_error: Exception | None = None + self._emitted = 0 + self._batches = 0 + self._decode_s = 0.0 + + def bind(self, sink: FirstAudioSink) -> None: + if not callable(getattr(sink, "prepare", None)): + raise TypeError("The codec first-chunk sink requires prepare(request_ids)") + with self._lock: + if self._closed: + raise RuntimeError("The codec first-chunk decoder is closed") + self._sink = sink + if self._thread is None: + self._thread = threading.Thread(target=self._run, name="moss-codec-first-chunk", daemon=True) + self._thread.start() + + def close(self) -> None: + with self._lock: + if not self._closed: + self._closed = True + self._jobs.put(None) + thread = self._thread + if thread is not None: + thread.join(timeout=5.0) + # A stalled decode may still own GPU state: keep its thread and + # handoff, rather than permitting a restart or slot reuse. + if thread.is_alive(): + logger.warning("Codec first-chunk decoder is still stopping; slot waits remain bounded") + + # ------------------------------------------------------------------ + # Receive thread + # ------------------------------------------------------------------ + def submit( + self, + request_id: str, + request_key: str, + codes_flat: Any, + req_slots: dict[str, int], + ) -> bool: + """Claim a first chunk for fast decode; False leaves it to the main path.""" + # Own the codes: the job outlives this call and the payload it came from. + codes = torch.as_tensor(codes_flat, dtype=torch.long).reshape(-1).clone() + if codes.numel() != self._n_vq * self.frames: + return False + with self._lock: + if ( + self._closed + or self._sink is None + or self._thread is None + or not self._thread.is_alive() + or request_key in req_slots + ): + return False + delivery = self._sink.prepare([request_id]) + if request_id not in delivery.routes: + return False + slot = self._session.acquire(oldest=True) + if slot is None: + return False + handoff = _SlotHandoff() + self._handoffs[slot] = handoff + self._decoded[request_key] = slot + req_slots[request_key] = slot + codes = codes.reshape(self._n_vq, self.frames) + # Enqueue under the admission lock: close's sentinel cannot + # overtake an accepted job. + self._jobs.put(_Job(request_id, request_key, delivery, slot, codes, handoff)) + return True + + # ------------------------------------------------------------------ + # Main (engine) thread + # ------------------------------------------------------------------ + def get_request_slot(self, request_key: str, req_slots: dict[str, int]) -> int | None: + """Observe a completed admission before deciding whether cleanup is needed. + + The scheduler removes cancelled routes before the runner's finish hook. + A submit that already froze its route must finish leasing its slot + before that hook checks it. Waiting for decode happens outside this + lock, so the worker can still complete the slot handoff. + """ + with self._lock: + return req_slots.get(request_key) + + def take_decoded(self, request_key: str) -> bool: + """True once per request whose first chunk this path already decoded.""" + with self._lock: + return self._decoded.pop(request_key, None) is not None + + def take_unsent_audio(self, request_key: str) -> torch.Tensor | None: + """Return first-chunk samples when the out-of-band sink rejected them.""" + with self._lock: + return self._unsent_audio.pop(request_key, None) + + def forget(self, request_key: str) -> None: + with self._lock: + self._decoded.pop(request_key, None) + self._unsent_audio.pop(request_key, None) + + def gate(self) -> None: + """Queue the current stream behind the latest fast decode, if still running.""" + event = self._inflight + if self._gate_main and event is not None and not event.query(): + torch.cuda.current_stream(self._device).wait_event(event) + + def order_after(self, slot: int) -> None: + """Make the current stream's next use of ``slot`` follow the fast decode.""" + with self._lock: + handoff = self._handoffs.get(slot) + if handoff is None: + return + if not handoff.done.wait(timeout=self._handoff_timeout_s): + # Keep the handoff. The worker may still be writing the slot, so + # neither cancellation nor a retry may release/re-decode it. + raise TimeoutError(f"Codec first-chunk slot {slot} did not complete within {self._handoff_timeout_s:g}s") + if handoff.error is not None: + raise RuntimeError(f"Codec first-chunk decode failed for slot {slot}") from handoff.error + if handoff.event is not None: + torch.cuda.current_stream(self._device).wait_event(handoff.event) + with self._lock: + self._handoffs.pop(slot, None) + + # ------------------------------------------------------------------ + # Decode thread + # ------------------------------------------------------------------ + def _run(self) -> None: + try: + current_omni_platform.set_device(self._device) + self._run_jobs() + except Exception as error: + logger.exception("MOSS codec first-chunk worker failed") + self._worker_error = error + finally: + with self._lock: + self._closed = True + for handoff in self._handoffs.values(): + if not handoff.done.is_set(): + handoff.error = self._worker_error or RuntimeError("Codec first-chunk worker stopped") + handoff.done.set() + # Accepted, unprocessed routes must not be left waiting for PCM. + while True: + try: + job = self._jobs.get_nowait() + except queue.Empty: + break + if job is not None: + self._fail_delivery(job) + + @staticmethod + def _fail_delivery(job: _Job) -> None: + try: + job.delivery.fail([job.request_id]) + except Exception: + logger.exception("Codec first-chunk error delivery failed for %s", job.request_id) + + def _run_jobs(self) -> None: + while True: + job = self._jobs.get() + if job is None: + return + jobs = [job] + while len(jobs) < self._max_batch: + try: + extra = self._jobs.get_nowait() + except queue.Empty: + break + if extra is None: + self._jobs.put(None) + break + jobs.append(extra) + try: + start = time.perf_counter() + self._decode(jobs) + self._decode_s += time.perf_counter() - start + self._batches += 1 + before, self._emitted = self._emitted, self._emitted + len(jobs) + if self._emitted // 1000 != before // 1000: + logger.info( + "MOSS codec first-chunk fast path: %d chunks in %d batches, %.2f ms/batch", + self._emitted, + self._batches, + 1e3 * self._decode_s / self._batches, + ) + except Exception as error: + logger.exception("MOSS codec first-chunk fast decode failed for %s", [j.request_id for j in jobs]) + # A partial graph replay may have advanced the slot. Surface + # the failure instead of re-decoding and advancing it twice. + for item in jobs: + item.handoff.error = error + item.handoff.done.set() + self._fail_delivery(item) + raise + + @torch.no_grad() + def _decode(self, jobs: list[_Job]) -> None: + n = len(jobs) + for row, item in enumerate(jobs): + self._host_codes[:, row].copy_(item.codes) + self._host_slots[row] = item.slot + stream = self._stream + with torch.cuda.stream(stream): + for item in jobs: + self._session.order_after_reset(item.slot, stream) + codes = self._host_codes[:, :n].to(self._device, non_blocking=True).clamp_(0, self._codebook_size - 1) + slots = self._host_slots[:n].to(self._device, non_blocking=True) + result = self._wrapper.decode(codes, slots) + if result is None: + raise RuntimeError(f"no first-chunk codec graph for B={n}, T={self.frames}") + audio = result[0][:n, ..., : self._samples].float() + state_written = torch.cuda.Event() + state_written.record(stream) + self._inflight = state_written + self._host_audio[:n].view_as(audio).copy_(audio, non_blocking=True) + self._copied.record(stream) + for item in jobs: + item.handoff.event = state_written + self._copied.synchronize() + for row, item in enumerate(jobs): + wav = self._host_audio[row].clone() + if self._n_channels == 1: + wav = wav.reshape(-1) + try: + item.delivery([item.request_id], [wav], self._sr) + except Exception: + logger.exception("MOSS codec first-chunk output failed for %s", item.request_id) + with self._lock: + self._unsent_audio[item.request_key] = wav + finally: + item.handoff.done.set() diff --git a/vllm_omni/model_executor/models/moss_tts/local_model_state.py b/vllm_omni/model_executor/models/moss_tts/local_model_state.py new file mode 100644 index 000000000..c382d4439 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/local_model_state.py @@ -0,0 +1,447 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""MOSS Local request-slot state, independent of each step's batch order. + +State writes, gathers and MTP replay run on the runner stream. Output gathers +own their storage; the runner retains its normal asynchronous output snapshot. +No request dictionary holds a view into a mutable slot or CUDA graph buffer. +""" + +import numpy as np +import torch +from vllm.config import CUDAGraphMode +from vllm.forward_context import set_forward_context +from vllm.logger import init_logger +from vllm.utils.torch_utils import async_tensor_h2d +from vllm.v1.worker.gpu.input_batch import get_num_sampled_and_rejected +from vllm.v1.worker.gpu.sample.output import SamplerOutput +from vllm.v1.worker.gpu.sample.sampler import Sampler + +from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot +from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState + +logger = init_logger(__name__) + + +def _metadata_to_device(data: np.ndarray, device: torch.device) -> torch.Tensor: + # CPU state tests need no CUDA allocator. The serving CUDA path retains + # vLLM's asynchronous pinned upload and its ownership/lifetime handling. + if device.type == "cpu": + return torch.from_numpy(data) + return async_tensor_h2d(data, device=device) + + +class _CodeRowsSnapshot(PackedOutputSnapshot): + """Owned per-request code rows backed by one gathered device tensor. + + The generic snapshot flattens every per-request leaf, concatenates them + into a slab and slices each back, then slices each again after the D2H. + Here the gathered ``[decode_rows, n_vq]`` tensor is already owned, so the + copy is a single D2H and the per-request views come from one ``split``. + """ + + def __init__(self, codes: torch.Tensor, decode_rows: list[int], num_reqs: int) -> None: + dict.__init__(self, {"codes": {"audio": self._rows(codes, decode_rows, num_reqs)}}) + self._codes = codes + self._decode_rows = decode_rows + self._num_reqs = num_reqs + self.producer_event = None + + @staticmethod + def _rows(codes: torch.Tensor, decode_rows: list[int], num_reqs: int) -> list[torch.Tensor]: + outputs = [codes.new_empty((0, codes.shape[-1]))] * num_reqs + for batch_row, row in zip(decode_rows, codes.split(1), strict=True): + outputs[batch_row] = row + return outputs + + def copy_to_cpu(self, copy_tensor): + cpu = copy_tensor(self._codes) + return {"codes": {"audio": self._rows(cpu, self._decode_rows, self._num_reqs)}} + + +class MossLocalModelState(OmniModelState): + def __init__(self, vllm_config, model, encoder_cache, device): + super().__init__(vllm_config, model, encoder_cache, device) + self._init_slot_buffers(self.scheduler_config.max_num_seqs, model.hidden_size, device, self.dtype) + self._batch_prefill = bool(getattr(model.config, "mrv2_batch_prefill", False)) + self._direct_tokens = bool(getattr(model.config, "mrv2_direct_tokens", False)) + self._local_eager_mtp = bool(getattr(model.config, "mrv2_eager_mtp", False)) + logger.info( + "MOSS Local MRV2 batch prefill=%s, determined tokens=%s, eager MTP=%s", + self._batch_prefill, + self._direct_tokens, + self._local_eager_mtp, + ) + logger.info("MOSS Local MRV2 GPU slot state enabled: capacity=%d", self.scheduler_config.max_num_seqs) + + def _init_slot_buffers(self, capacity, hidden_size, device, dtype): + self._hidden_pool = torch.zeros((capacity, hidden_size), dtype=dtype, device=device) + self._codes_pool = torch.full( + (capacity, self.model.n_vq), self.model.audio_pad_token_id, dtype=torch.long, device=device + ) + # Execution eligibility mirrors audio_state.is_stopping. Binary MTP + # continuation is a per-step output, not a sticky request-state update. + self._active_pool = torch.zeros(capacity, dtype=torch.bool, device=device) + self._active_host = [False] * capacity + self._decode_rows = [] + self._prefill_staging = [] + self._batch_prefill = False + self._direct_tokens = False + # Eager MTP: a frame is sampled right after the forward that produced + # its hidden state, one step earlier than the canonical preprocess + # path. Its input-embedding contribution (audio embedding times emit + # mask) is added to the next step's text embedding; ``_keep_pool`` + # chains emission so a stopped stream cannot emit again. + self._local_eager_mtp = False + self._eager_emb_pool = torch.zeros((capacity, hidden_size), dtype=dtype, device=device) + self._keep_pool = torch.zeros(capacity, dtype=torch.bool, device=device) + self._completing_rows = [] + self._local_eager_rows = [] + # A slot takes the eager path only after its final prefill chunk ran + # an eager MTP in this state; others keep the canonical path, so a + # reused slot never reads a previous request's frame state. + self._eager_host = [False] * capacity + self._mtp_dispatcher = None + + def sample_determined_tokens(self, batch, sampler): + """Skip full-vocabulary sampling only for unconstrained Local tokens. + + The audio/binary MTP sampler has already made the random decision. The + outer text distribution has exactly one finite logit. Keep the original + path for processors that can mask it or request distribution metadata. + """ + if not self._direct_tokens or type(sampler) is not Sampler: + return None + rows = batch.idx_mapping_np + if ( + sampler.compute_nans + or sampler.return_sampling_mask + or sampler.trace_replay_state is not None + or sampler.get_logprobs_dims(rows) is not None + or np.any(sampler.logit_bias_state.use_logit_bias[rows]) + or np.any(sampler.penalties_state.use_penalty[rows]) + or np.any(sampler.bad_words_state.num_bad_words.np[rows]) + or ( + sampler.thinking_budget_state.enabled + and np.any(sampler.thinking_budget_state.use_thinking_budget[rows]) + ) + ): + return None + keep = self.model._batch_should_continue + if keep is None or keep.numel() != batch.num_reqs: + return None + # Owned storage: subsequent forwards can replace the continue mask + # while these token IDs are being copied or delivered asynchronously. + tokens = torch.where(keep, self.model.audio_assistant_slot_token_id, self.model.im_end_token_id) + num_sampled, num_rejected = get_num_sampled_and_rejected( + batch.seq_lens.new_ones(batch.num_reqs), + batch.seq_lens, + batch.cu_num_logits, + batch.idx_mapping, + sampler.req_states.prefill_len.gpu, + ) + return SamplerOutput(tokens.reshape(-1, 1), None, None, num_sampled, num_rejected) + + def _apply_reference_batch(self, embeds, references): + """One owned pinned upload for reference codes and their token positions. + + Host storage may be reused only after its H2D event completes. Device + storage is consumed and overwritten in runner-stream order. If all + cached slots are busy, allocate a transient snapshot instead of waiting. + """ + if not references: + return + count = sum(codes.shape[0] for _, codes in references) + capacity = max(count, self._static_inputs_embeds.shape[0]) + staging = None + for item in self._prefill_staging: + if item[0].shape[0] >= count and (item[2] is None or item[2].query()): + staging = item + break + if staging is None: + host = torch.empty((capacity, self.model.n_vq + 1), dtype=torch.long, pin_memory=embeds.is_cuda) + device = torch.empty_like(host, device=embeds.device) if embeds.is_cuda else host + staging = [host, device, torch.cuda.Event() if embeds.is_cuda else None] + if len(self._prefill_staging) < 3: + self._prefill_staging.append(staging) + host, device, event = staging + offset = 0 + for start, codes in references: + end = offset + codes.shape[0] + host[offset:end, 0].numpy()[:] = np.arange(start, start + codes.shape[0]) + host[offset:end, 1:].copy_(codes) + offset = end + if embeds.is_cuda: + device[:count].copy_(host[:count], non_blocking=True) + event.record(torch.cuda.current_stream(embeds.device)) + packed = device[:count] + positions = packed[:, 0] + audio = self.model._audio_embed(packed[:, 1:]) + embeds.index_copy_(0, positions, embeds.index_select(0, positions) + audio) + + def add_request(self, req_index, new_req_data): + super().add_request(req_index, new_req_data) + self._hidden_pool[req_index].zero_() + self._codes_pool[req_index].fill_(self.model.audio_pad_token_id) + self._active_pool[req_index].fill_(True) + self._active_host[req_index] = True + self._eager_host[req_index] = False + self._keep_pool[req_index].fill_(False) + self._eager_emb_pool[req_index].zero_() + + def remove_request(self, req_index_or_id): + # Freeing is CPU bookkeeping. Reinitialization happens on admission, + # in stream order, so pending output snapshots remain independent. + super().remove_request(req_index_or_id) + + def _select_rows(self, tensor, rows): + if rows == list(range(len(rows))): + return tensor[: len(rows)] + indices = _metadata_to_device(np.asarray(rows, dtype=np.int64), tensor.device) + return tensor.index_select(0, indices) + + def run_preprocess(self, input_batch, model_inputs, req_states=None, mtp_batch_descriptor_dispatcher=None): + input_ids = model_inputs.get("input_ids") + if input_ids is None: + input_ids = input_batch.input_ids + embeds = model_inputs.get("inputs_embeds") + if embeds is None: + embeds = self.model.embed_input_ids(input_ids[: input_batch.num_tokens]) + model_inputs["inputs_embeds"] = embeds + elif self._static_inputs_embeds is not None and embeds.data_ptr() == self._static_inputs_embeds.data_ptr(): + embeds[: input_batch.num_tokens].copy_(self.model.embed_input_ids(input_ids[: input_batch.num_tokens])) + + decode_rows = [] + completing_rows = [] + references = [] + for row in range(input_batch.num_reqs): + slot = int(input_batch.idx_mapping_np[row]) + buf = self.intermediate_buffer.buffers[slot] + if not buf or str(buf.get("req_id", "")).startswith("_warmup_"): + continue + start = int(input_batch.query_start_loc_np[row]) + count = int(input_batch.num_scheduled_tokens[row]) + prompt_len = ( + None if req_states is None else self._get_req_state_value(getattr(req_states, "prompt_len", None), slot) + ) + computed = ( + None + if req_states is None + else self._get_req_state_value(getattr(req_states, "num_computed_tokens", None), slot) + ) + if computed is None: + computed = self._get_input_batch_num_computed(input_batch, slot, row) + prefill = computed < prompt_len if computed is not None and prompt_len is not None else count > 1 + if prefill and computed is not None and prompt_len is not None and computed + count >= prompt_len: + completing_rows.append(row) + if count == 1 and not prefill and isinstance(buf.get("audio_state"), dict): + active = not bool(buf["audio_state"].get("is_stopping")) + if active != self._active_host[slot]: + self._active_pool[slot].fill_(active) + self._active_host[slot] = active + decode_rows.append(row) + continue + if self._batch_prefill: + ref = (buf.get("codes", {}) or {}).get("ref") + ref_offset = int(buf.get("ref_offset", 0)) + # Device-side references keep the canonical path; do not add + # an implicit D2H just to place them in a host staging buffer. + if not isinstance(ref, torch.Tensor) or ref.device.type == "cpu": + if isinstance(ref, torch.Tensor) and ref.numel(): + if ref.dim() == 1 and ref.numel() % self.model.n_vq == 0: + ref = ref.view(-1, self.model.n_vq) + if ref.dim() == 2: + chunk = ref[ref_offset : ref_offset + count] + if chunk.numel() and chunk.shape[0] == count: + references.append((start, chunk)) + # Text embeddings have already been computed for all input + # tokens. Match the canonical prefill state transition. + self.intermediate_buffer.update( + slot, + { + "audio_state": {"is_stopping": False}, + "ref_offset": ref_offset + count, + }, + self.model.gpu_resident_buffer_keys, + ) + if not self._active_host[slot]: + self._active_pool[slot].fill_(True) + self._active_host[slot] = True + continue + # Reference conditioning and chunked-prefill coordinates retain + # the model's canonical scalar implementation. + info = dict(buf) + info["_omni_is_prefill"] = prefill + if prompt_len is not None: + info["_omni_prompt_len"] = prompt_len + if computed is not None: + info["_omni_num_computed_tokens"] = computed + ids = input_ids[start : start + count] + previous = embeds[start : start + count] + new_ids, new_embeds, updates = self.model.preprocess(ids, previous, **info) + if "mtp_inputs" in updates: + raise RuntimeError("MOSS slot state requires an admitted prefill before decode") + if self._preprocess_result_needs_writeback(ids, new_ids): + ids.copy_(new_ids) + if self._preprocess_result_needs_writeback(previous, new_embeds): + previous.copy_(new_embeds) + self.intermediate_buffer.update(slot, updates, self.model.gpu_resident_buffer_keys) + active = not bool(updates.get("audio_state", {}).get("is_stopping", False)) + if active != self._active_host[slot]: + self._active_pool[slot].fill_(active) + self._active_host[slot] = active + + self._apply_reference_batch(embeds, references) + self._decode_rows = decode_rows + if self._local_eager_mtp: + self._completing_rows = completing_rows + self._mtp_dispatcher = mtp_batch_descriptor_dispatcher + eager_rows = [row for row in decode_rows if self._eager_host[int(input_batch.idx_mapping_np[row])]] + canonical_rows = [row for row in decode_rows if not self._eager_host[int(input_batch.idx_mapping_np[row])]] + self._local_eager_rows = eager_rows + if eager_rows: + # Their frame was sampled after the previous forward. + slots = self._select_rows(input_batch.idx_mapping, eager_rows) + offsets = self._mtp_offsets[: len(eager_rows)] + offsets.copy_(self._select_rows(input_batch.query_start_loc, eager_rows)) + embeds.index_copy_( + 0, offsets, embeds.index_select(0, offsets) + self._eager_emb_pool.index_select(0, slots) + ) + if canonical_rows: + self._run_slot_mtp(input_batch, input_ids, embeds, canonical_rows, mtp_batch_descriptor_dispatcher) + elif decode_rows: + self._run_slot_mtp(input_batch, input_ids, embeds, decode_rows, mtp_batch_descriptor_dispatcher) + + def _run_slot_mtp(self, batch, input_ids, embeds, rows, dispatcher): + bsz = len(rows) + slots = self._select_rows(batch.idx_mapping, rows) + # Upstream query_start_loc is int32; index_copy_ requires int64. + offsets = self._mtp_offsets[:bsz] + offsets.copy_(self._select_rows(batch.query_start_loc, rows)) + ids, emb = self._mtp_input_ids[:bsz], self._mtp_input_embeds[:bsz] + hidden, step = self._mtp_hidden[:bsz], self._mtp_text_step[:bsz] + ids.copy_(input_ids.index_select(0, offsets)) # MRV2 token inputs are int32. + torch.index_select(embeds, 0, offsets, out=emb) + torch.index_select(self._hidden_pool, 0, slots, out=hidden) + step[:, 0].copy_(self._active_pool.index_select(0, slots)) + + buffers = [self.intermediate_buffer.buffers[int(batch.idx_mapping_np[row])] for row in rows] + req_ids = [str(buf["req_id"]) for buf in buffers] + generators = [ + self._get_mtp_generator(req_id, buf.get("sampling_params"), ids.device) + for req_id, buf in zip(req_ids, buffers, strict=True) + ] + can_graph = self._is_mtp_graph_runner() and not any(g is not None for g in generators) + desc = dispatcher(bsz) if can_graph and dispatcher is not None else None + mode = getattr(desc, "cg_mode", CUDAGraphMode.FULL) if can_graph else CUDAGraphMode.NONE + num_tokens = int(desc.num_tokens) if desc is not None else bsz + with set_forward_context( + None, self.vllm_config, num_tokens=num_tokens, cudagraph_runtime_mode=mode, batch_descriptor=desc + ): + new_emb, codes = self._call_mtp_with_sampling( + self._mtp_input_ids[:num_tokens], + self._mtp_input_embeds[:num_tokens], + self._mtp_hidden[:num_tokens], + self._mtp_text_step[:num_tokens], + buffers=buffers, + req_ids=req_ids, + generators=generators, + ) + embeds.index_copy_(0, offsets, new_emb[:bsz]) + self._codes_pool.index_copy_(0, slots, codes[:bsz]) + + def _run_local_eager_mtp(self, batch, rows, completing): + """Sample the next frame from this forward's hidden state. + + Same MTP inputs as the canonical next-step preprocess (hidden state + and active flag); the input embedding is zero, so the returned + embedding is exactly the audio contribution the next step adds to its + text embedding. A decode row emits only if its previous frame did. + """ + bsz = len(rows) + slots = self._select_rows(batch.idx_mapping, rows) + ids, emb = self._mtp_input_ids[:bsz], self._mtp_input_embeds[:bsz] + hidden, step = self._mtp_hidden[:bsz], self._mtp_text_step[:bsz] + ids.zero_() # the MTP does not read token ids + emb.zero_() + torch.index_select(self._hidden_pool, 0, slots, out=hidden) + gate = self._keep_pool.index_select(0, slots) + if completing: + first = _metadata_to_device( + np.asarray([row in completing for row in rows], dtype=np.bool_), + gate.device, + ) + gate = gate | first + step[:, 0].copy_(self._active_pool.index_select(0, slots) & gate) + + buffers = [self.intermediate_buffer.buffers[int(batch.idx_mapping_np[row])] for row in rows] + req_ids = [str(buf["req_id"]) for buf in buffers] + generators = [ + self._get_mtp_generator(req_id, buf.get("sampling_params"), ids.device) + for req_id, buf in zip(req_ids, buffers, strict=True) + ] + dispatcher = self._mtp_dispatcher + can_graph = self._is_mtp_graph_runner() and not any(g is not None for g in generators) + desc = dispatcher(bsz) if can_graph and dispatcher is not None else None + mode = getattr(desc, "cg_mode", CUDAGraphMode.FULL) if can_graph else CUDAGraphMode.NONE + num_tokens = int(desc.num_tokens) if desc is not None else bsz + with set_forward_context( + None, self.vllm_config, num_tokens=num_tokens, cudagraph_runtime_mode=mode, batch_descriptor=desc + ): + contribution, codes = self._call_mtp_with_sampling( + self._mtp_input_ids[:num_tokens], + self._mtp_input_embeds[:num_tokens], + self._mtp_hidden[:num_tokens], + self._mtp_text_step[:num_tokens], + buffers=buffers, + req_ids=req_ids, + generators=generators, + ) + codes = codes[:bsz] + self._eager_emb_pool.index_copy_(0, slots, contribution[:bsz]) + self._codes_pool.index_copy_(0, slots, codes) + self._keep_pool.index_copy_(0, slots, codes.ne(self.model.audio_pad_token_id).any(dim=-1)) + + def run_postprocess(self, hidden_states, input_batch): + # query_start_loc is already device-resident and authoritative for + # mixed/chunked-prefill batches. index_copy snapshots graph outputs. + last = input_batch.query_start_loc[1 : input_batch.num_reqs + 1] - 1 + self._hidden_pool.index_copy_( + 0, input_batch.idx_mapping[: input_batch.num_reqs], hidden_states.index_select(0, last) + ) + if self._local_eager_mtp: + completing = self._completing_rows + rows = sorted(self._local_eager_rows + completing) + # Canonical decode rows emit their preprocess frame; eager rows + # and completed prefills emit the frame sampled below. + self._decode_rows = sorted(set(self._decode_rows) | set(completing)) + if rows: + self._run_local_eager_mtp(input_batch, rows, set(completing)) + for row in completing: + self._eager_host[int(input_batch.idx_mapping_np[row])] = True + + def postprocess_model_output(self, model_output, input_batch, req_states): + # Local backbone returns a tensor; keep the standard runner's output + # snapshot/materializer contract and omit hidden payload transport. + if not isinstance(model_output, torch.Tensor): + raise TypeError("MOSS Local slot state requires a tensor backbone output") + model = self.model + model._batch_state = None + model._batch_state_spans = None + model._batch_should_continue = torch.ones(input_batch.num_reqs, device=model_output.device, dtype=torch.bool) + if not self._decode_rows: + return model_output, {} + # This gather is an owned snapshot. Later slot writes cannot change + # data already handed to the runner, even before its D2H completes. + slots = self._select_rows(input_batch.idx_mapping, self._decode_rows) + codes = self._codes_pool.index_select(0, slots) + keep = codes.ne(model.audio_pad_token_id).any(dim=-1) + if self._decode_rows == list(range(input_batch.num_reqs)): + model._batch_should_continue = keep + else: + rows = _metadata_to_device( + np.asarray(self._decode_rows, dtype=np.int64), + codes.device, + ) + model._batch_should_continue.index_copy_(0, rows, keep) + return model_output, _CodeRowsSnapshot(codes, list(self._decode_rows), input_batch.num_reqs) diff --git a/vllm_omni/model_executor/models/moss_tts/local_short_attention.py b/vllm_omni/model_executor/models/moss_tts/local_short_attention.py new file mode 100644 index 000000000..522d74c76 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/local_short_attention.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Single-query attention for the frame-local MOSS depth transformer. + +The caller has already written this token's K/V and provides only its valid +prefix. There is no persistent state here. BF16 probability rounding matches +the FlashAttention P/V multiply, but reduction order is not bit-identical. +""" + +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def _short_attention( + q_ptr, + k_ptr, + v_ptr, + out_ptr, + qs0: tl.constexpr, + qs1: tl.constexpr, + ks0: tl.constexpr, + ks1: tl.constexpr, + ks2: tl.constexpr, + vs0: tl.constexpr, + vs1: tl.constexpr, + vs2: tl.constexpr, + heads: tl.constexpr, + dim: tl.constexpr, + tokens: tl.constexpr, + block_d: tl.constexpr, + block_t: tl.constexpr, +): + row = tl.program_id(0) + batch, head = row // heads, row % heads + d, t = tl.arange(0, block_d), tl.arange(0, block_t) + query = tl.load(q_ptr + batch * qs0 + head * qs1 + d, d < dim, 0).to(tl.float32) + keys = tl.load( + k_ptr + batch * ks0 + head * ks1 + t[:, None] * ks2 + d[None, :], + (t[:, None] < tokens) & (d[None, :] < dim), + 0, + ).to(tl.float32) + logits = tl.sum(keys * query[None, :], axis=1) * (dim**-0.5 * 1.4426950408889634) + logits = tl.where(t < tokens, logits, float("-inf")) + probs = tl.exp2(logits - tl.max(logits, axis=0)) + denominator = tl.sum(probs, axis=0) + values = tl.load( + v_ptr + batch * vs0 + head * vs1 + t[:, None] * vs2 + d[None, :], + (t[:, None] < tokens) & (d[None, :] < dim), + 0, + ) + probs = probs.to(values.dtype).to(tl.float32) + output = tl.sum(values.to(tl.float32) * probs[:, None], axis=0) / denominator + tl.store(out_ptr + row * dim + d, output, d < dim) + + +@torch.library.custom_op("vllm_omni::moss_local_short_attention", mutates_args=()) +def local_short_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + batch, heads, seq_len, dim = query.shape + if not ( + query.is_cuda + and key.device == query.device + and value.device == query.device + and query.dtype == key.dtype == value.dtype == torch.bfloat16 + and seq_len == 1 + and 0 < dim <= 128 + and key.shape == value.shape + and key.shape[:2] == query.shape[:2] + and key.shape[-1] == dim + and 0 < key.shape[2] <= 16 + and query.stride(-1) == key.stride(-1) == value.stride(-1) == 1 + ): + raise ValueError("MOSS short attention requires CUDA BF16 single-query inputs and 1..16 valid keys") + output = query.new_empty((batch, heads, 1, dim)) + _short_attention[(batch * heads,)]( + query, + key, + value, + output, + query.stride(0), + query.stride(1), + key.stride(0), + key.stride(1), + key.stride(2), + value.stride(0), + value.stride(1), + value.stride(2), + heads, + dim, + key.shape[2], + triton.next_power_of_2(dim), + triton.next_power_of_2(key.shape[2]), + num_warps=1, + ) + return output + + +@local_short_attention.register_fake +def _fake_local_short_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + return query.new_empty(query.shape) diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py index d87f01b4b..a526d8e33 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py @@ -8,7 +8,10 @@ from __future__ import annotations +import threading from collections.abc import Iterable +from contextlib import nullcontext +from functools import partial from typing import Any import torch @@ -32,11 +35,16 @@ from vllm_omni.model_executor.models.moss_tts.configuration_moss_audio_tokenizer from vllm_omni.model_executor.models.moss_tts.cuda_graph_streaming_decoder_wrapper import ( CUDAGraphStreamingDecoderWrapper, ) +from vllm_omni.model_executor.models.moss_tts.first_chunk_fast_path import ( + MossFirstChunkFastPath, +) from vllm_omni.model_executor.models.moss_tts.moss_codec_cudagraph import ( MossTTSCUDAGraphCodecWrapper, ) from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot, pack_output_snapshot from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp +from vllm_omni.worker_v2.first_audio_sender import FirstAudioSink logger = init_logger(__name__) @@ -53,18 +61,46 @@ class _MossCodecStreamSession: vllm_config: VllmConfig, graph_batch_sizes: list[int] | None = None, graph_frame_sizes: list[int] | None = None, + gpu_output: bool = False, + chunk_frames: int = 0, + ring_headroom: bool = False, + fast_batch_sizes: list[int] | None = None, + fast_frames: int = 0, ) -> None: self._codec = codec self._state_capacity = int(state_capacity) + self._chunk_frames = int(chunk_frames) + self._ring_headroom = bool(ring_headroom) self._n_vq = int(n_vq) self._device = next(codec.parameters()).device + self._gpu_output = gpu_output and self._device.type == "cuda" self._free_stream_slots = list(reversed(range(self._state_capacity))) self._leased_slots: set[int] = set() + # The first-chunk fast path leases slots from its decode thread. + self._slot_lock = threading.Lock() + self._slot_reset_events: dict[int, torch.cuda.Event] = {} self._closed = False + self._metadata_index = 0 + self._metadata_ring = [] + if self._gpu_output: + # The metadata DMA is tiny, but a short ring forces the producer + # CPU thread to synchronize with the codec stream while it waits + # for a host slab to become reusable. Keep enough entries for a + # scheduler burst so staging remains enqueue-only in practice. + for _ in range(64): + host = torch.empty(2 * self._state_capacity, dtype=torch.long, pin_memory=True) + gpu = torch.empty_like(host, device=self._device) + self._metadata_ring.append((host, host.numpy(), gpu, torch.cuda.Event(), False)) self._cudagraph_wrapper: CUDAGraphStreamingDecoderWrapper | None = None + self._fast_wrapper: CUDAGraphStreamingDecoderWrapper | None = None batch_sizes = sorted({int(size) for size in (graph_batch_sizes or []) if 0 < int(size) <= self._state_capacity}) frame_sizes = sorted({int(size) for size in (graph_frame_sizes or []) if int(size) > 0}) scratch_capacity = max(batch_sizes, default=0) if self._device.type in ("cuda", "npu") else 0 + fast_batch_sizes = sorted({int(size) for size in (fast_batch_sizes or []) if 0 < int(size)}) + if not (fast_frames > 0 and fast_batch_sizes and batch_sizes and frame_sizes and self._device.type == "cuda"): + fast_batch_sizes = [] + fast_scratch_base = self._state_capacity + scratch_capacity + scratch_capacity += max(fast_batch_sizes, default=0) self._total_state_capacity = self._state_capacity + scratch_capacity self._state_slot_ids = torch.arange( self._total_state_capacity, @@ -76,7 +112,14 @@ class _MossCodecStreamSession: if not callable(initialize_state_pool): raise RuntimeError("The streaming codec does not implement a decoder state pool.") with torch.no_grad(): - initialize_state_pool(self._state_capacity, scratch_capacity) + if self._ring_headroom: + if self._chunk_frames <= 0: + raise ValueError("codec_ring_headroom requires a positive codec chunk frame count.") + # Size every attention ring for context + one full chunk so a + # chunk write never evicts keys inside its own causal window. + initialize_state_pool(self._state_capacity, scratch_capacity, chunk_frames=self._chunk_frames) + else: + initialize_state_pool(self._state_capacity, scratch_capacity) if batch_sizes and frame_sizes and self._device.type in ("cuda", "npu"): if self._device.type == "npu": from vllm_omni.platforms.npu.models.moss_tts_streaming_decode_wrapper import ( @@ -101,16 +144,34 @@ class _MossCodecStreamSession: vllm_config=vllm_config, ) self._cudagraph_wrapper.warmup(self._device) + if fast_batch_sizes and self._cudagraph_wrapper.is_ready: + main_compiled = self._cudagraph_wrapper._compiled_decode + self._fast_wrapper = CUDAGraphStreamingDecoderWrapper( + codec, + state_capacity=self._state_capacity, + batch_sizes=fast_batch_sizes, + frame_sizes=[int(fast_frames)], + num_quantizers=self._n_vq, + vllm_config=vllm_config, + scratch_base=fast_scratch_base, + private_pool=True, + compiled_decode=main_compiled if main_compiled is not None else codec.decode_streaming_tensors, + ) + self._fast_wrapper.warmup(self._device) + if not self._fast_wrapper.is_ready: + self._fast_wrapper = None self.reset_slots(list(range(self._state_capacity + scratch_capacity))) if not self._cudagraph_wrapper.is_ready: self._cudagraph_wrapper = None - def acquire(self) -> int | None: - if not self._free_stream_slots: - return None - slot = self._free_stream_slots.pop() - self._leased_slots.add(slot) - return slot + def acquire(self, *, oldest: bool = False) -> int | None: + """Lease a slot; ``oldest`` takes the one whose reset was queued first.""" + with self._slot_lock: + if not self._free_stream_slots: + return None + slot = self._free_stream_slots.pop(0 if oldest else -1) + self._leased_slots.add(slot) + return slot def release(self, slot: int, *, state_already_reset: bool = False) -> None: if self._closed: @@ -119,8 +180,23 @@ class _MossCodecStreamSession: return if not state_already_reset: self.reset_slots([slot]) - self._leased_slots.remove(slot) - self._free_stream_slots.append(slot) + reset_event = None + if self._fast_wrapper is not None: + # The fast path decodes on another stream; it must not read this + # slot before the reset queued on the current stream has run. + reset_event = torch.cuda.Event() + reset_event.record(torch.cuda.current_stream(self._device)) + with self._slot_lock: + if reset_event is not None: + self._slot_reset_events[slot] = reset_event + self._leased_slots.remove(slot) + self._free_stream_slots.append(slot) + + def order_after_reset(self, slot: int, stream: torch.cuda.Stream) -> None: + with self._slot_lock: + event = self._slot_reset_events.pop(slot, None) + if event is not None and not event.query(): + stream.wait_event(event) def _reset_slot_ids(self, slot_ids: torch.Tensor) -> None: reset_streaming_slots = getattr(self._codec, "reset_decoder_state_slots", None) @@ -139,9 +215,29 @@ class _MossCodecStreamSession: raise ValueError(f"Decoder state slots must be contiguous, got {slots}") self._reset_slot_ids(self._state_slot_ids[start : start + len(slots)]) + def _stage_slot_ids(self, slots: list[int], terminal_slots: set[int]) -> tuple[torch.Tensor, torch.Tensor]: + # Host slabs cannot be overwritten until their DMA completes. Device + # slabs are reused in producer-stream order, after all previous reads. + index = self._metadata_index + host, view, gpu, event, recorded = self._metadata_ring[index] + if recorded and not event.query(): + event.synchronize() + terminal_ids = [slot for slot in slots if slot in terminal_slots] + n, end = len(slots), len(slots) + len(terminal_ids) + view[:n] = slots + view[n:end] = terminal_ids + gpu[:end].copy_(host[:end], non_blocking=True) + event.record(torch.cuda.current_stream(self._device)) + self._metadata_ring[index] = (host, view, gpu, event, True) + self._metadata_index = (index + 1) % len(self._metadata_ring) + return gpu[:n], gpu[n:end] + def close(self) -> None: if self._closed: return + for _, _, _, event, recorded in self._metadata_ring: + if recorded: + event.synchronize() close_state_pool = getattr(self._codec, "close_decoder_state_pool", None) with torch.no_grad(): if callable(close_state_pool): @@ -154,6 +250,7 @@ class _MossCodecStreamSession: slot_codes: dict[int, torch.Tensor], *, terminal_slots: set[int] | None = None, + output_buffers: dict[int, torch.Tensor] | None = None, ) -> dict[int, torch.Tensor]: if not slot_codes: return {} @@ -178,17 +275,22 @@ class _MossCodecStreamSession: ], dim=1, ) - state_slot_ids = torch.tensor(slots, device=self._device, dtype=torch.long) + staged_terminal_ids = None + if self._metadata_ring: + state_slot_ids, staged_terminal_ids = self._stage_slot_ids(slots, terminal_slots) + else: + state_slot_ids = torch.tensor(slots, device=self._device, dtype=torch.long) graph_output: tuple[torch.Tensor, torch.Tensor, int] | None = None if self._cudagraph_wrapper is not None: - # A terminal tail can safely pad T to a larger graph: decoder - # attention is causal, audio_lengths trims the padding, and every - # affected state slot is reset immediately after this step. + # An all-terminal group may borrow the next larger frame graph. + # The decoder is causal, so cropping the padded tail on the time + # axis recovers the exact-T waveform up to attention-ring + # eviction (see codec_ring_headroom); the state is reset below. graph_output = self._cudagraph_wrapper.decode( codes_step, state_slot_ids, - allow_frame_padding=len(terminal_slots) == len(slots), + allow_frame_padding=bool(terminal_slots) and len(terminal_slots) == len(slots), ) used_cudagraph = graph_output is not None @@ -214,11 +316,69 @@ class _MossCodecStreamSession: audio_tensor = result.audio if terminal_slots: - terminal_rows = [row for row, slot in enumerate(slots) if slot in terminal_slots] - terminal_slot_ids = state_slot_ids if len(terminal_rows) == len(slots) else state_slot_ids[terminal_rows] + if staged_terminal_ids is not None: + terminal_slot_ids = staged_terminal_ids + else: + terminal_rows = [row for row, slot in enumerate(slots) if slot in terminal_slots] + terminal_slot_ids = ( + state_slot_ids if len(terminal_rows) == len(slots) else state_slot_ids[terminal_rows] + ) self._reset_slot_ids(terminal_slot_ids) - audio = audio_tensor.detach().to("cpu", torch.float32) + if self._gpu_output and output_buffers is not None: + # Write directly into the caller-owned per-forward slab. This + # preserves graph output before the next replay without creating + # a temporary waveform tensor that would be packed again later. + out: dict[int, torch.Tensor] = {} + row = 0 + while row < len(slots): + destination = output_buffers.get(slots[row]) + if destination is None: + row += 1 + continue + end = row + 1 + if end < len(slots) and destination.is_contiguous(): + storage = destination.untyped_storage().data_ptr() + offset = destination.storage_offset() + row_elements = destination.numel() + while end < len(slots): + following = output_buffers.get(slots[end]) + if not ( + following is not None + and following.shape == destination.shape + and following.dtype == destination.dtype + and following.device == destination.device + and following.is_contiguous() + and following.untyped_storage().data_ptr() == storage + and following.storage_offset() == offset + (end - row) * row_elements + ): + break + end += 1 + if end > row + 1: + # A mixed-frame forward can leave gaps between equal-sized + # rows in its slab. Copy each adjacent run as one batch, + # without overwriting the intervening requests' output. + destination = destination.as_strided( + (end - row, *destination.shape), (destination.numel(), *destination.stride()) + ) + source = audio_tensor[row:end] + else: + source = audio_tensor[row] + # Crop time before reshaping: flattening first mixes stereo + # channels whenever a row is shorter than the decode batch. + destination.copy_(source[..., : destination.shape[-1]].view_as(destination)) + out.update((slot, output_buffers[slot]) for slot in slots[row:end]) + row = end + return out + + if self._gpu_output: + # A graph's static output is overwritten on every replay, including + # multiple replays within one forward. Snapshot on the producer + # stream before returning; record_stream alone cannot prevent this + # overwrite. MRV2 owns the subsequent pinned, asynchronous D2H. + audio = audio_tensor.detach().to(dtype=torch.float32, copy=True) + else: + audio = audio_tensor.detach().to("cpu", torch.float32) out: dict[int, torch.Tensor] = {} for row, slot in enumerate(slots): audio_length = int(slot_codes[slot].shape[1]) * self._samples_per_frame @@ -270,6 +430,7 @@ class MossTTSCodecDecoder(nn.Module): requires_exact_input_shape: bool = True _OUTPUT_SAMPLE_RATE: int = 24_000 + _first_chunk_fast_path: MossFirstChunkFastPath | None = None def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() @@ -292,20 +453,38 @@ class MossTTSCodecDecoder(nn.Module): self._n_channels: int = 1 self._sr_tensor = torch.tensor(self._OUTPUT_SAMPLE_RATE, dtype=torch.int32) self._stream_session: _MossCodecStreamSession | None = None + self._codec_stream: torch.cuda.Stream | None = None scheduler_cfg = getattr(self.vllm_config, "scheduler_config", None) self._stream_state_capacity = max(1, int(getattr(scheduler_cfg, "max_num_seqs", 1) or 1)) self._initial_stream_chunk_frames: int = self._connector_int("initial_codec_chunk_frames", default=0) self._stream_chunk_frames: int = self._connector_int("codec_chunk_frames", default=0) self._stream_max_step_frames: int = self._stream_chunk_frames or 100 + # Opt-in: size codec attention rings as context + one chunk so chunked + # streaming reproduces whole-sequence decoding (costs ring memory). + self._stream_ring_headroom: bool = bool(self._connector_int("codec_ring_headroom", default=0)) self._stream_req_slots: dict[str, int] = {} + # Opt-in: decode each stream's first chunk from the receive thread. + self._first_chunk_fast = bool(self._connector_int("codec_first_chunk_fast_path", default=0)) + self._first_chunk_fast_path: MossFirstChunkFastPath | None = None self._async_chunk = bool(getattr(self.vllm_config.model_config, "async_chunk", False)) + self._gpu_stream_output = ( + self._async_chunk + and bool(getattr(vllm_config.model_config, "use_v2_model_runner", False)) + and getattr(cfg, "model_type", None) == "moss_tts_local" + # Diagnostic switch: 0 copies each decode group's batch to the host synchronously. + and bool(self._connector_int("codec_gpu_stream_output", default=1)) + ) self._streaming_graph_batch_sizes = self._streaming_graph_batch_sizes_from_compilation_config() - # codec_chunk_ramp ladder sizes join the streaming graph frame-size - # set so every ladder (B, T) combination is captured during warmup. + ramp = parse_chunk_ramp(self._connector_extra(), steady=self._stream_chunk_frames) + if ramp: + # The sender gives the ramp precedence over initial chunk size. + # The dedicated first-chunk graph must use that same first shape. + self._initial_stream_chunk_frames = ramp[0] + self._stream_max_step_frames = max( + self._stream_max_step_frames, self._initial_stream_chunk_frames, *(ramp or []) + ) self._streaming_graph_frame_sizes = _resolve_streaming_graph_frame_sizes( - self._initial_stream_chunk_frames, - self._stream_chunk_frames, - self._connector_extra_cfg(), + self._initial_stream_chunk_frames, self._stream_chunk_frames, self._connector_extra() ) # ------------------------------------------------------------------ @@ -337,6 +516,41 @@ class MossTTSCodecDecoder(nn.Module): inputs_embeds: torch.Tensor | None = None, runtime_additional_information: list[dict[str, Any]] | None = None, **kwargs: Any, + ) -> OmniOutput: + if self._gpu_stream_output and input_ids is not None and input_ids.is_cuda: + if self._codec_stream is None: + self._codec_stream = torch.cuda.Stream(device=input_ids.device) + main_stream = torch.cuda.current_stream(input_ids.device) + # Runner input storage is reused by the next batch on main_stream. + # Own it before allowing that preparation to overlap codec decode. + owned_input_ids = input_ids.clone() + with torch.cuda.stream(self._codec_stream): + self._codec_stream.wait_stream(main_stream) + owned_input_ids.record_stream(self._codec_stream) + output = self._forward( + owned_input_ids, + positions, + intermediate_tensors, + inputs_embeds, + runtime_additional_information, + **kwargs, + ) + if isinstance(output.multimodal_outputs, PackedOutputSnapshot): + output.multimodal_outputs.record_producer_event(self._codec_stream) + return output + return self._forward( + input_ids, positions, intermediate_tensors, inputs_embeds, runtime_additional_information, **kwargs + ) + + @torch.no_grad() + def _forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + intermediate_tensors: Any = None, + inputs_embeds: torch.Tensor | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, + **kwargs: Any, ) -> OmniOutput: """Decode audio VQ codes to waveform. @@ -493,13 +707,37 @@ class MossTTSCodecDecoder(nn.Module): audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav if streaming_work: - for i, wav in self._decode_streaming_batch(streaming_work).items(): + stream_output_slab = None + stream_output_buffers: dict[int, torch.Tensor] = {} + if self._gpu_stream_output and device.type == "cuda": + total_elements = 0 + for output_index, _, codes_nq_t, _ in streaming_work: + samples = int(codes_nq_t.shape[1]) * self._codec.downsample_rate + total_elements += samples * (self._n_channels if self._n_channels > 1 else 1) + if total_elements: + stream_output_slab = torch.empty(total_elements, dtype=torch.float32, device=device) + offset = 0 + for output_index, _, codes_nq_t, _ in streaming_work: + samples = int(codes_nq_t.shape[1]) * self._codec.downsample_rate + elements = samples * (self._n_channels if self._n_channels > 1 else 1) + shape = (self._n_channels, samples) if self._n_channels > 1 else (samples,) + stream_output_buffers[output_index] = stream_output_slab.narrow(0, offset, elements).view(shape) + offset += elements + for i, wav in self._decode_streaming_batch( + streaming_work, + output_buffers=stream_output_buffers or None, + ).items(): audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav - return OmniOutput( - text_hidden_states=None, - multimodal_outputs={"model_outputs": audios, "sr": srs}, - ) + payload = {"model_outputs": audios, "sr": srs} + if self._gpu_stream_output and device.type == "cuda": + # One transfer per dtype/device preserves ragged request lengths + # without issuing one D2H per waveform. A fresh slot owns each + # output until the runner completes its copy; no reusable ring. + packed = pack_output_snapshot(payload, {}, max_buckets=0, reuse_existing_storage=True) + if packed is not None: + payload = packed + return OmniOutput(text_hidden_states=None, multimodal_outputs=payload) @staticmethod def _normalize_seq_token_counts(value: Any) -> list[int] | None: @@ -546,25 +784,104 @@ class MossTTSCodecDecoder(nn.Module): vllm_config=self.vllm_config, graph_batch_sizes=self._streaming_graph_batch_sizes, graph_frame_sizes=self._streaming_graph_frame_sizes, + gpu_output=self._gpu_stream_output, + chunk_frames=self._stream_max_step_frames, + ring_headroom=self._stream_ring_headroom, + fast_batch_sizes=[1, 2, 4, 8] if self._first_chunk_fast else None, + fast_frames=self._initial_stream_chunk_frames if self._first_chunk_fast else 0, ) + if self._stream_session._fast_wrapper is not None: + self._first_chunk_fast_path = MossFirstChunkFastPath( + self._stream_session, + self._stream_session._fast_wrapper, + n_vq=self._n_vq, + frames=self._initial_stream_chunk_frames, + codebook_size=int(self._codec.config.codebook_size), + samples_per_frame=int(self._codec.downsample_rate), + n_channels=self._n_channels, + sample_rate=self._sr_tensor, + device=next(self._codec.parameters()).device, + gate_main=bool(self._connector_int("codec_first_chunk_gate", default=0)), + ) + logger.info( + "MOSS-TTS codec first-chunk fast path ready: B=%s T=%d", + self._stream_session._fast_wrapper.capture_sizes, + self._initial_stream_chunk_frames, + ) + elif self._first_chunk_fast: + logger.warning("MOSS-TTS codec first-chunk fast path requested but unavailable; using the main path.") + if self._stream_ring_headroom: + logger.info( + "MOSS-TTS codec attention rings sized with headroom for %d-frame chunks", + self._stream_max_step_frames, + ) return self._stream_session + def bind_first_chunk_fast_path(self, sink: FirstAudioSink) -> Any: + """Start the fast path with an output sink; returns the receive hook.""" + fast = self._first_chunk_fast_path + if fast is None: + return None + fast.bind(sink) + return self._submit_first_chunk + + def _submit_first_chunk(self, request_id: str, handle: Any, payload: Any) -> bool: + fast = self._first_chunk_fast_path + if fast is None or not isinstance(payload, dict): + return False + codes = payload.get("codes") + codes = codes.get("audio") if isinstance(codes, dict) else None + if codes is None: + return False + meta = payload.get("meta") if isinstance(payload.get("meta"), dict) else {} + key = meta.get("req_id") + if isinstance(key, (list, tuple)): + key = key[0] if key else None + if key is None: + key = getattr(handle, "external_req_id", None) or request_id + return fast.submit(request_id, str(key), codes, self._stream_req_slots) + def _decode_streaming_batch( self, items: list[tuple[int, str, torch.Tensor, bool]], + *, + output_buffers: dict[int, torch.Tensor] | None = None, ) -> dict[int, torch.Tensor]: session = self._ensure_stream_session() if session is None: return {} outputs: dict[int, torch.Tensor] = {} + unsent_first_audio: dict[int, torch.Tensor] = {} grouped: dict[int, list[tuple[int, str, int, torch.Tensor, bool]]] = {} max_step_frames = max(1, int(self._stream_max_step_frames)) graph = session._cudagraph_wrapper + # Terminal tails join the regular padded frame bucket when every + # execution batch width has a graph, so a step stays one replay + # instead of one eager or plain-graph decode per distinct tail length. coalesce_tails = graph is not None and max(graph.batch_sizes, default=0) >= self._stream_state_capacity + fast = self._first_chunk_fast_path for output_index, request_id, codes_nq_t, finished in items: slot = self._stream_req_slots.get(request_id) + if fast is not None and slot is not None: + fast.order_after(slot) + if fast.take_decoded(request_id): + unsent = fast.take_unsent_audio(request_id) + if unsent is not None: + unsent_first_audio[output_index] = unsent + # The fast path already emitted this chunk's audio and + # advanced the slot state past it. + codes_nq_t = codes_nq_t[:, fast.frames :] + if codes_nq_t.shape[1] == 0: + if unsent is not None: + outputs[output_index] = unsent + if finished: + self._finish_stream_request(request_id, session, slot) + continue + if output_buffers is not None: + # Sized for the untrimmed chunk; decode into a fresh tensor. + output_buffers.pop(output_index, None) if slot is None: slot = session.acquire() if slot is None: @@ -580,9 +897,11 @@ class MossTTSCodecDecoder(nn.Module): slot, codes_nq_t, terminal=finished, + output_buffer=output_buffers.get(output_index) if output_buffers is not None else None, ) if wav is not None: - outputs[output_index] = wav + prefix = unsent_first_audio.pop(output_index, None) + outputs[output_index] = torch.cat((prefix, wav), dim=-1) if prefix is not None else wav if finished: self._finish_stream_request( request_id, @@ -593,21 +912,33 @@ class MossTTSCodecDecoder(nn.Module): continue frame_count = int(codes_nq_t.shape[1]) - # Only combine tails that already used the same padded graph. - # Preserve exact-size graphs (including T=1) and the eager path. group_frames = frame_count if coalesce_tails and finished: + # Only terminal rows may be padded: their state is reset after + # the step, and the output is cropped to the exact frame count. group_frames = graph._select_frame_size(frame_count, allow_padding=True) or frame_count grouped.setdefault(group_frames, []).append((output_index, request_id, slot, codes_nq_t, finished)) + if fast is not None and grouped: + fast.gate() for group in grouped.values(): plan = {slot: codes_nq_t for _, _, slot, codes_nq_t, _ in group} terminal_slots = {slot for _, _, slot, _, finished in group if finished} - decoded = session.step(plan, terminal_slots=terminal_slots) + step_buffers = ( + { + slot: output_buffers[output_index] + for output_index, _, slot, _, _ in group + if output_buffers is not None and output_index in output_buffers + } + if output_buffers is not None + else None + ) + decoded = session.step(plan, terminal_slots=terminal_slots, output_buffers=step_buffers) for output_index, request_id, slot, _, finished in group: wav = decoded.get(slot) if wav is not None: - outputs[output_index] = wav + prefix = unsent_first_audio.pop(output_index, None) + outputs[output_index] = torch.cat((prefix, wav), dim=-1) if prefix is not None else wav if finished: self._finish_stream_request( request_id, @@ -625,6 +956,7 @@ class MossTTSCodecDecoder(nn.Module): codes_nq_t: torch.Tensor, *, terminal: bool = False, + output_buffer: torch.Tensor | None = None, ) -> torch.Tensor | None: if codes_nq_t.numel() == 0: return None @@ -633,13 +965,24 @@ class MossTTSCodecDecoder(nn.Module): for start in range(0, int(codes_nq_t.shape[1]), max_step_frames): chunk = codes_nq_t[:, start : start + max_step_frames] is_terminal_chunk = terminal and start + max_step_frames >= int(codes_nq_t.shape[1]) + chunk_samples = int(chunk.shape[1]) * self._codec.downsample_rate + chunk_buffer = ( + output_buffer[ + ..., start * self._codec.downsample_rate : start * self._codec.downsample_rate + chunk_samples + ] + if output_buffer is not None + else None + ) decoded = session.step( {slot: chunk}, terminal_slots={slot} if is_terminal_chunk else None, + output_buffers={slot: chunk_buffer} if chunk_buffer is not None else None, ) wav = decoded.get(slot) if wav is not None: parts.append(wav) + if output_buffer is not None: + return output_buffer if not parts: return None return torch.cat(parts, dim=-1) if len(parts) > 1 else parts[0] @@ -665,29 +1008,30 @@ class MossTTSCodecDecoder(nn.Module): terminal payload, so the runner calls this hook from its finished-request path to avoid leaking stream slots and buffered codes. """ - session = self._stream_session - for req_id in finished_req_ids: - request_id = str(req_id) - slot = self._stream_req_slots.get(request_id) - if slot is None: - continue - if session is not None: - self._finish_stream_request(request_id, session, slot) - else: - self._stream_req_slots.pop(request_id, None) + # Cancellation and terminal cleanup share the decoder's stream order; + # a reused slot must not be reset ahead of an outstanding decode. + context = torch.cuda.stream(self._codec_stream) if self._codec_stream is not None else nullcontext() + with context: + session = self._stream_session + for req_id in finished_req_ids: + request_id = str(req_id) + fast = self._first_chunk_fast_path + slot = ( + fast.get_request_slot(request_id, self._stream_req_slots) + if fast is not None + else self._stream_req_slots.get(request_id) + ) + if slot is None: + continue + if fast is not None: + fast.order_after(slot) + fast.forget(request_id) + if session is not None: + self._finish_stream_request(request_id, session, slot) + else: + self._stream_req_slots.pop(request_id, None) - def _connector_int(self, name: str, default: int = 0) -> int: - model_cfg = getattr(self.vllm_config, "model_config", None) - connector_cfg = getattr(model_cfg, "stage_connector_config", None) - if isinstance(connector_cfg, dict): - extra_cfg: dict | None = connector_cfg.get("extra", connector_cfg) - else: - extra_cfg = getattr(connector_cfg, "extra", None) - if isinstance(extra_cfg, dict) and name in extra_cfg: - return int(extra_cfg[name]) - return default - - def _connector_extra_cfg(self) -> dict: + def _connector_extra(self) -> dict[str, Any]: model_cfg = getattr(self.vllm_config, "model_config", None) connector_cfg = getattr(model_cfg, "stage_connector_config", None) if isinstance(connector_cfg, dict): @@ -696,6 +1040,9 @@ class MossTTSCodecDecoder(nn.Module): extra_cfg = getattr(connector_cfg, "extra", None) return extra_cfg if isinstance(extra_cfg, dict) else {} + def _connector_int(self, name: str, default: int = 0) -> int: + return int(self._connector_extra().get(name, default)) + def _streaming_graph_batch_sizes_from_compilation_config(self) -> list[int]: if getattr(self.vllm_config.model_config, "enforce_eager", True): return [] @@ -816,15 +1163,33 @@ class MossTTSCodecDecoder(nn.Module): codec.decoder.to(dtype=torch.bfloat16) attention_backend = getattr(self.vllm_config.model_config.hf_config, "codec_attention_backend", "sdpa") if attention_backend != "sdpa": - if attention_backend != "triton" or device.type != "cuda": + if attention_backend not in {"triton", "triton_slot"} or device.type != "cuda": raise ValueError(f"Unsupported codec attention backend/device: {attention_backend}/{device.type}") from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import MossAudioTokenizerMultiheadAttention from vllm_omni.model_executor.models.moss_tts.streaming_attention import masked_attention + if attention_backend == "triton_slot": + from vllm_omni.model_executor.models.moss_tts.slot_attention import ( + slot_ring_attention, + slot_ring_attention_rows, + ) + + skip_empty_tiles = bool(self._connector_int("codec_skip_empty_attention_tiles", default=0)) + fused_slots = bool(self._connector_int("codec_fused_slot_attention", default=0)) + if skip_empty_tiles and fused_slots: + raise ValueError("codec_skip_empty_attention_tiles requires unfused slot attention") + if skip_empty_tiles: + slot_ring_attention = partial(slot_ring_attention, skip_empty_tiles=True) + logger.info("MOSS codec slot attention: skipping empty tiles for T <= 32") + for module in codec.decoder.modules(): if isinstance(module, MossAudioTokenizerMultiheadAttention): module._streaming_attention = masked_attention - logger.info("Enabled Triton masked attention for the streaming codec decoder") + if attention_backend == "triton_slot": + module._slot_attention = slot_ring_attention + if self._connector_int("codec_fused_slot_attention", default=0): + module._slot_attention_rows = slot_ring_attention_rows + logger.info("Enabled codec attention backend=%s", attention_backend) build_decode_lut = getattr(codec.quantizer, "build_decode_lut", None) if callable(build_decode_lut): lut_dtype = torch.bfloat16 if device.type == "cuda" else torch.float32 diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py index c3f226fd1..ae837eddb 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py @@ -21,7 +21,8 @@ checkpoint 1:1 so ``load_weights()`` needs no remapping. from __future__ import annotations -from collections.abc import Sequence +import os +from collections.abc import Callable, Sequence import torch import torch.nn as nn @@ -52,6 +53,7 @@ class _MossTTSLocalAttention(nn.Module): self.n_head = n_head self.head_dim = hidden_size // n_head self.embed_dim = hidden_size + self._short_attention = os.environ.get("VLLM_OMNI_MOSS_LOCAL_SHORT_ATTN", "0") == "1" self.c_attn = nn.Linear(hidden_size, 3 * hidden_size, bias=True) self.c_proj = nn.Linear(hidden_size, hidden_size, bias=True) inv_freq = 1.0 / (rope_base ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32) / self.head_dim)) @@ -134,7 +136,19 @@ class _MossTTSLocalAttention(nn.Module): v_cache[:, :, position : position + 1].copy_(value) key = k_cache[:, :, : position + 1] value = v_cache[:, :, : position + 1] - attn_output = F.scaled_dot_product_attention(query, key, value, is_causal=kv_cache is None) + if ( + self._short_attention + and kv_cache is not None + and query.is_cuda + and query.dtype == torch.bfloat16 + and self.head_dim <= 128 + and key.shape[2] <= 16 + ): + from vllm_omni.model_executor.models.moss_tts.local_short_attention import local_short_attention + + attn_output = local_short_attention(query, key, value) + else: + attn_output = F.scaled_dot_product_attention(query, key, value, is_causal=kv_cache is None) attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, self.embed_dim) return self.c_proj(attn_output) @@ -193,6 +207,7 @@ class MossTTSLocalDepthTransformer(nn.Module): self.h = nn.ModuleList([_MossTTSLocalBlock(self.hidden_size, n_head, inner_size, rope_base, eps)]) self.ln_f = nn.LayerNorm(self.hidden_size, eps=eps) self._compiled_forward_prefix = None + self._compiled_audio_sampler = None def _forward_prefix( self, @@ -216,6 +231,14 @@ class MossTTSLocalDepthTransformer(nn.Module): dynamic=True, options={"epilogue_fusion": False}, ) + if os.environ.get("VLLM_OMNI_MOSS_LOCAL_COMPILE_AUDIO_SAMPLER", "0") == "1": + # Keep the existing top-k/top-p algorithm and torch RNG. Explicit + # per-request generators use the original helper below. The + # binary continue/stop head is deliberately unchanged. + self._compiled_audio_sampler = torch.compile( + _sample_token, fullgraph=True, dynamic=True, options={"fallback_random": True} + ) + logger.info("MOSS-TTS Local compiled audio-channel sampler enabled") logger.info("MOSS-TTS local depth frame-local KV execution enabled with torch.compile") def _run_prefix( @@ -239,6 +262,7 @@ class MossTTSLocalDepthTransformer(nn.Module): do_sample: bool, generator: torch.Generator | None, generators: Sequence[torch.Generator | None] | None = None, + sample_token: Callable[..., torch.Tensor] | None = None, ) -> torch.Tensor: """Compute channel logits, apply repetition penalty, sample, store.""" channel_logits = audio_lm_heads[channel_index](local_hidden).float() @@ -250,7 +274,7 @@ class MossTTSLocalDepthTransformer(nn.Module): pos = sel > 0 sel = torch.where(pos, sel / repetition_penalty, sel * repetition_penalty) channel_logits.index_copy_(-1, hist_t, sel) - channel_token = _sample_token( + channel_token = (sample_token or _sample_token)( channel_logits, temperature, top_k, @@ -351,6 +375,7 @@ class MossTTSLocalDepthTransformer(nn.Module): do_sample=do_sample, generator=generator, generators=generators, + sample_token=self._compiled_audio_sampler if generator is None and generators is None else None, ) if channel_index + 1 < n_vq: 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 bc4a5b716..303724e2e 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 @@ -37,6 +37,7 @@ from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth impo ) from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform +from vllm_omni.worker.sampling_utils import get_tts_local_seed logger = init_logger(__name__) @@ -1675,6 +1676,31 @@ class MossTTSLocalTalkerForGeneration(nn.Module): return input_embeds_out, output_codes # ------------------------------------------------------------------ + # MRV2 reads model-owned capabilities; retain the V1 hook and platform + # graph-safety setting as the canonical implementation. + mtp = talker_mtp + get_mtp_seed = staticmethod(get_tts_local_seed) + + def create_mrv2_model_state(self, vllm_config, encoder_cache, device): + from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState + + state_cls = OmniModelState + if getattr(self.config, "mrv2_gpu_slot_state", False): + from .local_model_state import MossLocalModelState + + state_cls = MossLocalModelState + return state_cls(vllm_config, self, encoder_cache, device) + + @property + def mtp_output_key(self) -> tuple[str, str]: + """Return the request-buffer destination for one local audio frame.""" + return self.talker_mtp_output_key + + @property + def mtp_graph_safe(self) -> bool: + """Honor the same platform graph-safety override as the V1 runner.""" + return self.talker_mtp_graph_safe + # Package runner-generated audio frames # ------------------------------------------------------------------ diff --git a/vllm_omni/model_executor/models/moss_tts/pipeline.py b/vllm_omni/model_executor/models/moss_tts/pipeline.py index 0d6151313..766c0d136 100644 --- a/vllm_omni/model_executor/models/moss_tts/pipeline.py +++ b/vllm_omni/model_executor/models/moss_tts/pipeline.py @@ -97,6 +97,7 @@ MOSS_TTS_LOCAL_PIPELINE = PipelineConfig( StagePipelineConfig( stage_id=0, model_stage="moss_tts_local", + supports_native_mrv2_data_plane=True, execution_type=StageExecutionType.LLM_AR, input_sources=(), owns_tokenizer=True, @@ -114,6 +115,7 @@ MOSS_TTS_LOCAL_PIPELINE = PipelineConfig( StagePipelineConfig( stage_id=1, model_stage="moss_tts_local_codec", + supports_native_mrv2_data_plane=True, execution_type=StageExecutionType.LLM_GENERATION, input_sources=(0,), final_output=True, 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 b1516d90a..e492cf4cd 100644 --- a/vllm_omni/model_executor/models/moss_tts/reference_encoder.py +++ b/vllm_omni/model_executor/models/moss_tts/reference_encoder.py @@ -25,7 +25,10 @@ so importing it from the API-server process does not pull the talker/codec. from __future__ import annotations import asyncio +import contextlib import hashlib +import os +from collections import OrderedDict from collections.abc import Awaitable, Callable from functools import lru_cache from typing import Any @@ -35,22 +38,93 @@ import torch from numpy.typing import NDArray from vllm.logger import init_logger +from vllm_omni.model_executor.models.moss_tts.reference_encoder_graphs import ( + FINE_BATCH_SIZES, + FINE_BUCKET_SECONDS, + MossReferenceEncoderGraphs, + install_windowed_attention, +) +from vllm_omni.model_executor.models.moss_tts.shared_reference_encoder import ( + SharedReferenceEncoderClient, + SharedReferenceEncoderHost, + SharedReferenceEncoderStartupError, + elect_host, +) + logger = init_logger(__name__) # Coalescing defaults. Kept as module constants (not env/CLI) to match the # ``_REF_AUDIO_RESOLVE_CACHE_MAX_*`` convention in serving_speech.py. _REF_ENCODE_BATCH_WINDOW_MS = 10.0 -_REF_ENCODE_MAX_BATCH = 8 +_REF_ENCODE_MAX_BATCH = 32 +# Experiment switches for cold-reference latency (see the A/B notes): +# batch window in ms; batches in flight per API process with a shared encoder; +# "high" puts reference encodes on the highest-priority CUDA stream. +_REF_ENCODE_WINDOW_ENV = "VLLM_OMNI_MOSS_REF_BATCH_WINDOW_MS" +_REF_ENCODE_INFLIGHT_ENV = "VLLM_OMNI_MOSS_REF_INFLIGHT" +_REF_ENCODE_PRIORITY_ENV = "VLLM_OMNI_MOSS_REF_STREAM_PRIORITY" +# "1" resamples and loudness-normalizes clips on the encoder's GPU instead of +# the API process's CPU (not bit-identical to the CPU path). +_REF_ENCODE_GPU_PREP_ENV = "VLLM_OMNI_MOSS_REF_GPU_PREP" _INT32_MAX = 2**31 +# Inline (data: URI) references already resolved and encoded, keyed by the +# serving resolve key. Codes are small; the URI itself is not retained. +_INLINE_INDEX_MAX_ENTRIES = 4096 +# Optional directory (e.g. under /dev/shm) that API processes of one server +# share, so an inline reference encoded by any frontend is reused by all. +_SHARED_CODES_DIR_ENV = "VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR" +# "0" keeps one reference encoder per API process even with a shared directory. +_SHARED_ENCODER_ENV = "VLLM_OMNI_MOSS_REF_SHARED_ENCODER" +_shared_role: str | None = None +# Concurrent encode workers in the shared host, each with its own CUDA graphs. +_SHARED_WORKERS_ENV = "VLLM_OMNI_MOSS_REF_ENCODER_WORKERS" +_SHARED_WORKER_BATCH_SECONDS = 96.0 + + +def shared_encoder_role() -> str: + """``"host"``, ``"client"`` or ``""`` (no shared encoder) for this process. + + Decided once, before the processor places its tokenizer on the GPU, so + client processes can leave it on the CPU. + """ + global _shared_role + if _shared_role is None: + shared_dir = os.environ.get(_SHARED_CODES_DIR_ENV) + if not shared_dir or os.environ.get(_SHARED_ENCODER_ENV, "1") == "0": + _shared_role = "" + else: + os.makedirs(shared_dir, exist_ok=True) + _shared_role = "host" if elect_host(shared_dir) else "client" + return _shared_role + + +# Opt in to reference graphs; existing deployments keep the eager encoder. +_REF_ENCODE_GRAPHS_ENV = "VLLM_OMNI_MOSS_REF_GRAPHS" + + +def reference_graphs_enabled() -> bool: + return os.environ.get(_REF_ENCODE_GRAPHS_ENV, "0") == "1" + + +def reference_stream_priority() -> int: + """CUDA stream priority for reference encodes (lower is more urgent).""" + if os.environ.get(_REF_ENCODE_PRIORITY_ENV, "") == "high" and torch.cuda.is_available(): + return torch.cuda.Stream.priority_range()[1] + return 0 + + +# "sdpa" keeps the tokenizer's dense masked SDPA attention. By default the +# reference encoder attends through a local-window flash kernel, the backend +# the tokenizer config selects (flash_attention_2) when flash-attn is present. +_REF_ENCODE_ATTN_ENV = "VLLM_OMNI_MOSS_REF_ATTN" +# "0" captures the encoder's eager kernels. By default the captured encoder is +# compiled first, which fuses its many small elementwise kernels. +_REF_ENCODE_COMPILE_ENV = "VLLM_OMNI_MOSS_REF_COMPILE" _Waveform = NDArray[np.float32] | list[float] -def _sha1(s: str) -> str: - return hashlib.sha1((s or "").encode("utf-8")).hexdigest() - - @lru_cache(maxsize=16) def _reference_resampler(sr: int, sr_target: int) -> torch.nn.Module: import torchaudio @@ -72,6 +146,13 @@ def _prep_wav_sync(waveform: _Waveform, sr: int, sr_target: int) -> torch.Tensor return wav +def _inline_reference_key(ref_str: str) -> str | None: + """SHA-1 of a ``data:`` URI, matching the serving resolve key for it.""" + if not isinstance(ref_str, str) or ref_str[:5].lower() != "data:": + return None + return hashlib.sha1(ref_str.encode("utf-8")).hexdigest() + + def _to_compact_codes(codes: torch.Tensor) -> torch.Tensor: """Downcast RVQ codes to int32 for compact caching.""" codes = codes.detach().cpu().contiguous() @@ -119,18 +200,26 @@ class _RefEncodeBatcher: *, window_ms: float, max_batch: int, + max_inflight: int = 1, ): # encode_batch_fn: sync, takes [(waveform, sr), ...], returns a list of # (codes_tensor | Exception) aligned to the input order. self._encode_batch_fn = encode_batch_fn self._window_s = max(0.0, float(window_ms) / 1000.0) self._max_batch = max(1, int(max_batch)) + # Batches encoded at once. Clips that arrive while every slot is busy + # are queued and form the next batch. + self._max_inflight = max(1, int(max_inflight)) self._queue: asyncio.Queue | None = None + self._slots: asyncio.Semaphore | None = None + self._inflight: set[asyncio.Task] = set() self._drainer: asyncio.Task | None = None def _ensure_started(self) -> None: if self._queue is None: self._queue = asyncio.Queue() + if self._slots is None: + self._slots = asyncio.Semaphore(self._max_inflight) if self._drainer is None or self._drainer.done(): self._drainer = asyncio.create_task(self._drain_loop()) @@ -141,16 +230,27 @@ class _RefEncodeBatcher: return await fut async def _drain_loop(self) -> None: - assert self._queue is not None + assert self._queue is not None and self._slots is not None while True: - first = await self._queue.get() - jobs = None + await self._slots.acquire() try: - jobs = await self._coalesce(first) - await self._run_batch(jobs) - finally: - # Do not retain the completed batch while waiting for new work. - del first, jobs + jobs = await self._coalesce(await self._queue.get()) + except BaseException: + self._slots.release() + raise + task = asyncio.create_task(self._run_and_release(jobs)) + self._inflight.add(task) + task.add_done_callback(self._inflight.discard) + # Do not retain the batch here while waiting for new work. + del jobs + + async def _run_and_release(self, jobs: list[tuple[_Waveform, int, asyncio.Future]]) -> None: + try: + await self._run_batch(jobs) + finally: + del jobs + assert self._slots is not None + self._slots.release() async def _coalesce(self, first: tuple) -> list[tuple]: """Group ``first`` with jobs arriving within the batch window.""" @@ -200,6 +300,8 @@ class _RefEncodeBatcher: except (asyncio.CancelledError, Exception): # noqa: BLE001 pass self._drainer = None + for task in list(self._inflight): + task.cancel() class MossReferenceEncoder: @@ -213,7 +315,7 @@ class MossReferenceEncoder: n_vq: int, sr_target: int, speaker_cache: Any, - window_ms: float = _REF_ENCODE_BATCH_WINDOW_MS, + window_ms: float | None = None, max_batch: int = _REF_ENCODE_MAX_BATCH, ): self._processor = processor @@ -224,8 +326,33 @@ class MossReferenceEncoder: # model_type namespaces the whole family so a moss_tts server never # collides with another model's speaker-cache entries. self._model_type = f"moss_tts_{variant}_nq{int(n_vq)}" - self._inflight: dict[str, asyncio.Task] = {} - self._batcher = _RefEncodeBatcher(self._encode_batch_sync, window_ms=window_ms, max_batch=max_batch) + self._inflight: dict[str | tuple[str, str], asyncio.Task] = {} + self._inline_index: OrderedDict[str, torch.Tensor] = OrderedDict() + self._shared_codes_dir = os.environ.get(_SHARED_CODES_DIR_ENV) or None + if self._shared_codes_dir: + os.makedirs(self._shared_codes_dir, exist_ok=True) + # Realtime encodes with its own codec and never shares the encoder. + self._shares_encoder = bool(self._shared_codes_dir) and type(self)._encode_prepared is ( + MossReferenceEncoder._encode_prepared + ) + self._shared_host: SharedReferenceEncoderHost | None = None + self._shared_client: SharedReferenceEncoderClient | None = None + self._graphs: MossReferenceEncoderGraphs | None = None + self._attention_installed = False + self._graphs_enabled = reference_graphs_enabled() + if window_ms is None: + window_ms = float(os.environ.get(_REF_ENCODE_WINDOW_ENV, _REF_ENCODE_BATCH_WINDOW_MS)) + # A shared encoder serves several batches at once (one per host + # worker); a process's own graphs replay one batch at a time. + inflight = 1 + if self._shares_encoder and os.environ.get(_SHARED_ENCODER_ENV, "1") != "0": + inflight = max(1, int(os.environ.get(_REF_ENCODE_INFLIGHT_ENV, "1"))) + self._local_stream: torch.cuda.Stream | None = None + self._gpu_prep = os.environ.get(_REF_ENCODE_GPU_PREP_ENV, "0") == "1" + self._gpu_resamplers: dict[torch.device, torch.nn.Module] = {} + self._batcher = _RefEncodeBatcher( + self._encode_batch_sync, window_ms=window_ms, max_batch=max_batch, max_inflight=inflight + ) def _make_cache_key(self, name: str, created_at: int) -> tuple: return self._speaker_cache.make_cache_key(name, model_type=self._model_type, created_at=int(created_at)) @@ -251,6 +378,7 @@ class MossReferenceEncoder: caller salts those requests with ``voice_created_at`` instead). """ voice_name, created_at = _registered_voice(voice_name, voice_created_at) + inline_key: str | None = None if voice_name: # A named voice has a stable key that does not depend on the @@ -260,6 +388,16 @@ class MossReferenceEncoder: if cached is not None: return _clone_out(cached["codes"]), None else: + inline_key = _inline_reference_key(ref_str) + if inline_key is not None: + codes = self._inline_index.get(inline_key) + if codes is None: + codes = self._load_shared_codes(inline_key) + if codes is not None: + # A data: URI is immutable content, and its serving + # resolve key is this same digest: skip the resolve hop. + self._inline_index.move_to_end(inline_key) + return _clone_out(codes), inline_key # Anonymous refs have no pre-resolve hot path: the content key must # come from the resolve itself (mtime/size from a single stat) so an # on-disk edit invalidates the cached codes; the flight body @@ -267,17 +405,55 @@ class MossReferenceEncoder: # when the resolve cache is warm. The flight key is the request-side # reference (not the content hash), so concurrent requests for the # same ref_str also share the resolve/download, not just the encode. - flight_key = "ref:" + _sha1(ref_str) + # Key by the locator itself: equality is exact, and hashing a + # large inline data URI once per request object is much cheaper + # than a separate SHA-1 of its UTF-8 copy on the event loop. + flight_key = ("ref", ref_str) codes, resolve_key = await self._single_flight( flight_key, lambda: self._resolve_and_encode(ref_str, resolve_ref_audio, get_artifact_key, voice_name, created_at), ) + if not voice_name and inline_key is not None and resolve_key == inline_key: + # Only index when the resolver's key is exactly this digest. + self._remember_inline(inline_key, codes) + self._store_shared_codes(inline_key, codes) return _clone_out(codes), resolve_key + def _remember_inline(self, key: str, codes: torch.Tensor) -> None: + self._inline_index[key] = codes + self._inline_index.move_to_end(key) + while len(self._inline_index) > _INLINE_INDEX_MAX_ENTRIES: + self._inline_index.popitem(last=False) + + def _load_shared_codes(self, key: str) -> torch.Tensor | None: + if not self._shared_codes_dir: + return None + try: + array = np.load(os.path.join(self._shared_codes_dir, key + ".npy"), allow_pickle=False) + except (OSError, ValueError): + return None + if array.ndim != 2 or array.shape[1] != self._n_vq or array.dtype not in (np.int32, np.int64): + return None + codes = torch.from_numpy(array) + self._remember_inline(key, codes) + return codes + + def _store_shared_codes(self, key: str, codes: torch.Tensor) -> None: + if not self._shared_codes_dir: + return + path = os.path.join(self._shared_codes_dir, key + ".npy") + tmp = f"{path}.{os.getpid()}.tmp" + try: + with open(tmp, "wb") as handle: + np.save(handle, codes.detach().cpu().numpy(), allow_pickle=False) + os.replace(tmp, path) + except OSError: + logger.warning("MOSS ref codes: could not write shared entry %s", path, exc_info=True) + async def _single_flight( self, - flight_key: str, + flight_key: str | tuple[str, str], start_flight: Callable[[], Awaitable[tuple[torch.Tensor, str]]], ) -> tuple[torch.Tensor, str]: """Join the in-flight encode for ``flight_key``, starting one if absent. @@ -300,7 +476,7 @@ class MossReferenceEncoder: # a duplicate resolve/encode instead of joining this one. Identity # guard: only drop the slot if it still holds *our* task (a later # request may have replaced it after ours completed). - def _retire(t: asyncio.Task, key: str = flight_key) -> None: + def _retire(t: asyncio.Task, key: str | tuple[str, str] = flight_key) -> None: if self._inflight.get(key) is t: self._inflight.pop(key, None) @@ -350,29 +526,228 @@ class MossReferenceEncoder: prepared_idx: list[int] = [] for i, (waveform, sr) in enumerate(payload): try: - prepared.append(_prep_wav_sync(waveform, sr, self._sr_target)) + prepared.append(self._encoder_input(_prep_wav_sync(waveform, sr, self._sr_target))) prepared_idx.append(i) except Exception as exc: # noqa: BLE001 — isolate this clip's failure results[i] = exc if not prepared: return results - try: - codes_list = self._encode_prepared(prepared) - for local_i, orig_i in enumerate(prepared_idx): - results[orig_i] = codes_list[local_i] - except Exception: # noqa: BLE001 — batch forward failed; retry item-by-item - logger.warning("MOSS ref batch encode (n=%d) failed; falling back to per-item", len(prepared)) - for local_i, orig_i in enumerate(prepared_idx): - try: - results[orig_i] = self._encode_prepared([prepared[local_i]])[0] - except Exception as exc: # noqa: BLE001 — isolate this clip's failure - results[orig_i] = exc + for orig_i, codes in zip(prepared_idx, self._encode_shared_or_local(prepared)): + results[orig_i] = codes return results - def _encode_prepared(self, prepared: list[torch.Tensor]) -> list[torch.Tensor]: - with torch.no_grad(): - return self._processor.encode_audios_from_wav(prepared, sampling_rate=self._sr_target, n_vq=self._n_vq) + def prepare(self) -> None: + """Capture the CUDA graphs (or start the shared encoder) before the first request.""" + role = shared_encoder_role() if self._shares_encoder else "" + if role == "host": + self._host().wait_until_ready() + elif role == "client": + if self._shared_client is None: + self._shared_client = SharedReferenceEncoderClient(self._shared_codes_dir) + self._shared_client.wait_until_ready() + elif role == "": + self._reference_graphs() + + def _host(self) -> SharedReferenceEncoderHost: + if self._shared_host is None: + # Graph capture runs in the background; clips queue until it ends. + self._shared_host = SharedReferenceEncoderHost( + self._shared_codes_dir, + self._host_workers(), + stream_priority=reference_stream_priority(), + batch_window_ms=float(os.environ.get("VLLM_OMNI_MOSS_REF_HOST_WINDOW_MS", "0")), + ) + return self._shared_host + + def _encode_shared_or_local(self, prepared: list[torch.Tensor]) -> list: + role = shared_encoder_role() if self._shares_encoder else "" + if role == "host": + return self._host().encode(prepared) + if role == "client": + if self._shared_client is None: + self._shared_client = SharedReferenceEncoderClient(self._shared_codes_dir) + try: + return self._shared_client.encode(prepared) + except SharedReferenceEncoderStartupError: + raise + except Exception: # noqa: BLE001 — the host is unreachable: encode here + logger.warning("MOSS shared reference encoder unavailable; encoding locally", exc_info=True) + self._place_tokenizer_for_local_encode() + with self._local_stream_context(): + return self._encode_local(prepared) + + def _local_stream_context(self): + tokenizer = getattr(self._processor, "audio_tokenizer", None) + params = next(tokenizer.parameters(), None) if isinstance(tokenizer, torch.nn.Module) else None + if params is None or params.device.type != "cuda": + return contextlib.nullcontext() + if self._local_stream is None: + self._local_stream = torch.cuda.Stream(device=params.device, priority=reference_stream_priority()) + return torch.cuda.stream(self._local_stream) + + def _host_workers(self) -> list: + """Encode workers of the shared host, each with its own graphs and stream.""" + count = max(1, int(os.environ.get(_SHARED_WORKERS_ENV, "4"))) + if not self._graph_capable(): + return [(self._encode_local, None)] * count + workers = [] + for _ in range(count): + graphs = MossReferenceEncoderGraphs( + self._processor.audio_tokenizer, + n_vq=self._n_vq, + batch_sizes=FINE_BATCH_SIZES, + bucket_seconds=FINE_BUCKET_SECONDS, + max_batch_seconds=_SHARED_WORKER_BATCH_SECONDS, + compile_core=self._compile_graphs(), + batched_transfer=os.environ.get("VLLM_OMNI_MOSS_REF_BATCHED_TRANSFER", "0") == "1", + singleton_bucket_seconds=self._singleton_bucket_seconds(), + ) + workers.append(((lambda wavs, g=graphs: self._encode_local(wavs, g)), graphs.capture)) + return workers + + def _encode_local(self, prepared: list[torch.Tensor], graphs: MossReferenceEncoderGraphs | None = None) -> list: + """One batched forward, retried per clip so one bad clip fails alone.""" + try: + return list(self._encode_prepared(prepared, graphs)) + except Exception: # noqa: BLE001 — batch forward failed; retry item-by-item + logger.warning("MOSS ref batch encode (n=%d) failed; falling back to per-item", len(prepared)) + results: list = [] + for wav in prepared: + try: + results.append(self._encode_prepared([wav], graphs)[0]) + except Exception as exc: # noqa: BLE001 — isolate this clip's failure + results.append(exc) + return results + + def _place_tokenizer_for_local_encode(self) -> None: + """A client process keeps the tokenizer off the GPU until it must encode.""" + tokenizer = getattr(self._processor, "audio_tokenizer", None) + device = getattr(self._processor, "_vllm_omni_ref_encoder_device", None) + if tokenizer is not None and device is not None and next(tokenizer.parameters()).device != device: + self._processor.audio_tokenizer = tokenizer.to(device) + + def _encode_prepared( + self, prepared: list[torch.Tensor], graphs: MossReferenceEncoderGraphs | None = None + ) -> list[torch.Tensor]: + # ``prepared`` clips come from _encoder_input, at ``_clip_rate()``. + graphs = graphs or self._reference_graphs() + if graphs is None: + with torch.no_grad(): + return self._processor.encode_audios_from_wav( + prepared, sampling_rate=self._clip_rate(), n_vq=self._n_vq + ) + codes = graphs.encode([self._finish_prepare(wav) for wav in prepared]) + missing = [i for i, c in enumerate(codes) if c is None] + if missing: + with torch.no_grad(): + eager = self._processor.encode_audios_from_wav( + [prepared[i] for i in missing], sampling_rate=self._clip_rate(), n_vq=self._n_vq + ) + for i, c in zip(missing, eager): + codes[i] = c + return codes + + def _graph_capable(self) -> bool: + """True when the tokenizer can run graphs; installs windowed attention once.""" + if not self._graphs_enabled: + return False + tokenizer = getattr(self._processor, "audio_tokenizer", None) + device = next(tokenizer.parameters()).device if tokenizer is not None else None + if device is None or device.type != "cuda" or not hasattr(tokenizer, "_codec_inference_autocast"): + self._graphs_enabled = False + return False + if not self._attention_installed and os.environ.get(_REF_ENCODE_ATTN_ENV, "flash") == "flash": + logger.info( + "MOSS reference encoder: %d attention layers use local-window flash attention", + install_windowed_attention( + tokenizer, + fa_version=int(os.environ.get("VLLM_OMNI_MOSS_REF_FA_VERSION", "2")), + cache_rope=os.environ.get("VLLM_OMNI_MOSS_REF_CACHE_ROPE", "0") == "1", + skip_padded_query_mask=os.environ.get("VLLM_OMNI_MOSS_REF_SKIP_PADDED_QUERY_MASK", "0") == "1", + ), + ) + self._attention_installed = True + return True + + @staticmethod + def _compile_graphs() -> bool: + return os.environ.get(_REF_ENCODE_COMPILE_ENV, "1") != "0" + + @staticmethod + def _singleton_bucket_seconds() -> tuple[float, ...] | None: + step_ms = int(os.environ.get("VLLM_OMNI_MOSS_REF_SINGLETON_BUCKET_MS", "0")) + if step_ms == 0: + return None + if not 80 <= step_ms <= 1000: + raise ValueError("MOSS reference singleton bucket step must be 80..1000 ms, or 0 to disable") + return tuple(ms / 1000 for ms in range(3000, 10001, step_ms)) + + def _reference_graphs(self) -> MossReferenceEncoderGraphs | None: + """Capture on first use, on the encode worker thread that replays them.""" + if self._graphs is None: + if not self._graph_capable(): + return None + graphs = MossReferenceEncoderGraphs( + self._processor.audio_tokenizer, + n_vq=self._n_vq, + compile_core=self._compile_graphs(), + batched_transfer=os.environ.get("VLLM_OMNI_MOSS_REF_BATCHED_TRANSFER", "0") == "1", + singleton_bucket_seconds=self._singleton_bucket_seconds(), + ) + graphs.capture() + self._graphs = graphs + return self._graphs + + def _tokenizer_rate(self) -> int: + return int(getattr(getattr(self._processor, "model_config", None), "sampling_rate", self._sr_target)) + + def _encoder_input(self, wav: torch.Tensor) -> torch.Tensor: + """A prepped clip at the tokenizer rate, one or two channels. + + This is the costly half of the processor's per-clip preparation; it runs + in the process that received the request, even with a shared encoder. + """ + import torchaudio + + if wav.ndim == 1: + wav = wav.unsqueeze(0) + if wav.shape[0] > 2: + wav = wav[:2] + if self._clip_rate() != self._sr_target: + # Channels resample independently, so a mono clip is resampled + # once before it is duplicated; the result is the same. + wav = torchaudio.functional.resample(wav, self._sr_target, self._tokenizer_rate()) + return wav + + def _clip_rate(self) -> int: + """Sample rate of the clips _encoder_input hands to the encoder.""" + return self._sr_target if self._gpu_prep else self._tokenizer_rate() + + def _finish_prepare(self, wav: torch.Tensor) -> torch.Tensor: + """The rest of the processor's preparation: two channels, normalized loudness.""" + if self._gpu_prep: + device = next(self._processor.audio_tokenizer.parameters()).device + wav = wav.to(device) + if self._clip_rate() != self._tokenizer_rate(): + wav = self._gpu_resampler(device)(wav) + if wav.shape[0] == 1: + wav = wav.repeat(2, 1) + return self._processor.loudness_normalize(wav) + + def _gpu_resampler(self, device: torch.device) -> torch.nn.Module: + resampler = self._gpu_resamplers.get(device) + if resampler is None: + import torchaudio + + # The same windowed-sinc filter as torchaudio.functional.resample's defaults. + resampler = torchaudio.transforms.Resample(self._sr_target, self._tokenizer_rate()).to(device) + self._gpu_resamplers[device] = resampler + return resampler + + def _processor_prepare(self, wav: torch.Tensor) -> torch.Tensor: + """Per-clip preparation of the processor's ``encode_audios_from_wav``.""" + return self._finish_prepare(self._encoder_input(wav)) async def aclose(self) -> None: """Release the batcher drainer + any in-flight encodes (tests/shutdown).""" @@ -385,7 +760,10 @@ class MossReferenceEncoder: 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]: + def _encoder_input(self, wav: torch.Tensor) -> torch.Tensor: + return wav # the realtime codec takes the clip at the working rate + + def _encode_prepared(self, prepared: list[torch.Tensor], graphs: Any = None) -> 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 [ diff --git a/vllm_omni/model_executor/models/moss_tts/reference_encoder_graphs.py b/vllm_omni/model_executor/models/moss_tts/reference_encoder_graphs.py new file mode 100644 index 000000000..52dead755 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/reference_encoder_graphs.py @@ -0,0 +1,399 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CUDA graphs for batched MOSS reference-audio encoding. + +Reference clips are encoded by the upstream MOSS audio tokenizer that the +processor loads. Its eager forward is dominated by kernel launches for the +small batches that serving produces. This runner replays graphs of the same +encoder and quantizer modules, captured per ``(batch, length)`` bucket. For a +given padded input the codes are bit-identical to the tokenizer's own +``_encode_frame``. Only the per-clip crop, which needs host values, runs +outside the graph. + +The encoder is causal, so padding a clip to its length bucket leaves its valid +frames unchanged up to kernel selection for the padded shape, the same effect +that batching already has when the tokenizer pads a batch to its longest clip. +""" + +from __future__ import annotations + +import math +import types +from dataclasses import dataclass + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +DEFAULT_BATCH_SIZES = (1, 2, 4, 8, 16, 32) +DEFAULT_BUCKET_SECONDS = (4.0, 6.0, 8.0, 10.0, 12.0, 16.0, 20.0) +# Fine steps keep padding small: a padded row or padded time is encoded in +# full. Worth the extra graphs where several encoders share one GPU. +FINE_BATCH_SIZES = (1, 2, 3, 4, 5, 6, 8) +FINE_BUCKET_SECONDS = (3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 10.0, 12.0, 14.0, 16.0, 20.0) +# Graphs are captured only while batch * bucket stays within this much audio, +# which bounds the shared graph memory pool by its largest batch. +DEFAULT_MAX_BATCH_SECONDS = 256.0 + + +@torch.compile(dynamic=False) +def _reference_rope_factors(seqlen: int, head_dim: int, max_period: float, device: torch.device): + # Use the same compiled FP32 arithmetic as the encoder. Eager libdevice + # sin/cos and compiled sin/cos can round differently near a VQ boundary. + ds = torch.arange(head_dim // 2, device=device, dtype=torch.float32) + freqs = torch.exp(ds * (-math.log(max_period) * 2 / head_dim)) + angles = torch.arange(seqlen, device=device, dtype=torch.float32)[:, None] * freqs[None, :] + return torch.cos(angles), torch.sin(angles) + + +def _cached_reference_rope(self, q: torch.Tensor, k: torch.Tensor, cache: dict): + if self.rope is None: + return q, k + batch, heads, seqlen, dim = q.shape + key = (dim, self.rope.max_period, q.device) + if key not in cache or cache[key][0].shape[0] < seqlen: + cache[key] = _reference_rope_factors(seqlen, dim, self.rope.max_period, q.device) + cos, sin = cache[key] + cos, sin = cos[:seqlen], sin[:seqlen] + qr, qi = q.view(batch, heads, seqlen, dim // 2, 2).unbind(-1) + kr, ki = k.view(batch, heads, seqlen, dim // 2, 2).unbind(-1) + qr, qi, kr, ki = qr.float(), qi.float(), kr.float(), ki.float() + # Preserve the original arithmetic, casts, and B/H/T/D output layout. + qo = torch.stack(((qr * cos - qi * sin).to(q.dtype), (qr * sin + qi * cos).to(q.dtype)), dim=-1) + ko = torch.stack(((kr * cos - ki * sin).to(k.dtype), (kr * sin + ki * cos).to(k.dtype)), dim=-1) + return qo.view(batch, heads, seqlen, dim), ko.view(batch, heads, seqlen, dim) + + +def _windowed_attention( + self, + x: torch.Tensor, + input_lengths: torch.Tensor, + *, + sdpa, + fa_version: int = 2, + rope_cache: dict | None = None, + skip_padded_query_mask: bool = False, +) -> torch.Tensor: + """The tokenizer's non-streaming attention with a local-window flash kernel. + + Same projection, RoPE and padded-row zeroing as the tokenizer's SDPA path. + That path masks each query to ``context`` keys (distance 0 .. context-1) + inside a dense block; the flash kernel visits only that window. + """ + batch, seqlen, _ = x.shape + q, k, v = self._project_qkv(x) + q, k = self._apply_dense_rope(q, k) if rope_cache is None else _cached_reference_rope(self, q, k, rope_cache) + if q.dtype not in (torch.bfloat16, torch.float16): + return sdpa(x, input_lengths) + from vllm.vllm_flash_attn import flash_attn_varlen_func + + heads, dim = q.shape[1], q.shape[-1] + + def packed(t: torch.Tensor) -> torch.Tensor: + return t.transpose(1, 2).reshape(batch * seqlen, heads, dim) + + starts = torch.arange(0, (batch + 1) * seqlen, seqlen, device=x.device, dtype=torch.int32) + window = [self.context - 1, 0] if self.causal and self.context is not None else [-1, -1] + out = flash_attn_varlen_func( + packed(q), + packed(k), + packed(v), + max_seqlen_q=seqlen, + cu_seqlens_q=starts, + max_seqlen_k=seqlen, + cu_seqlens_k=starts, + causal=self.causal, + window_size=window, + fa_version=fa_version, + ).view(batch, seqlen, heads, dim) + if not skip_padded_query_mask: + valid = (torch.arange(seqlen, device=x.device).view(1, seqlen) < input_lengths.view(-1, 1)).view( + batch, seqlen, 1, 1 + ) + out = torch.where(valid, out, torch.zeros((), device=out.device, dtype=out.dtype)) + return out.reshape(batch, seqlen, self.embed_dim) + + +def _validate_unmasked_reference_encoder(tokenizer: torch.nn.Module) -> None: + """Restrict the optimization to the audited HF encoder's valid-prefix contract. + + Causal attention and pointwise transformer operations cannot read future + padded rows. Its patch-downsample floors lengths, so every valid output + group consists entirely of valid input rows. The quantizer is pointwise + and the caller crops codes by those lengths. Padded *outputs* may differ. + This argument does not apply to the streaming decoder or to upsampling. + """ + transformers = 0 + for module in tokenizer.encoder: + name = type(module).__name__ + if name == "MossAudioTokenizerPatchedPretransform" and module.is_downsample: + continue + if name == "MossAudioTokenizerProjectedTransformer": + layers = module.transformer.layers + if layers and all(layer.self_attn.causal for layer in layers): + transformers += 1 + continue + raise ValueError("Unmasked reference queries require the causal MOSS patch-downsample encoder") + if not transformers: + raise ValueError("Unmasked reference queries require at least one causal MOSS transformer") + if type(tokenizer.quantizer).__name__ not in { + "MossAudioTokenizerResidualVQ", + "MossAudioTokenizerResidualLFQ", + }: + raise ValueError("Unmasked reference queries require the pointwise MOSS quantizer") + for module in tokenizer.quantizer.modules(): + if isinstance(module, torch.nn.Conv1d) and module.kernel_size != (1,): + raise ValueError("Unmasked reference queries do not support temporal quantizer convolutions") + + +def install_windowed_attention( + tokenizer: torch.nn.Module, + *, + fa_version: int = 2, + cache_rope: bool = False, + skip_padded_query_mask: bool = False, +) -> int: + """Route the encoder's non-streaming attention through a local-window flash kernel.""" + if skip_padded_query_mask: + _validate_unmasked_reference_encoder(tokenizer) + logger.info("MOSS reference encoder: skipping padded query masks (valid-prefix outputs only)") + count = 0 + rope_cache = {} if cache_rope else None + if cache_rope: + tokenizer._moss_reference_rope_cache = rope_cache + tokenizer._moss_reference_rope_prepared_samples = 0 + for module in tokenizer.encoder.modules(): + sdpa = getattr(module, "_forward_non_streaming_sdpa", None) + if sdpa is None or not hasattr(module, "_project_qkv") or not getattr(module, "causal", False): + continue + + def attention(self, x, input_lengths, _sdpa=sdpa): + return _windowed_attention( + self, + x, + input_lengths, + sdpa=_sdpa, + fa_version=fa_version, + rope_cache=rope_cache, + skip_padded_query_mask=skip_padded_query_mask, + ) + + module._forward_non_streaming_sdpa = types.MethodType(attention, module) + count += 1 + return count + + +@dataclass +class _Graph: + graph: torch.cuda.CUDAGraph + audio: torch.Tensor # (B, C, T) float32 input + lengths: torch.Tensor # (B,) int64 input, in samples per channel + codes: torch.Tensor # (n_vq, B, frames) output + code_lengths: torch.Tensor # (B,) output + + +class MossReferenceEncoderGraphs: + """Encode prepared reference waveforms with captured tokenizer graphs.""" + + def __init__( + self, + tokenizer: torch.nn.Module, + *, + n_vq: int, + batch_sizes: tuple[int, ...] = DEFAULT_BATCH_SIZES, + bucket_seconds: tuple[float, ...] = DEFAULT_BUCKET_SECONDS, + max_batch_seconds: float = DEFAULT_MAX_BATCH_SECONDS, + compile_core: bool = False, + batched_transfer: bool = False, + singleton_bucket_seconds: tuple[float, ...] | None = None, + ) -> None: + self._tokenizer = tokenizer + self._compile_core = compile_core + self._batched_transfer = batched_transfer + self._n_vq = int(n_vq) + self._device = next(tokenizer.parameters()).device + self._channels = int(tokenizer.number_channels) + frame = int(tokenizer.downsample_rate) + sampling_rate = int(tokenizer.sampling_rate) + # Whole frames, so the tokenizer's own frame padding is a no-op. + self._buckets = sorted({int(math.ceil(s * sampling_rate / frame)) * frame for s in bucket_seconds}) + self._singleton_buckets = sorted( + set(self._buckets) + | {int(math.ceil(s * sampling_rate / frame)) * frame for s in (singleton_bucket_seconds or ())} + ) + self._batch_sizes = sorted({int(b) for b in batch_sizes if int(b) > 0}) + self._max_batch_samples = float(max_batch_seconds) * sampling_rate + self._graphs: dict[tuple[int, int], _Graph] = {} + self._pool = None + self._transfer_buffers = None + + @property + def max_batch(self) -> int: + return self._batch_sizes[-1] + + @property + def captured(self) -> list[tuple[int, int]]: + return sorted(self._graphs) + + def _core(self, audio: torch.Tensor, lengths: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + # Mirrors the tokenizer's _encode_frame without its host-side crop. + tokenizer = self._tokenizer + values, value_lengths = tokenizer._flatten_channels_for_codec(audio, lengths) + with tokenizer._codec_inference_autocast(): + hidden, hidden_lengths = values, value_lengths + for module in tokenizer.encoder: + hidden, hidden_lengths = module(hidden, hidden_lengths) + _, codes, code_lengths = tokenizer.quantizer(hidden.float(), hidden_lengths, self._n_vq) + return codes, code_lengths + + @torch.no_grad() + def capture(self) -> None: + if self._device.type != "cuda" or self._graphs: + return + self._pool = torch.cuda.graph_pool_handle() + stream = torch.cuda.Stream(device=self._device) + prepared = getattr(self._tokenizer, "_moss_reference_rope_prepared_samples", None) + max_samples = max(self._singleton_buckets) + if prepared is not None and prepared < max_samples: + # Prime immutable FP32 factors at the largest length before any + # capture or compilation. Smaller lengths use a prefix of the + # same table; the dictionary stays fixed throughout capture and + # workers share tables without introducing dynamic-shape guards + # for a separate dictionary key at each sequence length. + audio = torch.zeros(1, self._channels, max_samples, device=self._device) + lengths = torch.full((1,), max_samples, dtype=torch.long, device=self._device) + self._core(audio, lengths) + self._tokenizer._moss_reference_rope_prepared_samples = max_samples + core = self._core + if self._compile_core: + # Fuses the encoder's many small elementwise kernels (casts, RoPE, + # residuals). Only capture runs the compiled code; replays do not. + core = torch.compile(self._core, dynamic=True) + # Each encoder stage is its own module shape: allow a compile per stage. + with torch._dynamo.config.patch(recompile_limit=64): + self._capture_all(core, stream) + if self._batched_transfer and self._graphs: + # One reusable slab per worker, rather than one per graph bucket. + entries = list(self._graphs.values()) + self._transfer_buffers = ( + torch.empty(max(e.audio.numel() for e in entries), dtype=entries[0].audio.dtype, pin_memory=True), + torch.empty(max(e.lengths.numel() for e in entries), dtype=entries[0].lengths.dtype, pin_memory=True), + torch.empty(max(e.codes.numel() for e in entries), dtype=entries[0].codes.dtype, pin_memory=True), + torch.empty( + max(e.code_lengths.numel() for e in entries), dtype=entries[0].code_lengths.dtype, pin_memory=True + ), + torch.cuda.Event(), + ) + logger.info( + "MOSS reference encoder CUDA graphs: batch=%s buckets(s)=%s compiled=%s captured=%d", + self._batch_sizes, + [round(s / int(self._tokenizer.sampling_rate), 2) for s in self._buckets], + self._compile_core, + len(self._graphs), + ) + if self._singleton_buckets != self._buckets: + logger.info( + "MOSS reference encoder singleton buckets(s): %s", + [round(s / int(self._tokenizer.sampling_rate), 2) for s in self._singleton_buckets], + ) + + def _capture_all(self, core, stream: torch.cuda.Stream) -> None: + for batch in reversed(self._batch_sizes): + buckets = self._singleton_buckets if batch == 1 else self._buckets + for samples in reversed(buckets): + if batch > 1 and batch * samples > self._max_batch_samples: + continue + audio = torch.zeros(batch, self._channels, samples, device=self._device) + lengths = torch.full((batch,), samples, dtype=torch.long, device=self._device) + try: + stream.wait_stream(torch.cuda.current_stream(self._device)) + with torch.cuda.stream(stream): + # Lazy library setup must happen before capture; one run does it. + core(audio, lengths) + torch.cuda.current_stream(self._device).wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=self._pool, stream=stream): + codes, code_lengths = core(audio, lengths) + except Exception: + logger.warning( + "MOSS reference encoder graph (B=%d, samples=%d) not captured; using eager", + batch, + samples, + exc_info=True, + ) + continue + self._graphs[(batch, samples)] = _Graph(graph, audio, lengths, codes, code_lengths) + + @torch.no_grad() + def encode(self, wavs: list[torch.Tensor]) -> list[torch.Tensor | None]: + """Codes ``(frames, n_vq)`` int64 on CPU per clip; None where no graph fits. + + ``wavs`` are prepared the way the processor prepares them: ``(C, T)`` + float32 after channel handling, resampling and loudness normalization. + Clips are grouped longest first, so each replay pads to the bucket of + similar-length clips. + """ + out: list[torch.Tensor | None] = [None] * len(wavs) + order = sorted(range(len(wavs)), key=lambda i: -int(wavs[i].shape[-1])) + pos = 0 + while pos < len(order): + longest = int(wavs[order[pos]].shape[-1]) + # Fine singleton buckets reduce padding without fragmenting an + # already assembled batch into separate graph replays. + buckets = self._singleton_buckets if len(order) - pos == 1 else self._buckets + samples = next((s for s in buckets if s >= longest and (1, s) in self._graphs), None) + if samples is None: + pos += 1 # longer than every bucket: left to the caller + continue + widest = max(b for b in self._batch_sizes if (b, samples) in self._graphs) + rows = order[pos : pos + widest] + batch = next(b for b in self._batch_sizes if b >= len(rows) and (b, samples) in self._graphs) + for row, codes in zip(rows, self._replay(self._graphs[(batch, samples)], [wavs[i] for i in rows])): + out[row] = codes + pos += len(rows) + return out + + def _replay(self, entry: _Graph, wavs: list[torch.Tensor]) -> list[torch.Tensor]: + if self._transfer_buffers is not None and all(w.device.type == "cpu" for w in wavs): + return self._replay_batched_transfer(entry, wavs) + count, batch, samples = len(wavs), entry.audio.shape[0], entry.audio.shape[-1] + # Unused rows encode full-length silence: a fully masked row would be + # all-NaN inside attention, and rows never mix. + entry.audio.zero_() + lengths = [int(w.shape[-1]) for w in wavs] + [samples] * (batch - count) + entry.lengths.copy_(torch.tensor(lengths, dtype=torch.long)) + for row, wav in enumerate(wavs): + entry.audio[row, :, : wav.shape[-1]].copy_(wav) + entry.graph.replay() + codes = entry.codes[:, :count].cpu() + code_lengths = entry.code_lengths[:count].tolist() + return [codes[:, row, : int(code_lengths[row])].transpose(0, 1).contiguous().long() for row in range(count)] + + def _replay_batched_transfer(self, entry: _Graph, wavs: list[torch.Tensor]) -> list[torch.Tensor]: + """Stage one batch on pinned host memory, then wait once for its outputs.""" + audio_slab, lengths_slab, codes_slab, code_lengths_slab, copied = self._transfer_buffers + count, batch, samples = len(wavs), entry.audio.shape[0], entry.audio.shape[-1] + audio = audio_slab[: entry.audio.numel()].view_as(entry.audio) + lengths = lengths_slab[:batch] + audio.zero_() + lengths.fill_(samples) + for row, wav in enumerate(wavs): + audio[row, :, : wav.shape[-1]].copy_(wav) + lengths[row] = wav.shape[-1] + entry.audio.copy_(audio, non_blocking=True) + entry.lengths.copy_(lengths, non_blocking=True) + entry.graph.replay() + result = entry.codes[:, :count] + codes = codes_slab[: result.numel()].view_as(result) + code_lengths = code_lengths_slab[:count] + codes.copy_(result, non_blocking=True) + code_lengths.copy_(entry.code_lengths[:count], non_blocking=True) + copied.record(torch.cuda.current_stream(self._device)) + copied.synchronize() + # Returned codes must outlive this worker's next replay, including + # one-frame clips whose transpose is already contiguous. + return [ + codes[:, row, : int(code_lengths[row])].transpose(0, 1).to(dtype=torch.long, copy=True).contiguous() + for row in range(count) + ] diff --git a/vllm_omni/model_executor/models/moss_tts/shared_reference_encoder.py b/vllm_omni/model_executor/models/moss_tts/shared_reference_encoder.py new file mode 100644 index 000000000..fec5c2860 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/shared_reference_encoder.py @@ -0,0 +1,334 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""One reference-audio encoder shared by the API processes of a server. + +Each API process otherwise loads its own copy of the audio tokenizer on the +GPU and encodes only the clips it receives. When the processes share a +directory (``VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR``), the first one to take the lock in it +hosts the encoder and serves the others over a Unix socket in that directory. +Requests from all processes are merged into batches, and only the host keeps +the tokenizer and its CUDA graphs on the GPU. + +Messages carry raw arrays and a JSON header; nothing is unpickled. A client +that cannot reach the host encodes locally instead. +""" + +from __future__ import annotations + +import collections +import fcntl +import json +import os +import threading +import time +from collections.abc import Callable +from contextlib import nullcontext +from multiprocessing.connection import Client, Connection, Listener + +import numpy as np +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_LOCK_NAME = "ref-encoder.lock" +_SOCKET_NAME = "ref-encoder.sock" +# Clips one worker merges into an encode, across all API processes. Several +# workers encode concurrently on their own streams: one serial queue of large +# batches makes a burst of new references wait behind each other. A free +# worker takes what is queued without waiting for more: API processes already +# gather their clips before sending them. +_MAX_BATCH = 8 +# A client gives up on the host after this long and encodes locally. +_CONNECT_TIMEOUT_S = 120.0 +_REQUEST_TIMEOUT_S = 120.0 +# First-time compilation can outlast a normal encode request. Readiness has a +# separate deadline and must finish before the API starts accepting requests. +_STARTUP_TIMEOUT_S = 1200.0 + +EncodeFn = Callable[[list[torch.Tensor]], list] + +# Open for the life of the process: closing it would release the host lock. +_host_lock_handle = None + + +class SharedReferenceEncoderStartupError(RuntimeError): + """The shared encoder cannot safely accept requests yet.""" + + +def elect_host(shared_dir: str) -> bool: + """True in the one process that holds the encoder lock for ``shared_dir``. + + The lock is held for the life of the process, so a restarted server + elects a new host even if the old socket file is still present. + """ + global _host_lock_handle + path = os.path.join(shared_dir, _LOCK_NAME) + handle = open(path, "a+") # noqa: SIM115 - kept open to hold the lock + try: + fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB) + except OSError: + handle.close() + return False + _host_lock_handle = handle + return True + + +def _send(conn: Connection, header: dict, arrays: list[np.ndarray]) -> None: + conn.send_bytes(json.dumps(header).encode()) + for array in arrays: + conn.send_bytes(np.ascontiguousarray(array).tobytes()) + + +def _recv(conn: Connection) -> tuple[dict, list[bytes]]: + header = json.loads(conn.recv_bytes().decode()) + return header, [conn.recv_bytes() for _ in range(int(header["count"]))] + + +class _Job: + __slots__ = ("wav", "done", "result", "queued_at") + + def __init__(self, wav: torch.Tensor) -> None: + self.wav = wav + self.done = threading.Event() + self.result: torch.Tensor | BaseException | None = None + self.queued_at = time.monotonic() + + +class SharedReferenceEncoderHost: + """Serves reference encodes for every API process sharing ``shared_dir``. + + ``workers`` pairs each worker's encode function with an optional warmup + (for example, capturing that worker's CUDA graphs). + """ + + def __init__( + self, + shared_dir: str, + workers: list[tuple[EncodeFn, Callable[[], object] | None]], + *, + stream_priority: int = 0, + batch_window_ms: float = 0.0, + ) -> None: + self._stream_priority = stream_priority + self._batch_window_s = max(0.0, float(batch_window_ms)) / 1000 + self._pending: collections.deque[_Job] = collections.deque() + # Includes jobs already taken by a worker, so shutdown can release + # every caller without waiting for an in-flight GPU operation. + self._outstanding: set[_Job] = set() + self._ready = threading.Condition() + self._path = os.path.join(shared_dir, _SOCKET_NAME) + if os.path.exists(self._path): + os.unlink(self._path) # left by a previous host; we hold the lock + self._listener = Listener(self._path, family="AF_UNIX") + os.chmod(self._path, 0o600) + self._closed = False + self._started = threading.Event() + self._startup_error: Exception | None = None + threading.Thread( + target=self._start_workers, args=(workers,), name="moss-ref-encoder-start", daemon=True + ).start() + threading.Thread(target=self._accept_loop, name="moss-ref-encoder-accept", daemon=True).start() + logger.info("MOSS shared reference encoder serving at %s with %d workers", self._path, len(workers)) + + def encode(self, wavs: list[torch.Tensor]) -> list: + """Encode on behalf of the host process itself.""" + self.wait_until_ready() + jobs = [_Job(wav) for wav in wavs] + with self._ready: + if self._closed: + raise RuntimeError("shared reference encoder is closed") + self._outstanding.update(jobs) + self._pending.extend(jobs) + self._ready.notify(len(jobs)) + for job in jobs: + job.done.wait() + return [job.result for job in jobs] + + def close(self) -> None: + with self._ready: + self._closed = True + self._started.set() + for job in self._outstanding: + job.result = RuntimeError("shared reference encoder closed before encoding completed") + job.done.set() + self._outstanding.clear() + self._pending.clear() + self._ready.notify_all() + try: + self._listener.close() + except OSError: + pass + + def wait_until_ready(self) -> None: + if not self._started.wait(_STARTUP_TIMEOUT_S): + raise SharedReferenceEncoderStartupError("shared reference encoder warmup timed out") + if self._closed or self._startup_error is not None: + raise SharedReferenceEncoderStartupError("shared reference encoder startup failed") from self._startup_error + + def _accept_loop(self) -> None: + while not self._closed: + try: + conn = self._listener.accept() + except OSError: + return + threading.Thread(target=self._serve, args=(conn,), name="moss-ref-encoder-conn", daemon=True).start() + + def _serve(self, conn: Connection) -> None: + with conn: + while True: + try: + header, payloads = _recv(conn) + except (EOFError, OSError): + return + if header.get("op") == "ready": + try: + self.wait_until_ready() + error = None + except SharedReferenceEncoderStartupError as exc: + error = str(exc) + try: + _send(conn, {"count": 0, "ready": error is None, "error": error}, []) + except OSError: + return + continue + wavs = [ + torch.from_numpy(np.frombuffer(data, dtype=np.float32).reshape(shape).copy()) + for data, shape in zip(payloads, header["shapes"]) + ] + results = self.encode(wavs) + errors = [None if isinstance(r, torch.Tensor) else repr(r) for r in results] + codes = [r.to(torch.int32).numpy() for r in results if isinstance(r, torch.Tensor)] + reply = {"count": len(codes), "shapes": [list(c.shape) for c in codes], "errors": errors} + try: + _send(conn, reply, codes) + except OSError: + return + + def _start_workers(self, workers: list[tuple[EncodeFn, Callable[[], object] | None]]) -> None: + # Every warmup (graph capture) finishes before any worker encodes: CUDA + # work in another thread would invalidate a capture in progress. + try: + for _, warmup in workers: + if warmup is not None: + warmup() + if self._closed: + return + for index, (encode_fn, _) in enumerate(workers): + threading.Thread( + target=self._encode_loop, args=(encode_fn,), name=f"moss-ref-encoder-{index}", daemon=True + ).start() + logger.info("MOSS shared reference encoder ready: %d workers", len(workers)) + except Exception as exc: # noqa: BLE001 - wake readiness waiters with the failure + self._startup_error = exc + logger.exception("MOSS shared reference encoder startup failed") + finally: + self._started.set() + + def _encode_loop(self, encode_fn: EncodeFn) -> None: + stream = torch.cuda.Stream(priority=self._stream_priority) if torch.cuda.is_available() else None + with torch.cuda.stream(stream) if stream is not None else nullcontext(): + self._serve_jobs(encode_fn) + + def _serve_jobs(self, encode_fn: EncodeFn) -> None: + while True: + with self._ready: + while True: + if self._closed: + return + if not self._pending: + self._ready.wait() + continue + remaining = self._pending[0].queued_at + self._batch_window_s - time.monotonic() + if remaining > 0 and len(self._pending) < _MAX_BATCH: + # Keep clips in the common queue while collecting; + # workers must not each reserve a singleton first. + self._ready.wait(remaining) + continue + break + jobs = [self._pending.popleft() for _ in range(min(_MAX_BATCH, len(self._pending)))] + try: + results = encode_fn([job.wav for job in jobs]) + if len(results) != len(jobs): + raise RuntimeError( + f"shared reference encoder returned {len(results)} results for {len(jobs)} clips" + ) + except Exception as error: # noqa: BLE001 - every waiter gets the failure + results = [error] * len(jobs) + with self._ready: + for job, result in zip(jobs, results, strict=True): + # close() may already have completed this caller with an + # error. A late GPU result must not replace that outcome. + if not job.done.is_set(): + job.result = result + job.done.set() + self._outstanding.discard(job) + + +class SharedReferenceEncoderClient: + """Sends reference clips to the host process of ``shared_dir``.""" + + def __init__(self, shared_dir: str) -> None: + self._path = os.path.join(shared_dir, _SOCKET_NAME) + # One connection per request in flight; idle ones are reused. + self._idle: list[Connection] = [] + self._lock = threading.Lock() + self._connected = False + + def _connect(self) -> Connection: + # The host may still be starting until the first connection succeeds; + # later reconnects are quick. + deadline = time.monotonic() + (1.0 if self._connected else _CONNECT_TIMEOUT_S) + while True: + try: + conn = Client(self._path, family="AF_UNIX") + self._connected = True + break + except (FileNotFoundError, ConnectionRefusedError): + if time.monotonic() >= deadline: + raise + time.sleep(0.2) + try: + _send(conn, {"op": "ready", "count": 0}, []) + if not conn.poll(_STARTUP_TIMEOUT_S): + raise SharedReferenceEncoderStartupError("shared reference encoder readiness timed out") + header, _ = _recv(conn) + if not header.get("ready"): + raise SharedReferenceEncoderStartupError(header.get("error", "shared reference encoder not ready")) + return conn + except BaseException: + conn.close() + raise + + def wait_until_ready(self) -> None: + """Startup handshake, without encoding or populating the reference cache.""" + try: + conn = self._connect() + except Exception as exc: + raise SharedReferenceEncoderStartupError("shared reference encoder is not ready") from exc + with self._lock: + self._idle.append(conn) + + def encode(self, wavs: list[torch.Tensor]) -> list: + """Codes (or an exception) per clip; raises when the host is unreachable.""" + with self._lock: + conn = self._idle.pop() if self._idle else None + try: + if conn is None: + conn = self._connect() + arrays = [w.detach().to(torch.float32).cpu().numpy() for w in wavs] + _send(conn, {"count": len(arrays), "shapes": [list(a.shape) for a in arrays]}, arrays) + if not conn.poll(_REQUEST_TIMEOUT_S): + raise TimeoutError(f"no reply from the shared reference encoder in {_REQUEST_TIMEOUT_S}s") + header, payloads = _recv(conn) + except BaseException: + if conn is not None: + conn.close() + raise + with self._lock: + self._idle.append(conn) + codes = iter( + torch.from_numpy(np.frombuffer(data, dtype=np.int32).reshape(shape).copy()).long() + for data, shape in zip(payloads, header["shapes"]) + ) + return [next(codes) if error is None else RuntimeError(error) for error in header["errors"]] diff --git a/vllm_omni/model_executor/models/moss_tts/slot_attention.py b/vllm_omni/model_executor/models/moss_tts/slot_attention.py new file mode 100644 index 000000000..f7f6d2783 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/slot_attention.py @@ -0,0 +1,475 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Direct-slot ring attention for MOSS streaming codec. + +Inputs q/k/v are (execution_batch, heads, padded_T, head_dim). cache is +contiguous (2, state_capacity, heads, ring_capacity, head_dim). end_offset +contains absolute *pre-call* positions per state slot. slot_ids and valid_lengths +contain one entry per execution row. valid_lengths=0 is a scratch/padding row. + +This op writes only valid tokens, attends to the resulting ring with the same +position semantics as RingKVCache.complete, then advances only active offsets. +Slots must be unique and in bounds; 0 <= valid_lengths <= padded_T. Chunks +longer than the ring retain only their final ring_capacity tokens. +These device-value preconditions are caller-owned to avoid host synchronization. +The op deliberately preserves chunk-complete ring overwrite semantics; it does +not recover entries overwritten by later tokens of the same input chunk. +""" + +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def _write( + k_ptr, + v_ptr, + cache_ptr, + end_ptr, + slots_ptr, + lengths_ptr, + ks0: tl.constexpr, + ks1: tl.constexpr, + ks2: tl.constexpr, + ks3: tl.constexpr, + vs0: tl.constexpr, + vs1: tl.constexpr, + vs2: tl.constexpr, + vs3: tl.constexpr, + end_stride: tl.constexpr, + slot_stride: tl.constexpr, + length_stride: tl.constexpr, + state_capacity: tl.constexpr, + num_heads: tl.constexpr, + num_frames: tl.constexpr, + ring_capacity: tl.constexpr, + head_dim: tl.constexpr, + block_size: tl.constexpr, +): + row = tl.program_id(0) + slot = tl.load(slots_ptr + row * slot_stride) + length = tl.load(lengths_ptr + row * length_stride) + start = tl.load(end_ptr + slot * end_stride) + x = tl.program_id(1) * block_size + tl.arange(0, block_size) + dim = x % head_dim + token = (x // head_dim) % num_frames + head = x // (num_frames * head_dim) + # For chunks longer than the ring, write only the final C tokens. This + # avoids duplicate physical addresses (and write races) in one launch. + skip = tl.maximum(length - ring_capacity, 0) + source_token = token + skip + mask = (head < num_heads) & (source_token < length) + pos = (start + source_token) % ring_capacity + dst = ((slot * num_heads + head) * ring_capacity + pos) * head_dim + dim + kval = tl.load(k_ptr + row * ks0 + head * ks1 + source_token * ks2 + dim * ks3, mask, 0) + vval = tl.load(v_ptr + row * vs0 + head * vs1 + source_token * vs2 + dim * vs3, mask, 0) + tl.store(cache_ptr + dst, kval, mask) + tl.store(cache_ptr + state_capacity * num_heads * ring_capacity * head_dim + dst, vval, mask) + + +@triton.jit +def _attend( + q_ptr, + cache_ptr, + end_ptr, + slots_ptr, + lengths_ptr, + out_ptr, + qs0: tl.constexpr, + qs1: tl.constexpr, + qs2: tl.constexpr, + qs3: tl.constexpr, + end_stride: tl.constexpr, + slot_stride: tl.constexpr, + length_stride: tl.constexpr, + state_capacity: tl.constexpr, + num_heads: tl.constexpr, + num_frames: tl.constexpr, + ring_capacity: tl.constexpr, + head_dim: tl.constexpr, + context_size: tl.constexpr, + block_m: tl.constexpr, + block_n: tl.constexpr, + skip_empty_tiles: tl.constexpr = False, +): + rows = tl.program_id(0) * block_m + tl.arange(0, block_m) + bh = tl.program_id(1) + batch, head = bh // num_heads, bh % num_heads + slot = tl.load(slots_ptr + batch * slot_stride) + length = tl.load(lengths_ptr + batch * length_stride) + start = tl.load(end_ptr + slot * end_stride) + end = start + length + dims = tl.arange(0, head_dim) + q = tl.load( + q_ptr + batch * qs0 + head * qs1 + rows[:, None] * qs2 + dims[None, :] * qs3, rows[:, None] < num_frames, 0 + ) + acc = tl.full((block_m, head_dim), 0.0, tl.float32) + maximum = tl.full((block_m,), -float("inf"), tl.float32) + denominator = tl.full((block_m,), 0.0, tl.float32) + # Padding rows own scratch slots and must emit zeros without reading KV. + # Keep valid-row accumulation and the ring overwrite semantics unchanged. + if skip_empty_tiles: + block_start = tl.program_id(0) * block_m + live = (length > 0) & (block_start < length) + live = live & (block_start + block_m > tl.maximum(length - ring_capacity, 0)) + else: + live = length > 0 + if live: + # Physical-ring order matches the existing masked_attention accumulation. + last = end - 1 + last_index = last % ring_capacity + if skip_empty_tiles: + # Before the first wrap, only physical keys [0, end) exist. + # Omitted blocks have zero probability for every query; keep the + # physical order and arithmetic of all remaining blocks unchanged. + key_blocks = tl.cdiv(tl.minimum(end, ring_capacity), block_n) + else: + key_blocks = tl.cdiv(ring_capacity, block_n) + for first in range(key_blocks): + cols = first * block_n + tl.arange(0, block_n) + delta = cols - last_index + positions = tl.where(delta <= 0, last + delta, last + delta - ring_capacity) + distance = start + rows[:, None] - positions[None, :] + allowed = ( + (rows[:, None] < length) + & (cols[None, :] < ring_capacity) + & (positions[None, :] >= 0) + & (positions[None, :] < end) + & (distance >= 0) + ) + if context_size > 0: + allowed = allowed & (distance < context_size) + base = (slot * num_heads + head) * ring_capacity * head_dim + k = tl.load(cache_ptr + base + cols[None, :] * head_dim + dims[:, None], cols[None, :] < ring_capacity, 0) + score = tl.dot(q, k).to(tl.float32) * (head_dim**-0.5) + score = tl.where(allowed, score, -float("inf")) + next_max = tl.maximum(maximum, tl.max(score, 1)) + safe_max = tl.where(next_max == -float("inf"), 0.0, next_max) + rescale = tl.exp(maximum - safe_max) + p = tl.exp(score - safe_max[:, None]) + denominator = denominator * rescale + tl.sum(p, 1) + v = tl.load( + cache_ptr + + state_capacity * num_heads * ring_capacity * head_dim + + base + + cols[:, None] * head_dim + + dims[None, :], + cols[:, None] < ring_capacity, + 0, + ) + acc = acc * rescale[:, None] + tl.dot(p.to(v.dtype), v) + maximum = next_max + out = acc / tl.where(denominator > 0, denominator, 1.0)[:, None] + tl.store( + out_ptr + ((batch * num_heads + head) * num_frames + rows[:, None]) * head_dim + dims[None, :], + out, + rows[:, None] < num_frames, + ) + + +@triton.jit +def _advance( + end_ptr, + slots_ptr, + lengths_ptr, + end_stride: tl.constexpr, + slot_stride: tl.constexpr, + length_stride: tl.constexpr, + batch_size: tl.constexpr, + block_size: tl.constexpr, +): + row = tl.arange(0, block_size) + slot = tl.load(slots_ptr + row * slot_stride, row < batch_size, 0) + length = tl.load(lengths_ptr + row * length_stride, row < batch_size, 0) + active = (row < batch_size) & (length > 0) + start = tl.load(end_ptr + slot * end_stride, active, 0) + tl.store(end_ptr + slot * end_stride, start + length, active) + + +@triton.jit +def _write_attend( + q_ptr, + k_ptr, + v_ptr, + cache_ptr, + starts_ptr, + slots_ptr, + lengths_ptr, + out_ptr, + qs0: tl.constexpr, + qs1: tl.constexpr, + qs2: tl.constexpr, + qs3: tl.constexpr, + ks0: tl.constexpr, + ks1: tl.constexpr, + ks2: tl.constexpr, + ks3: tl.constexpr, + vs0: tl.constexpr, + vs1: tl.constexpr, + vs2: tl.constexpr, + vs3: tl.constexpr, + start_stride: tl.constexpr, + slot_stride: tl.constexpr, + length_stride: tl.constexpr, + state_capacity: tl.constexpr, + num_heads: tl.constexpr, + num_frames: tl.constexpr, + ring_capacity: tl.constexpr, + head_dim: tl.constexpr, + context_size: tl.constexpr, + write_frames: tl.constexpr, + block_m: tl.constexpr, + block_n: tl.constexpr, + block_w: tl.constexpr, +): + rows = tl.program_id(0) * block_m + tl.arange(0, block_m) + bh = tl.program_id(1) + batch, head = bh // num_heads, bh % num_heads + slot = tl.load(slots_ptr + batch * slot_stride) + length = tl.load(lengths_ptr + batch * length_stride) + start = tl.load(starts_ptr + batch * start_stride) + end = start + length + dims = tl.arange(0, head_dim) + base = (slot * num_heads + head) * ring_capacity * head_dim + v_base = state_capacity * num_heads * ring_capacity * head_dim + base + # Every program of this (row, head) writes the whole chunk for its head, + # so it only reads ring entries it wrote itself (duplicate writes across + # query blocks store identical values). Chunks longer than the ring keep + # only their final ring_capacity tokens, as in _write. + skip = tl.maximum(length - ring_capacity, 0) + for first in range(tl.cdiv(write_frames, block_w)): + token = first * block_w + tl.arange(0, block_w) + source = token + skip + mask = (source < length)[:, None] + pos = (start + source) % ring_capacity + kval = tl.load(k_ptr + batch * ks0 + head * ks1 + source[:, None] * ks2 + dims[None, :] * ks3, mask, 0) + vval = tl.load(v_ptr + batch * vs0 + head * vs1 + source[:, None] * vs2 + dims[None, :] * vs3, mask, 0) + dst = pos[:, None] * head_dim + dims[None, :] + tl.store(cache_ptr + base + dst, kval, mask) + tl.store(cache_ptr + v_base + dst, vval, mask) + tl.debug_barrier() + q = tl.load( + q_ptr + batch * qs0 + head * qs1 + rows[:, None] * qs2 + dims[None, :] * qs3, rows[:, None] < num_frames, 0 + ) + acc = tl.full((block_m, head_dim), 0.0, tl.float32) + maximum = tl.full((block_m,), -float("inf"), tl.float32) + denominator = tl.full((block_m,), 0.0, tl.float32) + if length > 0: + last = end - 1 + last_index = last % ring_capacity + for first in range(tl.cdiv(ring_capacity, block_n)): + cols = first * block_n + tl.arange(0, block_n) + delta = cols - last_index + positions = tl.where(delta <= 0, last + delta, last + delta - ring_capacity) + distance = start + rows[:, None] - positions[None, :] + allowed = ( + (rows[:, None] < length) + & (cols[None, :] < ring_capacity) + & (positions[None, :] >= 0) + & (positions[None, :] < end) + & (distance >= 0) + ) + if context_size > 0: + allowed = allowed & (distance < context_size) + k = tl.load(cache_ptr + base + cols[None, :] * head_dim + dims[:, None], cols[None, :] < ring_capacity, 0) + score = tl.dot(q, k).to(tl.float32) * (head_dim**-0.5) + score = tl.where(allowed, score, -float("inf")) + next_max = tl.maximum(maximum, tl.max(score, 1)) + safe_max = tl.where(next_max == -float("inf"), 0.0, next_max) + rescale = tl.exp(maximum - safe_max) + p = tl.exp(score - safe_max[:, None]) + denominator = denominator * rescale + tl.sum(p, 1) + v = tl.load( + cache_ptr + v_base + cols[:, None] * head_dim + dims[None, :], + cols[:, None] < ring_capacity, + 0, + ) + acc = acc * rescale[:, None] + tl.dot(p.to(v.dtype), v) + maximum = next_max + out = acc / tl.where(denominator > 0, denominator, 1.0)[:, None] + tl.store( + out_ptr + ((batch * num_heads + head) * num_frames + rows[:, None]) * head_dim + dims[None, :], + out, + rows[:, None] < num_frames, + ) + + +@torch.library.custom_op("vllm_omni::moss_slot_ring_attention", mutates_args=("cache", "end_offset")) +def slot_ring_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cache: torch.Tensor, + end_offset: torch.Tensor, + slot_ids: torch.Tensor, + valid_lengths: torch.Tensor, + context: int = -1, + skip_empty_tiles: bool = False, +) -> torch.Tensor: + """Fused address/mask path, with three ordered launches and no KV gather. + + q/k are already rotary-position encoded. The caller still advances its + separate MHA/RoPE offset state; this op advances only end_offset in place. + context=-1 means unbounded causal attention over available ring entries. + Outputs beyond each row's valid_lengths are exactly zero. + skip_empty_tiles skips wholly masked tiles for small (T <= 32) chunks. + Larger chunks keep the static loop, which benchmarks faster after wrap. + """ + b, h, t, d = q.shape + if k.shape != q.shape or v.shape != q.shape: + raise ValueError("q/k/v must have identical shapes") + if q.dtype not in (torch.bfloat16, torch.float16) or k.dtype != q.dtype or v.dtype != q.dtype: + raise ValueError("q/k/v must share bf16 or fp16 dtype") + if d not in (32, 64, 128) or t < 1: + raise ValueError("requires head_dim 32/64/128 and positive padded_T") + if cache.ndim != 5 or cache.shape[0] != 2 or cache.shape[2] != h or cache.shape[4] != d: + raise ValueError("cache must be (2, state_capacity, heads, capacity, head_dim)") + s, c = cache.shape[1], cache.shape[3] + if not cache.is_contiguous() or cache.dtype != q.dtype: + raise ValueError("cache must be contiguous and share q dtype") + if end_offset.shape != (s,) or slot_ids.shape != (b,) or valid_lengths.shape != (b,): + raise ValueError("invalid metadata shapes") + if any(x.dtype not in (torch.int32, torch.int64) for x in (end_offset, slot_ids, valid_lengths)): + raise ValueError("metadata must be int32/int64") + if q.device.type != "cuda" or any(x.device != q.device for x in (k, v, cache, end_offset, slot_ids, valid_lengths)): + raise ValueError("all tensors must share one CUDA device") + if context == 0 or context < -1: + raise ValueError("context must be -1 or positive") + out = torch.empty(q.shape, device=q.device, dtype=q.dtype) + if b == 0: + return out + strides = (end_offset.stride(0), slot_ids.stride(0), valid_lengths.stride(0)) + _write[(b, triton.cdiv(h * t * d, 256))]( + k, v, cache, end_offset, slot_ids, valid_lengths, *k.stride(), *v.stride(), *strides, s, h, t, c, d, 256 + ) + bm = 16 if t <= 32 else 64 + _attend[(triton.cdiv(t, bm), b * h)]( + q, + cache, + end_offset, + slot_ids, + valid_lengths, + out, + *q.stride(), + *strides, + s, + h, + t, + c, + d, + context, + bm, + 64, + skip_empty_tiles=skip_empty_tiles and t <= 32, + num_warps=8 if c <= 256 else 4, + num_stages=2, + ) + _advance[(1,)](end_offset, slot_ids, valid_lengths, *strides, b, triton.next_power_of_2(b)) + return out + + +@slot_ring_attention.register_fake +def _( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cache: torch.Tensor, + end_offset: torch.Tensor, + slot_ids: torch.Tensor, + valid_lengths: torch.Tensor, + context: int = -1, + skip_empty_tiles: bool = False, +) -> torch.Tensor: + return torch.empty(q.shape, device=q.device, dtype=q.dtype) + + +def _check_inputs(q, k, v, cache, metadata, context): + b, h, t, d = q.shape + if k.shape != q.shape or v.shape != q.shape: + raise ValueError("q/k/v must have identical shapes") + if q.dtype not in (torch.bfloat16, torch.float16) or k.dtype != q.dtype or v.dtype != q.dtype: + raise ValueError("q/k/v must share bf16 or fp16 dtype") + if d not in (32, 64, 128) or t < 1: + raise ValueError("requires head_dim 32/64/128 and positive padded_T") + if cache.ndim != 5 or cache.shape[0] != 2 or cache.shape[2] != h or cache.shape[4] != d: + raise ValueError("cache must be (2, state_capacity, heads, capacity, head_dim)") + if not cache.is_contiguous() or cache.dtype != q.dtype: + raise ValueError("cache must be contiguous and share q dtype") + if any(x.dtype not in (torch.int32, torch.int64) for x in metadata): + raise ValueError("metadata must be int32/int64") + if q.device.type != "cuda" or any(x.device != q.device for x in (k, v, cache, *metadata)): + raise ValueError("all tensors must share one CUDA device") + if context == 0 or context < -1: + raise ValueError("context must be -1 or positive") + + +@torch.library.custom_op("vllm_omni::moss_slot_ring_attention_rows", mutates_args=("cache",)) +def slot_ring_attention_rows( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cache: torch.Tensor, + starts: torch.Tensor, + slot_ids: torch.Tensor, + valid_lengths: torch.Tensor, + context: int = -1, +) -> torch.Tensor: + """``slot_ring_attention`` as one launch, with per-row start positions. + + ``starts`` holds each execution row's pre-call absolute position (the + caller's gathered offset). The ring is written and attended exactly as by + ``slot_ring_attention``; no offset is advanced here, so the caller owns + every offset update and can fuse them with its own. + """ + b, h, t, d = q.shape + _check_inputs(q, k, v, cache, (starts, slot_ids, valid_lengths), context) + if starts.shape != (b,) or slot_ids.shape != (b,) or valid_lengths.shape != (b,): + raise ValueError("invalid metadata shapes") + out = torch.empty(q.shape, device=q.device, dtype=q.dtype) + if b == 0: + return out + s, c = cache.shape[1], cache.shape[3] + bm = 16 if t <= 32 else 64 + _write_attend[(triton.cdiv(t, bm), b * h)]( + q, + k, + v, + cache, + starts, + slot_ids, + valid_lengths, + out, + *q.stride(), + *k.stride(), + *v.stride(), + starts.stride(0), + slot_ids.stride(0), + valid_lengths.stride(0), + s, + h, + t, + c, + d, + context, + min(t, c), + bm, + 64, + 16, + num_warps=8 if c <= 256 else 4, + num_stages=2, + ) + return out + + +@slot_ring_attention_rows.register_fake +def _( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cache: torch.Tensor, + starts: torch.Tensor, + slot_ids: torch.Tensor, + valid_lengths: torch.Tensor, + context: int = -1, +) -> torch.Tensor: + return torch.empty(q.shape, device=q.device, dtype=q.dtype) diff --git a/vllm_omni/model_executor/output_snapshot.py b/vllm_omni/model_executor/output_snapshot.py new file mode 100644 index 000000000..93ab0491e --- /dev/null +++ b/vllm_omni/model_executor/output_snapshot.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Owned, dtype-grouped snapshots for asynchronous runner outputs.""" + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import torch +from torch.utils._pytree import tree_flatten, tree_unflatten + + +@dataclass(frozen=True) +class RequestOutputSnapshot: + """Already partitioned, CPU-owned payloads in the current request order. + + Model finalizers may return this after D2H completes to avoid recursively + cloning and partitioning data whose request ownership is already known. + """ + + inter_stage: list[dict[str, Any] | None] + client: list[dict[str, Any] | None] | None = None + + +class PackedOutputSnapshot(dict): + """A normal payload mapping with an internal batched-copy plan. + + The runner's snapshot-slot event protects the device slabs until D2H is + complete. Host slabs are allocated per output, so downstream consumers can + retain their views without depending on reuse of the device ring slot. + """ + + def __init__(self, leaves: list[Any], spec: Any, groups: list[tuple[list[int], torch.Tensor]]) -> None: + super().__init__(tree_unflatten(leaves, spec)) + self._leaves = leaves + self._spec = spec + self._groups = groups + self.producer_event: torch.cuda.Event | None = None + + def record_producer_event(self, stream: torch.cuda.Stream) -> None: + """Publish readiness when a model produces on a separate CUDA stream.""" + self.producer_event = torch.cuda.Event() + self.producer_event.record(stream) + + def copy_to_cpu(self, copy_tensor: Callable[[torch.Tensor], torch.Tensor]) -> dict[str, Any]: + leaves = list(self._leaves) + for indices, slab in self._groups: + cpu = copy_tensor(slab) + offset = 0 + for index in indices: + value = leaves[index] + size = value.numel() + leaves[index] = cpu.narrow(0, offset, size).view(value.shape) + offset += size + return tree_unflatten(leaves, self._spec) + + +def pack_output_snapshot( + payload: dict[str, Any], + slot: dict[tuple[Any, ...], torch.Tensor], + *, + max_buckets: int, + reuse_existing_storage: bool = False, +) -> PackedOutputSnapshot | None: + """Copy tensor leaves once per dtype/device, retaining their exact values. + + No CUDA synchronization is introduced here. The caller must select the + producer stream and wait for the previous consumer before reusing a slot. + """ + leaves, spec = tree_flatten(payload) + grouped: dict[tuple[Any, ...], list[int]] = defaultdict(list) + for index, value in enumerate(leaves): + if isinstance(value, torch.Tensor): + # Preserve the existing per-tensor path for unusual layouts. + if value.layout != torch.strided or value.is_quantized: + return None + grouped[(value.dtype, value.device)].append(index) + groups = [] + for (dtype, device), indices in grouped.items(): + values = [leaves[index].detach().reshape(-1) for index in indices] + size = sum(value.numel() for value in values) + # A producer may have already written the leaves into one contiguous + # per-forward slab. Reuse that storage directly instead of copying + # it into a second packed slab. The strict offset/stride check keeps + # this optimization limited to views whose order matches the copy + # plan; arbitrary payloads retain the existing copy path. + if reuse_existing_storage and values and all(value.is_contiguous() for value in values): + storage = values[0].untyped_storage() + offset = values[0].storage_offset() + cursor = offset + shares_storage = True + for value in values: + if value.untyped_storage().data_ptr() != storage.data_ptr() or value.storage_offset() != cursor: + shares_storage = False + break + cursor += value.numel() + if shares_storage: + slab = values[0].as_strided((size,), (1,), storage_offset=offset) + offset = 0 + for index, value in zip(indices, values): + shape = leaves[index].shape + leaves[index] = slab.narrow(0, offset, value.numel()).view(shape) + offset += value.numel() + groups.append((indices, slab)) + continue + key = ("packed-output", dtype, device, size) + slab = slot.get(key) + if slab is None: + slab = torch.empty(size, dtype=dtype, device=device) + if len(slot) < max_buckets: + slot[key] = slab + if len(values) == 1: + slab.copy_(values[0]) + else: + torch.cat(values, out=slab) + offset = 0 + for index, value in zip(indices, values): + shape = leaves[index].shape + leaves[index] = slab.narrow(0, offset, value.numel()).view(shape) + offset += value.numel() + groups.append((indices, slab)) + return PackedOutputSnapshot(leaves, spec, groups) diff --git a/vllm_omni/model_executor/stage_input_processors/moss_tts.py b/vllm_omni/model_executor/stage_input_processors/moss_tts.py index e07f552b1..bdb7f8e92 100644 --- a/vllm_omni/model_executor/stage_input_processors/moss_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/moss_tts.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"). @@ -9,15 +12,13 @@ from collections import defaultdict from collections.abc import Mapping from typing import Any +import numpy as np import torch from vllm.inputs import TokensPrompt as OmniTokensPrompt from vllm.logger import init_logger from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct -from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import ( - parse_chunk_ramp, - ramp_chunk_size, -) +from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp, ramp_chunk_size logger = init_logger(__name__) @@ -237,7 +238,8 @@ def talker2codec_raw_async_chunk( Stage 0 emits newly generated raw codec rows shaped ``[T, n_vq]`` (normally ``[1, n_vq]`` per decode step). This processor buffers those new rows until a codec chunk is ready, then forwards the chunk to Stage 1. No delay-pattern - de-delay is applied on this path. + de-delay is applied on this path. An optional ``codec_chunk_ramp`` selects + successive chunk sizes and takes precedence over the initial chunk size. """ external_req_id = getattr(request, "external_req_id", None) req_id = str(external_req_id if external_req_id is not None else getattr(request, "request_id", id(request))) @@ -248,6 +250,8 @@ def talker2codec_raw_async_chunk( transfer_manager.request_payload = {} if not hasattr(transfer_manager, "put_req_chunk"): transfer_manager.put_req_chunk = defaultdict(int) + if not hasattr(transfer_manager, "ramp_chunk_count"): + transfer_manager.ramp_chunk_count = defaultdict(int) pending_frames = transfer_manager.code_prompt_token_ids[req_id] @@ -255,14 +259,16 @@ def talker2codec_raw_async_chunk( codes_dict = multimodal_output.get("codes", {}) or {} new_frames = codes_dict.get("audio") if isinstance(new_frames, torch.Tensor) and new_frames.numel() > 0: - frames_cpu = new_frames.detach().to("cpu", torch.long).contiguous() - if frames_cpu.ndim == 1: - frames_cpu = frames_cpu.reshape(1, -1) - if frames_cpu.ndim != 2: - raise ValueError(f"MOSS raw codec frames must be 2-D, got {tuple(frames_cpu.shape)}") - valid_rows = frames_cpu.ne(_MOSS_AUDIO_PAD_CODE).any(dim=1) - for frame in frames_cpu[valid_rows]: - pending_frames.append(frame.clone()) + # This runs per request on every decode step. Tiny torch ops cost + # tens of microseconds each here; NumPy on the host copy does not. + frames_np = new_frames.detach().cpu().numpy() + if frames_np.ndim == 1: + frames_np = frames_np.reshape(1, -1) + if frames_np.ndim != 2: + raise ValueError(f"MOSS raw codec frames must be 2-D, got {tuple(frames_np.shape)}") + valid_rows = (frames_np != _MOSS_AUDIO_PAD_CODE).any(axis=1) + owned = np.array(frames_np if valid_rows.all() else frames_np[valid_rows], dtype=np.int64) + pending_frames.extend(owned) # Raw/local streaming should mirror the non-streaming path: the codec # decodes only generated audio rows. Reference audio conditions the # talker, but feeding its codes into the codec streaming state adds a @@ -288,26 +294,27 @@ def talker2codec_raw_async_chunk( initial_chunk_frames = chunk_frames pending = len(pending_frames) - # codec_chunk_ramp (opt-in): static ladder of chunk sizes replacing the - # initial_codec_chunk_frames -> codec_chunk_frames cliff. The chunk index - # is maintained by the connector (ramp_chunk_count, incremented per - # emitted chunk); past the ladder, ramp_chunk_size falls back to steady. - if not hasattr(transfer_manager, "_moss_ramp_parsed"): - transfer_manager._moss_ramp_parsed = parse_chunk_ramp(cfg, steady=chunk_frames) - _ramp = transfer_manager._moss_ramp_parsed - if _ramp is not None: - if not hasattr(transfer_manager, "ramp_chunk_count"): - transfer_manager.ramp_chunk_count = defaultdict(int) - _chunk_idx = int(transfer_manager.ramp_chunk_count.get(req_id, 0)) - threshold = ramp_chunk_size(_chunk_idx, _ramp, chunk_frames) - else: - emitted_any = int(transfer_manager.put_req_chunk.get(req_id, 0)) > 0 - threshold = initial_chunk_frames if initial_chunk_frames > 0 and not emitted_any else chunk_frames + emitted_chunks = int(transfer_manager.put_req_chunk.get(req_id, 0)) + emitted_any = emitted_chunks > 0 + threshold = initial_chunk_frames if initial_chunk_frames > 0 and not emitted_any else chunk_frames + # The connector configuration is fixed for this transfer manager's lifetime. + if not hasattr(transfer_manager, "_moss_chunk_ramp"): + transfer_manager._moss_chunk_ramp = parse_chunk_ramp(cfg, steady=chunk_frames) + ramp = transfer_manager._moss_chunk_ramp + if ramp is not None: + # The ladder is indexed by the connector's segment-local counter, which + # restarts at each segment boundary; put_req_chunk is request-global. + ramp_index = int(transfer_manager.ramp_chunk_count.get(req_id, 0)) + threshold = ramp_chunk_size(ramp_index, ramp, chunk_frames) if pending <= 0: if is_finished: transfer_manager.code_prompt_token_ids.pop(req_id, None) transfer_manager.request_payload.pop(req_id, None) return OmniPayloadStruct( + # Explicitly replace the previous chunk (or the prewarmed + # placeholder) with no codes. Metadata alone is not an + # empty audio snapshot in the native connector contract. + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), meta=MetaStruct( req_id=[req_id], left_context_size=0, @@ -325,13 +332,10 @@ def talker2codec_raw_async_chunk( emit_frames = pending if is_finished else threshold chunk_rows = pending_frames[:emit_frames] del pending_frames[:emit_frames] - chunk_codes = torch.stack( - [row.to(torch.long).cpu() for row in chunk_rows], - dim=0, - ).contiguous() + chunk_np = np.stack([np.asarray(row, dtype=np.int64) for row in chunk_rows]) finished = bool(is_finished and len(pending_frames) == 0) - codec_flat = chunk_codes.transpose(0, 1).contiguous().reshape(-1).to(torch.long) + codec_flat = torch.from_numpy(np.ascontiguousarray(chunk_np.T).reshape(-1)) if finished: transfer_manager.code_prompt_token_ids.pop(req_id, None) @@ -342,7 +346,7 @@ def talker2codec_raw_async_chunk( meta=MetaStruct( req_id=[req_id], left_context_size=0, - codec_chunk_frames=int(chunk_codes.shape[0]), + codec_chunk_frames=int(chunk_np.shape[0]), codec_left_context_frames=0, code_flat_numel=int(codec_flat.numel()), stream_finished=torch.tensor(finished, dtype=torch.bool), diff --git a/vllm_omni/worker_v2/first_audio_sender.py b/vllm_omni/worker_v2/first_audio_sender.py index 7e9246b5b..e1d09f367 100644 --- a/vllm_omni/worker_v2/first_audio_sender.py +++ b/vllm_omni/worker_v2/first_audio_sender.py @@ -111,9 +111,10 @@ class _PreparedDelivery: class _EngineOutputSink: - def __init__(self, output_queue: Any, scheduler: Any) -> None: + def __init__(self, output_queue: Any, scheduler: Any, *, upstream_first_audio: bool = True) -> None: self.output_queue = output_queue self.scheduler = scheduler + self.upstream_first_audio = upstream_first_audio def prepare(self, request_ids: list[str]) -> _PreparedDelivery: routes = {} @@ -130,15 +131,14 @@ class _EngineOutputSink: for request_id, pcm in zip(request_ids, pcm_rows, strict=True): if request_id not in routes: continue + payload = {"model_outputs": pcm, "sr": sample_rate} + if self.upstream_first_audio: + payload[FIRST_AUDIO_KEY] = torch.tensor(True) by_client.setdefault(routes[request_id], []).append( OmniEngineCoreOutput( request_id=request_id, new_token_ids=[], - multimodal_output={ - "model_outputs": pcm, - "sr": sample_rate, - FIRST_AUDIO_KEY: torch.tensor(True), - }, + multimodal_output=payload, ) ) for client_index, outputs in by_client.items(): @@ -163,5 +163,8 @@ class _EngineOutputSink: routes.pop(request_id) -def engine_output_queue_sink(output_queue: Any, scheduler: Any) -> _EngineOutputSink: - return _EngineOutputSink(output_queue, scheduler) +def engine_output_queue_sink( + output_queue: Any, scheduler: Any, *, upstream_first_audio: bool = True +) -> _EngineOutputSink: + """Route prepared audio; final-stage PCM must not carry the upstream marker.""" + return _EngineOutputSink(output_queue, scheduler, upstream_first_audio=upstream_first_audio) diff --git a/vllm_omni/worker_v2/model_states/__init__.py b/vllm_omni/worker_v2/model_states/__init__.py index be7e19dbb..7475d3d7d 100644 --- a/vllm_omni/worker_v2/model_states/__init__.py +++ b/vllm_omni/worker_v2/model_states/__init__.py @@ -38,6 +38,12 @@ def init_omni_model_state( getattr(model, flag, False) is True for flag in ("has_preprocess", "has_postprocess", "have_multimodal_outputs") ) if uses_omni_lifecycle: + factory = getattr(type(model), "create_mrv2_model_state", None) + if callable(factory): + state = factory(model, vllm_config, encoder_cache, device) + if not isinstance(state, ModelState): + raise TypeError("create_mrv2_model_state must return a ModelState") + return state from vllm_omni.worker_v2.model_states.omni_model_state import ( OmniModelState, ) diff --git a/vllm_omni/worker_v2/omni_ar_model_runner.py b/vllm_omni/worker_v2/omni_ar_model_runner.py index 73434b8a1..c9ef50236 100644 --- a/vllm_omni/worker_v2/omni_ar_model_runner.py +++ b/vllm_omni/worker_v2/omni_ar_model_runner.py @@ -138,6 +138,23 @@ class OmniARModelRunner(OmniGPUModelRunner): # sample_tokens: OmniOutput handling + pooler_output + async D2H # ------------------------------------------------------------------ + def sample(self, hidden_states, input_batch, grammar_output): + # An explicitly declared model-state hook may return already determined + # tokens. Unsupported sampling features retain the upstream path. + sample_determined = getattr(type(self.model_state), "sample_determined_tokens", None) + if ( + sample_determined is not None + and grammar_output is None + and self.batch_sharder is None + and self.pp_handler is None + and input_batch.num_draft_tokens == 0 + and input_batch.num_reqs > 0 + ): + output = sample_determined(self.model_state, input_batch, self.sampler) + if output is not None: + return output, output.num_sampled, output.num_rejected + return super().sample(hidden_states, input_batch, grammar_output) + @torch.inference_mode() @step_eplb_after() def sample_tokens( @@ -311,6 +328,11 @@ class OmniARModelRunner(OmniGPUModelRunner): # Freshly allocated each step and filled after sampling; nothing # overwrites them before the host copy. return outputs + # A producer-side PackedOutputSnapshot already owns its device slabs + # for this forward and carries the event that makes them visible to + # the output copy stream. Repacking it would add a second GPU copy. + if isinstance(outputs, PackedOutputSnapshot): + return outputs slot_index = self._async_mm_snapshot_cursor if self._async_mm_snapshot_pending[slot_index]: raise RuntimeError("Async multimodal snapshot ring exhausted before sample_tokens().") @@ -567,6 +589,8 @@ def _async_copy_mm( if not mm_outputs: return {} if isinstance(mm_outputs, PackedOutputSnapshot): + if mm_outputs.producer_event is not None: + mm_outputs.producer_event.wait(copy_stream) return mm_outputs.copy_to_cpu( lambda tensor: _async_copy_tensor(tensor, copy_stream=copy_stream, pin_memory=pin_memory) ) diff --git a/vllm_omni/worker_v2/output_snapshot.py b/vllm_omni/worker_v2/output_snapshot.py index a8ae384db..e496a986a 100644 --- a/vllm_omni/worker_v2/output_snapshot.py +++ b/vllm_omni/worker_v2/output_snapshot.py @@ -1,91 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Owned, dtype-grouped snapshots for asynchronous runner outputs.""" +"""Compatibility imports for the shared model/runner output snapshot contract.""" -from __future__ import annotations +from vllm_omni.model_executor.output_snapshot import ( + PackedOutputSnapshot, + RequestOutputSnapshot, + pack_output_snapshot, +) -from collections import defaultdict -from collections.abc import Callable -from dataclasses import dataclass -from typing import Any - -import torch -from torch.utils._pytree import tree_flatten, tree_unflatten - - -@dataclass(frozen=True) -class RequestOutputSnapshot: - """Already partitioned, CPU-owned payloads in the current request order. - - Model finalizers may return this after D2H completes to avoid recursively - cloning and partitioning data whose request ownership is already known. - """ - - inter_stage: list[dict[str, Any] | None] - client: list[dict[str, Any] | None] | None = None - - -class PackedOutputSnapshot(dict): - """A normal payload mapping with an internal batched-copy plan. - - The runner's snapshot-slot event protects the device slabs until D2H is - complete. Host slabs are allocated per output, so downstream consumers can - retain their views without depending on reuse of the device ring slot. - """ - - def __init__(self, leaves: list[Any], spec: Any, groups: list[tuple[list[int], torch.Tensor]]) -> None: - super().__init__(tree_unflatten(leaves, spec)) - self._leaves = leaves - self._spec = spec - self._groups = groups - - def copy_to_cpu(self, copy_tensor: Callable[[torch.Tensor], torch.Tensor]) -> dict[str, Any]: - leaves = list(self._leaves) - for indices, slab in self._groups: - cpu = copy_tensor(slab) - offset = 0 - for index in indices: - value = leaves[index] - size = value.numel() - leaves[index] = cpu.narrow(0, offset, size).view(value.shape) - offset += size - return tree_unflatten(leaves, self._spec) - - -def pack_output_snapshot( - payload: dict[str, Any], slot: dict[tuple[Any, ...], torch.Tensor], *, max_buckets: int -) -> PackedOutputSnapshot | None: - """Copy tensor leaves once per dtype/device, retaining their exact values. - - No CUDA synchronization is introduced here. The caller must select the - producer stream and wait for the previous consumer before reusing a slot. - """ - leaves, spec = tree_flatten(payload) - grouped: dict[tuple[Any, ...], list[int]] = defaultdict(list) - for index, value in enumerate(leaves): - if isinstance(value, torch.Tensor): - # Preserve the existing per-tensor path for unusual layouts. - if value.layout != torch.strided or value.is_quantized: - return None - grouped[(value.dtype, value.device)].append(index) - groups = [] - for (dtype, device), indices in grouped.items(): - values = [leaves[index].detach().reshape(-1) for index in indices] - size = sum(value.numel() for value in values) - key = ("packed-output", dtype, device, size) - slab = slot.get(key) - if slab is None: - slab = torch.empty(size, dtype=dtype, device=device) - if len(slot) < max_buckets: - slot[key] = slab - if len(values) == 1: - slab.copy_(values[0]) - else: - torch.cat(values, out=slab) - offset = 0 - for index, value in zip(indices, values): - shape = leaves[index].shape - leaves[index] = slab.narrow(0, offset, value.numel()).view(shape) - offset += value.numel() - groups.append((indices, slab)) - return PackedOutputSnapshot(leaves, spec, groups) +__all__ = ["PackedOutputSnapshot", "RequestOutputSnapshot", "pack_output_snapshot"] From ee8fdab1dc26ea8f09d7910cdc98a17c4b0ea02a Mon Sep 17 00:00:00 2001 From: andyluo7 <43718156+andyluo7@users.noreply.github.com> Date: Fri, 2 Oct 2026 18:13:51 -0700 Subject: [PATCH 22/22] [CI/Build] Fix MOSS CPU tests on ROCm workers (#8447) Signed-off-by: andyluo7 --- .../sched/test_generation_scheduler_init.py | 15 ++++++--- .../models/moss_tts/test_reference_encoder.py | 32 +++++++++++-------- .../moss_tts/test_shared_reference_encoder.py | 6 ++-- 3 files changed, 32 insertions(+), 21 deletions(-) diff --git a/tests/core/sched/test_generation_scheduler_init.py b/tests/core/sched/test_generation_scheduler_init.py index ade94af17..db0285e38 100644 --- a/tests/core/sched/test_generation_scheduler_init.py +++ b/tests/core/sched/test_generation_scheduler_init.py @@ -8,8 +8,9 @@ import pytest from tests.helpers.stage_config import get_deploy_config_path from vllm_omni.config.pipeline_registry import resolve_pipeline_config -from vllm_omni.config.stage_config import _apply_platform_overrides, load_deploy_config, merge_pipeline_deploy +from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler, VLLMScheduler +from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -45,10 +46,14 @@ def _config(*, native=True, stateful=True, tp=1, pp=1, extras=None, capacity=128 @pytest.mark.parametrize("profile", ["high_concurrency", "low_latency"]) @pytest.mark.parametrize("platform", ["cuda", "npu", "xpu", "rocm", "musa"]) -def test_moss_profile_generation_constructor_after_platform_resolution(construct_scheduler, mocker, profile, platform): - deploy = _apply_platform_overrides( - load_deploy_config(get_deploy_config_path(f"moss_tts_local_mrv2_{profile}.yaml")), platform=platform - ) +def test_moss_profile_generation_constructor_after_platform_resolution( + construct_scheduler, monkeypatch, mocker, profile, platform +): + # merge_pipeline_deploy owns platform resolution. Mock the detected device + # instead of applying an override first and then resolving the worker's + # actual platform a second time. + monkeypatch.setattr(current_omni_platform, "device_name", platform) + deploy = load_deploy_config(get_deploy_config_path(f"moss_tts_local_mrv2_{profile}.yaml")) pipeline = resolve_pipeline_config("moss_tts_local") codec = merge_pipeline_deploy(pipeline, deploy)[1] extras = deploy.connectors["shm"]["extra"] diff --git a/tests/model_executor/models/moss_tts/test_reference_encoder.py b/tests/model_executor/models/moss_tts/test_reference_encoder.py index 06efaaf25..5e03d40ec 100644 --- a/tests/model_executor/models/moss_tts/test_reference_encoder.py +++ b/tests/model_executor/models/moss_tts/test_reference_encoder.py @@ -6,6 +6,7 @@ single-flight, and micro-batched encoding.""" import asyncio import threading import time +from tempfile import TemporaryDirectory import pytest import torch @@ -650,22 +651,25 @@ async def test_inline_index_requires_matching_resolve_key(make_encoder): assert audio.resolve_calls == [ref, ref] -async def test_inline_reference_shared_across_encoders(make_encoder, tmp_path_factory, monkeypatch): - shared_dir = tmp_path_factory.mktemp("ref") - monkeypatch.setenv("VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR", str(shared_dir)) - ref = "data:audio/wav;base64,R0hJ" - first_audio, second_audio = _DigestAudio(), _DigestAudio() - first_audio.register(ref, 9) - second_audio.register(ref, 9) - first, second = make_encoder(_FakeProcessor()), make_encoder(_FakeProcessor()) +async def test_inline_reference_shared_across_encoders(make_encoder, monkeypatch): + # Keep the AF_UNIX socket path independent of pytest's nested base path. + with TemporaryDirectory(prefix="moss-ref-") as shared_dir: + monkeypatch.setenv("VLLM_OMNI_MOSS_REF_CODES_SHARED_DIR", shared_dir) + ref = "data:audio/wav;base64,R0hJ" + first_audio, second_audio = _DigestAudio(), _DigestAudio() + first_audio.register(ref, 9) + second_audio.register(ref, 9) + first, second = make_encoder(_FakeProcessor()), make_encoder(_FakeProcessor()) - a, key_a = await first.encode(ref, resolve_ref_audio=first_audio.resolve, get_artifact_key=first_audio.artifact_key) - b, key_b = await second.encode( - ref, resolve_ref_audio=second_audio.resolve, get_artifact_key=second_audio.artifact_key - ) + a, key_a = await first.encode( + ref, resolve_ref_audio=first_audio.resolve, get_artifact_key=first_audio.artifact_key + ) + b, key_b = await second.encode( + ref, resolve_ref_audio=second_audio.resolve, get_artifact_key=second_audio.artifact_key + ) - assert first_audio.resolve_calls == [ref] and second_audio.resolve_calls == [] - assert key_a == key_b and torch.equal(a, b) + assert first_audio.resolve_calls == [ref] and second_audio.resolve_calls == [] + assert key_a == key_b and torch.equal(a, b) # --------------------------------------------------------------------------- # diff --git a/tests/model_executor/models/moss_tts/test_shared_reference_encoder.py b/tests/model_executor/models/moss_tts/test_shared_reference_encoder.py index b8bd50184..d59fb5f6f 100644 --- a/tests/model_executor/models/moss_tts/test_shared_reference_encoder.py +++ b/tests/model_executor/models/moss_tts/test_shared_reference_encoder.py @@ -3,6 +3,7 @@ """API processes share one reference encoder through a Unix socket.""" import threading +from tempfile import TemporaryDirectory import pytest import torch @@ -30,10 +31,11 @@ class _Recorder: @pytest.fixture -def host_dir(tmp_path_factory, monkeypatch): +def host_dir(monkeypatch): monkeypatch.setattr(shared, "_host_lock_handle", None) # AF_UNIX limits the full socket path, including pytest directory names. - return str(tmp_path_factory.mktemp("ref")) + with TemporaryDirectory(prefix="moss-ref-") as path: + yield path def test_only_one_process_per_directory_hosts(host_dir):