From a31690041de9a19acfec8239c4bbb64bc3e051d4 Mon Sep 17 00:00:00 2001 From: Yancy <138764723+Asthenia0412@users.noreply.github.com> Date: Tue, 8 Sep 2026 21:56:27 +0800 Subject: [PATCH 01/93] [Bugfix][Examples] Use --profiler-config flag in offline TTS examples (#6763) Signed-off-by: Asthenia Co-authored-by: Asthenia --- .../offline_inference/qwen2_5_omni/end2end.py | 33 +++++++++++++++---- .../text_to_speech/cosyvoice3/end2end.py | 30 +++++++++++++++-- 2 files changed, 54 insertions(+), 9 deletions(-) diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index 6dbafdfc1a4..bbeef32b01c 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -5,10 +5,12 @@ with the correct prompt format on Qwen2.5-Omni """ +import argparse +import functools import json import os import time -from typing import NamedTuple +from typing import Any, NamedTuple import numpy as np import soundfile as sf @@ -26,6 +28,20 @@ SEED = 42 +def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]: + """Parse a CLI value as a JSON object, attributing errors to ``flag_name``.""" + try: + config = json.loads(value) + except json.JSONDecodeError as e: + raise argparse.ArgumentTypeError(f"{flag_name} must be valid JSON: {e}") from e + if not isinstance(config, dict): + raise argparse.ArgumentTypeError(f"{flag_name} must be a JSON object") + return config + + +parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config") + + class QueryResult(NamedTuple): inputs: dict limit_mm_per_prompt: dict[str, int] @@ -378,11 +394,10 @@ def main(args): for i, prompt in enumerate(prompts): prompt["modalities"] = output_modalities - profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR")) - if profiler_enabled and hasattr(omni, "start_profile"): + profiler_enabled = args.profiler_config is not None + if profiler_enabled: + print("[Profiler] Starting profiling...") omni.start_profile(stages=[0]) - elif profiler_enabled: - print("[Warn] VLLM_TORCH_PROFILER_DIR is set, but current engine does not support profiler controls.") omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator) # Determine output directory: prefer --output-dir; fallback to --output-wav @@ -418,7 +433,7 @@ def main(args): print(f"Request ID: {request_id}, Saved audio to {output_wav}") processed_count += 1 - if profiler_enabled and hasattr(omni, "stop_profile") and processed_count >= total_requests: + if profiler_enabled and processed_count >= total_requests: print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...") # Stop the profiler while workers are still alive omni.stop_profile() @@ -476,6 +491,12 @@ def parse_args(): default=300, help="Timeout for initializing stages in seconds (default: 300)", ) + parser.add_argument( + "--profiler-config", + type=parse_profiler_config, + default=None, + help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', + ) parser.add_argument( "--shm-threshold-bytes", type=int, diff --git a/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py b/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py index 62917de078c..40e19412059 100644 --- a/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py +++ b/examples/offline_inference/text_to_speech/cosyvoice3/end2end.py @@ -1,9 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse +import functools +import json import os import urllib.request from pathlib import Path +from typing import Any import numpy as np import soundfile as sf @@ -29,6 +32,20 @@ def _default_ref_audio() -> str: return str(dest) +def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]: + """Parse a CLI value as a JSON object, attributing errors to ``flag_name``.""" + try: + config = json.loads(value) + except json.JSONDecodeError as e: + raise argparse.ArgumentTypeError(f"{flag_name} must be valid JSON: {e}") from e + if not isinstance(config, dict): + raise argparse.ArgumentTypeError(f"{flag_name} must be a JSON object") + return config + + +parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config") + + def run_e2e(): parser = argparse.ArgumentParser() # ""FunAudioLLM/Fun-CosyVoice3-0.5B-2512 @@ -64,6 +81,12 @@ def run_e2e(): required=True, help="Path to tokenizer directory (e.g., /CosyVoice-BlankEN).", ) + parser.add_argument( + "--profiler-config", + type=parse_profiler_config, + default=None, + help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', + ) args = parser.parse_args() # Ensure tokenizer directory exists if not os.path.exists(args.tokenizer): @@ -79,6 +102,7 @@ def run_e2e(): deploy_config=args.deploy_config, tokenizer=args.tokenizer, log_stats=True, + profiler_config=args.profiler_config, ) sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1} @@ -149,8 +173,8 @@ def run_e2e(): sampling_params_list = [gpt_sampling, s2mel_sampling] - # Start profiling (requires VLLM_TORCH_PROFILER_DIR env var) - if os.environ.get("VLLM_TORCH_PROFILER_DIR"): + profiler_enabled = args.profiler_config is not None + if profiler_enabled: print("Starting profiler...") omni.start_profile() @@ -158,7 +182,7 @@ def run_e2e(): outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list[:2])) # Stop profiling and get results - if os.environ.get("VLLM_TORCH_PROFILER_DIR"): + if profiler_enabled: print("Stopping profiler...") profile_results = omni.stop_profile() print(f"Profile traces saved to: {profile_results}") From 7422a85cc7a156a9131938ebe3afb6c8025c2a40 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Tue, 8 Sep 2026 22:54:49 +0800 Subject: [PATCH 02/93] [Bugfix] Skip HWR store-size scans when no limit is configured (#7131) Signed-off-by: Hongsheng Liu --- vllm_omni/host_weight_runtime/filesystem/store.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/vllm_omni/host_weight_runtime/filesystem/store.py b/vllm_omni/host_weight_runtime/filesystem/store.py index 0f01e36dac5..2b58b77e874 100644 --- a/vllm_omni/host_weight_runtime/filesystem/store.py +++ b/vllm_omni/host_weight_runtime/filesystem/store.py @@ -1247,8 +1247,10 @@ def _check_capacity(self, projected_artifact_bytes: int, additional_bytes: int) retryable=True, ) ) - store_bytes = self._tree_bytes(self.root) - if policy.max_store_bytes is not None and store_bytes + additional_bytes > policy.max_store_bytes: + if ( + policy.max_store_bytes is not None + and self._tree_bytes(self.root) + additional_bytes > policy.max_store_bytes + ): raise HostWeightError( _failure( ResolutionStage.CAPACITY, From de1b0708bdb912397ffa016bbbca4ccd42f48592 Mon Sep 17 00:00:00 2001 From: andyluo7 <43718156+andyluo7@users.noreply.github.com> Date: Tue, 8 Sep 2026 07:59:37 -0700 Subject: [PATCH 03/93] [CI][ROCm] Route LTX2 Ulysses parity to two-GPU lane (#7234) Signed-off-by: andyluo7 --- .buildkite/amd/test-amd-ready.yml | 13 ++++++++++++- .../models/ltx2/test_ltx2_transformer_ulysses.py | 6 +++++- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/.buildkite/amd/test-amd-ready.yml b/.buildkite/amd/test-amd-ready.yml index 695f47dc1a5..585f089a271 100644 --- a/.buildkite/amd/test-amd-ready.yml +++ b/.buildkite/amd/test-amd-ready.yml @@ -107,6 +107,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/models/ltx2/test_ltx2_transformer_ulysses.py + -m 'core_model and rocm and cards_2' + --run-level "core_model" - label: "Diffusion GPU Worker Test" agent_pool: mi300_2 @@ -122,7 +127,13 @@ steps: mirror_hardwares: [amdproduction] grade: Blocking commands: - - timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model" + # Preserve legacy one-GPU tests without cards_1 while excluding every + # registered multi-GPU marker from this one-GPU job. + - >- + timeout 15m pytest -sv tests/diffusion/models/ + -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" - label: "Diffusion · Wan22 Test" agent_pool: mi300_1 diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer_ulysses.py b/tests/diffusion/models/ltx2/test_ltx2_transformer_ulysses.py index 9f09b72a76b..0e080e015e5 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_transformer_ulysses.py +++ b/tests/diffusion/models/ltx2/test_ltx2_transformer_ulysses.py @@ -156,10 +156,14 @@ def _run_transformer_parity(rank: int, world_size: int, master_port: int) -> Non @pytest.mark.full_model -@hardware_test(res={"cuda": "L4"}, num_cards=2) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2) def test_ltx_ulysses_transformer_matches_sp1(unused_tcp_port) -> None: """Exercise real plan hooks, RoPE/timestep sharding, and attention routing.""" world_size = 2 + device_count = current_omni_platform.device_count() + assert device_count >= world_size, ( + f"LTX Ulysses parity requires {world_size} accelerator devices; only {device_count} are visible" + ) torch.multiprocessing.spawn( _run_transformer_parity, args=(world_size, unused_tcp_port), From 7a80b796986b15840bc30ef134f0d7158dce82ea Mon Sep 17 00:00:00 2001 From: liangmenghuang Date: Tue, 8 Sep 2026 23:53:20 +0800 Subject: [PATCH 04/93] [Bugfix][Model] GR00T-N1.7: honor the per-request seed for flow-matching noise (#7253) Signed-off-by: liangmengh Co-authored-by: Claude Fable 5.1 --- docs/serving/openpi_api.md | 8 +- tests/config/test_environment_variables.py | 4 +- .../diffusion/models/gr00t/test_noise_seed.py | 177 ++++++++++++++++++ tests/diffusion/models/gr00t/test_pipeline.py | 11 +- .../test_gr00t_openpi_expansion.py | 9 +- .../openai_api/test_openpi_serving.py | 15 ++ tests/helpers/client.py | 12 +- .../config/environment_variable_inventory.py | 1 - .../models/gr00t/modeling/gr00t_n1d7.py | 29 ++- .../diffusion/models/gr00t/pipeline_gr00t.py | 39 +++- vllm_omni/diffusion/models/gr00t/policy.py | 7 +- vllm_omni/entrypoints/openpi/serving.py | 14 +- 12 files changed, 285 insertions(+), 41 deletions(-) create mode 100644 tests/diffusion/models/gr00t/test_noise_seed.py diff --git a/docs/serving/openpi_api.md b/docs/serving/openpi_api.md index 27f092d87f6..52d94cbbf0d 100644 --- a/docs/serving/openpi_api.md +++ b/docs/serving/openpi_api.md @@ -64,7 +64,7 @@ connect <- msgpack(policy_server_config) infer - -> msgpack({"endpoint": "infer", "session_id": "...", ...observation}) + -> msgpack({"endpoint": "infer", "session_id": "...", "seed": , ...observation}) <- msgpack(ndarray | dict[str, ndarray]) reset @@ -92,6 +92,12 @@ frames are not observation messages. - The policy pipeline owns observation transforms and persistent model state, normally keyed by `session_id`; the API layer forwards the raw observation dictionary. +- An optional integer `seed` in an inference message becomes the engine + request's `sampling_params.seed`. A policy whose sampling consumes the + request generator (GR00T-N1.7) returns the same action chunk for the same + observation and seed. If omitted, the engine assigns a random per-request + seed, so unseeded inferences are independent of each other. The API layer + consumes `seed`; it is not forwarded inside the observation. ## Limits and Errors diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index 52ff1e27489..7603945a0bf 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: 24, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, - EnvironmentVariableCategory.MODEL_SPECIFIC: 62, + EnvironmentVariableCategory.MODEL_SPECIFIC: 61, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20, EnvironmentVariableCategory.INTERNAL: 2, } @@ -179,7 +179,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): ) assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == { ModelEnvironmentVariableDisposition.PROMOTE: 37, - ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 7, + ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6, ModelEnvironmentVariableDisposition.EXTERNAL: 0, ModelEnvironmentVariableDisposition.INTERNALIZE: 13, ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 5, diff --git a/tests/diffusion/models/gr00t/test_noise_seed.py b/tests/diffusion/models/gr00t/test_noise_seed.py new file mode 100644 index 00000000000..a865d9a05a2 --- /dev/null +++ b/tests/diffusion/models/gr00t/test_noise_seed.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The request's seed must reach GR00T's flow-matching noise. + +The runner builds ``sampling_params.generator`` from ``sampling_params.seed``, the pipeline hands +that generator to the policy, and the action head draws its initial noise from it, so the same +seed reproduces the same action chunk and different seeds give independent ones. These tests pin +both ends of that plumbing. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from torch import nn +from transformers.feature_extraction_utils import BatchFeature + +from vllm_omni.diffusion.models.gr00t import pipeline_gr00t +from vllm_omni.diffusion.models.gr00t.configs.gr00t_n1d7 import Gr00tN1d7Config +from vllm_omni.diffusion.models.gr00t.modeling.gr00t_n1d7 import Gr00tN1d7ActionHead +from vllm_omni.diffusion.models.gr00t.pipeline_gr00t import Gr00tN1d7Pipeline +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +ROBOT_OBS = { + "images": {"cam": np.zeros((1, 1, 8, 8, 3), dtype=np.uint8)}, + "state": {"joint": np.zeros((1, 1, 2), dtype=np.float32)}, + "prompt": "pick the cube", +} + + +class _OptionsRecordingPolicy: + """Stand-in for ``Gr00tPolicy``; the pipeline only needs ``language_key`` and ``get_action``.""" + + def __init__(self, *, model_path: str, embodiment_tag: str, device: str, strict: bool) -> None: + del model_path, embodiment_tag, device, strict # the pipeline's constructor contract; unused here + self.language_key = "annotation.language.language_instruction" + self.seen_options: dict[str, object] | None = None + + def get_action(self, obs: dict[str, object], options: dict[str, object] | None = None): + self.seen_options = options + return {"arm": np.zeros((1, 2, 2), dtype=np.float32)}, {} + + +@pytest.fixture +def pipeline(monkeypatch) -> Gr00tN1d7Pipeline: + monkeypatch.setattr(pipeline_gr00t, "Gr00tPolicy", _OptionsRecordingPolicy) + od_config = SimpleNamespace( + model="nvidia/GR00T-N1.7-3B", + model_config={"embodiment_tag": "LIBERO_PANDA", "strict": False}, + custom_pipeline_args={}, + ) + return Gr00tN1d7Pipeline(od_config=od_config) + + +def _request(**sampling_kwargs) -> OmniDiffusionRequest: + return OmniDiffusionRequest( + prompt="pick", + request_id="req", + sampling_params=OmniDiffusionSamplingParams(extra_args={"robot_obs": ROBOT_OBS}, **sampling_kwargs), + ) + + +def test_runner_generator_is_passed_through_untouched(pipeline): + """In serving the runner has already built ``sampling_params.generator``; the head must get that object.""" + generator = torch.Generator().manual_seed(99) + + assert pipeline.forward(_request(generator=generator)).error is None + + assert pipeline.policy.seen_options["generator"] is generator + + +def test_seed_without_generator_builds_one_on_the_pipeline_device(pipeline): + """Callers that bypass the runner (direct use, tests) still get the seed honoured.""" + pipeline.forward(_request(seed=1234)) + + generator = pipeline.policy.seen_options["generator"] + assert isinstance(generator, torch.Generator) + assert generator.initial_seed() == 1234 + assert generator.device.type == pipeline.device + + +def test_auto_assigned_seed_reaches_the_head(pipeline): + """``OmniDiffusionRequest`` seeds every request, so the head is never left on the global RNG.""" + req = _request() + assert req.sampling_params.seed is not None + + pipeline.forward(req) + + assert pipeline.policy.seen_options["generator"].initial_seed() == req.sampling_params.seed + + +def _captured_warnings(monkeypatch) -> list[tuple]: + """vLLM loggers do not propagate to the root logger, so ``caplog`` cannot see them; record calls instead.""" + captured: list[tuple] = [] + monkeypatch.setattr(pipeline_gr00t.logger, "warning", lambda *args, **kwargs: captured.append(args)) + return captured + + +def test_unexpected_generator_form_warns_and_falls_back_to_the_seed(pipeline, monkeypatch): + """GR00T draws one noise tensor per request, so a list of several generators cannot be honoured.""" + warnings = _captured_warnings(monkeypatch) + generators = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)] + + pipeline.forward(_request(generator=generators, seed=7)) + + assert len(warnings) == 1 + assert "torch.Generator" in warnings[0][0] + assert pipeline.policy.seen_options["generator"].initial_seed() == 7 + + +def test_unexpected_generator_form_without_a_seed_warns_and_leaves_the_global_rng(pipeline, monkeypatch): + warnings = _captured_warnings(monkeypatch) + generators = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)] + + pipeline.forward(_request(generator=generators)) + + assert len(warnings) == 1 + assert pipeline.policy.seen_options is None + + +def _noise_only_head(action_horizon: int, action_dim: int) -> Gr00tN1d7ActionHead: + """An action head whose single denoising step is the identity, so ``action_pred`` is the initial noise.""" + config = Gr00tN1d7Config( + max_action_dim=action_dim, + action_horizon=action_horizon, + num_inference_timesteps=1, + add_pos_embed=False, + use_alternate_vl_dit=False, + ) + head = Gr00tN1d7ActionHead.__new__(Gr00tN1d7ActionHead) + nn.Module.__init__(head) + head.config = config + head.action_dim = config.max_action_dim + head.action_horizon = config.action_horizon + head.num_inference_timesteps = config.num_inference_timesteps + head.num_timestep_buckets = config.num_timestep_buckets + head.action_encoder = lambda actions, timesteps, embodiment_id: actions + head.model = lambda hidden_states, encoder_hidden_states, timestep: hidden_states + head.action_decoder = lambda hidden_states, embodiment_id: torch.zeros_like(hidden_states) + return head + + +def _predict(head: Gr00tN1d7ActionHead, options) -> torch.Tensor: + batch, seq, dim = 2, 3, head.action_dim + return head.get_action_with_features( + backbone_features=torch.zeros(batch, seq, dim), + state_features=torch.zeros(batch, 1, dim), + embodiment_id=torch.zeros(batch, dtype=torch.long), + backbone_output=BatchFeature(data={}), + action_input=BatchFeature(data={}), + options=options, + )["action_pred"] + + +def test_head_draws_its_noise_from_the_request_generator(): + head = _noise_only_head(action_horizon=4, action_dim=3) + expected = torch.randn(2, 4, 3, generator=torch.Generator().manual_seed(5)) + + assert torch.equal(_predict(head, {"generator": torch.Generator().manual_seed(5)}), expected) + assert not torch.equal(_predict(head, {"generator": torch.Generator().manual_seed(6)}), expected) + + +def test_head_without_generator_stays_on_the_global_rng(): + """Upstream Isaac-GR00T parity: no generator means a bare ``torch.randn`` on the global RNG.""" + head = _noise_only_head(action_horizon=4, action_dim=3) + + torch.manual_seed(5) + expected = torch.randn(2, 4, 3) + torch.manual_seed(5) + + assert torch.equal(_predict(head, None), expected) diff --git a/tests/diffusion/models/gr00t/test_pipeline.py b/tests/diffusion/models/gr00t/test_pipeline.py index 3221121074f..3310d049ec1 100644 --- a/tests/diffusion/models/gr00t/test_pipeline.py +++ b/tests/diffusion/models/gr00t/test_pipeline.py @@ -1,4 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from types import SimpleNamespace +from typing import Any import numpy as np import pytest @@ -12,12 +16,12 @@ class FakeGr00tPolicy: - instances = [] + instances: list["FakeGr00tPolicy"] = [] def __init__(self, **kwargs): self.kwargs = kwargs self.reset_calls = 0 - self.seen_obs = None + self.seen_obs: dict[str, Any] | None = None self.embodiment_tag = SimpleNamespace(value="fake_embodiment") self.language_key = "annotation.language.language_instruction" self.modality_configs = { @@ -40,7 +44,7 @@ def __init__(self, **kwargs): ) FakeGr00tPolicy.instances.append(self) - def get_action(self, obs): + def get_action(self, obs, options=None): self.seen_obs = obs return { "arm": np.array([[[1.0, 2.0]]], dtype=np.float64), @@ -107,6 +111,7 @@ def test_forward_returns_dict_actions_in_output(): assert actions["arm"].dtype == np.float32 np.testing.assert_allclose(actions["arm"], np.array([[[1.0, 2.0]]], dtype=np.float32)) policy = FakeGr00tPolicy.instances[0] + assert policy.seen_obs is not None assert "video" in policy.seen_obs assert policy.seen_obs["language"] == {"annotation.language.language_instruction": [["pick the cube"]]} assert "images" not in policy.seen_obs diff --git a/tests/e2e/online_serving/test_gr00t_openpi_expansion.py b/tests/e2e/online_serving/test_gr00t_openpi_expansion.py index 70e11088d6a..2f81ab3fa6d 100644 --- a/tests/e2e/online_serving/test_gr00t_openpi_expansion.py +++ b/tests/e2e/online_serving/test_gr00t_openpi_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """End-to-end online serving test for GR00T N1.7 through the OpenPI robot endpoint.""" import os @@ -29,8 +29,7 @@ OmniServerParams( model=MODEL, stage_config_path=get_deploy_config_path("Gr00tN1d7.yaml"), - server_args=["--disable-log-stats"], - env_dict={"VLLM_DISABLE_COMPILE_CACHE": "1", "GR00T_NOISE_SEED": "42"}, + env_dict={"VLLM_DISABLE_COMPILE_CACHE": "1"}, init_timeout=1200, stage_init_timeout=900, ), @@ -162,14 +161,14 @@ def test_gr00t_n1d7_openpi_online(omni_server, openai_client) -> None: @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_gr00t_n1d7_openpi_precision(omni_server, openai_client) -> None: - """Assert actions match Isaac-GR00T reference (GR00T_NOISE_SEED=42, zero inputs).""" + """Assert actions match Isaac-GR00T reference (request seed=42, zero inputs).""" response = openai_client.send_robot_openpi_ws_request( { "operations": [ {"endpoint": "reset", "payload": {}}, { "endpoint": "infer", - "payload": build_openpi_droid_observation(session_id="gr00t-precision-e2e"), + "payload": build_openpi_droid_observation(session_id="gr00t-precision-e2e", seed=42), }, ], } diff --git a/tests/entrypoints/openai_api/test_openpi_serving.py b/tests/entrypoints/openai_api/test_openpi_serving.py index cdd833db083..072e734e48f 100644 --- a/tests/entrypoints/openai_api/test_openpi_serving.py +++ b/tests/entrypoints/openai_api/test_openpi_serving.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import asyncio import json import threading @@ -218,6 +221,18 @@ def test_build_request_uses_unique_engine_request_id_per_inference(): assert request_a.request_id != request_b.request_id +def test_build_request_forwards_seed_to_sampling_params(): + """``seed`` in the inference message is the engine-level seed; omitted, the request auto-seeds.""" + serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=_engine_with_policy_config()) + + seeded = serving._build_request({"prompt": "pick up the object", "seed": 42}, session_id="s", reset=True) + unseeded = serving._build_request({"prompt": "pick up the object"}, session_id="s", reset=False) + + assert seeded.sampling_params.seed == 42 + assert "seed" not in seeded.sampling_params.extra_args["robot_obs"] + assert isinstance(unseeded.sampling_params.seed, int) + + def test_infer_keeps_session_state_but_uses_unique_engine_request_ids(): engine = RecordingEngine() serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=engine) diff --git a/tests/helpers/client.py b/tests/helpers/client.py index 0f834620073..5339f7195b9 100644 --- a/tests/helpers/client.py +++ b/tests/helpers/client.py @@ -152,11 +152,14 @@ class OpenPIWebSocketResponse: action_tensors: list[np.ndarray] | None = None -def build_openpi_droid_observation(*, session_id: str = "gr00t-smoke") -> dict[str, Any]: - """Build a minimal DROID-style observation payload for the OpenPI robot endpoint.""" +def build_openpi_droid_observation(*, session_id: str = "gr00t-smoke", seed: int | None = None) -> dict[str, Any]: + """Build a minimal DROID-style observation payload for the OpenPI robot endpoint. + + ``seed`` pins the policy's sampling noise for this request (see ``docs/serving/openpi_api.md``). + """ identity_eef_9d = np.zeros((1, 1, 9), dtype=np.float32) identity_eef_9d[..., 3:] = np.array([1, 0, 0, 0, 1, 0], dtype=np.float32) - return { + observation: dict[str, Any] = { "session_id": session_id, "video": { "exterior_image_1_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8), @@ -169,6 +172,9 @@ def build_openpi_droid_observation(*, session_id: str = "gr00t-smoke") -> dict[s }, "language": {"annotation.language.language_instruction": [["pick up the object"]]}, } + if seed is not None: + observation["seed"] = seed + return observation DREAMZERO_DEFAULT_PROMPT = ( diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index 1a1d4014177..c1e380a22ae 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -175,7 +175,6 @@ def is_public_omni(self) -> bool: _MODEL_REQUEST_SCOPE = ( "DOTS_TTS_DIT_NUM_STEPS", - "GR00T_NOISE_SEED", "MAGI2_NEGATIVE_PROMPT", "STEP_AUDIO2_DEFAULT_PROMPT_WAV", "VLLM_GEPARD_END_FADE_MS", diff --git a/vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py b/vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py index 3fcbce203f5..6116572ee96 100644 --- a/vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py +++ b/vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py @@ -1,8 +1,7 @@ # 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 -import os from typing import Any import torch @@ -138,26 +137,22 @@ def get_action_with_features( state_features: [B, state_horizon, input_embedding_dim] embodiment_id: [B] (embodiment IDs) backbone_output: Output from the backbone model + options: ``generator`` (a ``torch.Generator`` on the model device) draws + the initial action noise; ``None`` leaves ``torch.randn`` on the + global RNG, as upstream Isaac-GR00T does. The ``rtc_*`` keys drive + real-time chunking when ``action_input`` carries a previous chunk. """ vl_embeds = backbone_features batch_size = vl_embeds.shape[0] device = vl_embeds.device - _seed_env = os.environ.get("GR00T_NOISE_SEED") - if _seed_env is not None: - _gen = torch.Generator(device=device).manual_seed(int(_seed_env)) - actions = torch.randn( - size=(batch_size, self.config.action_horizon, self.action_dim), - dtype=vl_embeds.dtype, - device=device, - generator=_gen, - ) - else: - actions = torch.randn( - size=(batch_size, self.config.action_horizon, self.action_dim), - dtype=vl_embeds.dtype, - device=device, - ) + generator = options.get("generator") if options else None + actions = torch.randn( + size=(batch_size, self.config.action_horizon, self.action_dim), + dtype=vl_embeds.dtype, + device=device, + generator=generator, + ) dt = 1.0 / self.num_inference_timesteps vel_strength = torch.ones_like(actions) diff --git a/vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py b/vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py index 1154dcf7145..87108522b61 100644 --- a/vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py +++ b/vllm_omni/diffusion/models/gr00t/pipeline_gr00t.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 collections.abc import Iterable, Mapping from typing import Any @@ -13,6 +13,7 @@ from vllm_omni.diffusion.models.gr00t.policy import Gr00tPolicy from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.inputs.data import OmniDiffusionSamplingParams logger = init_logger(__name__) @@ -30,6 +31,10 @@ class Gr00tN1d7Pipeline(nn.Module): vLLM-Omni owns the serving integration: OpenPI observations arrive through `sampling_params.extra_args["robot_obs"]`, this pipeline runs GR00T policy inference, and actions are returned through `DiffusionOutput.output["actions"]`. + + The request's `seed` (materialised by the runner into `sampling_params.generator`) + draws the flow-matching noise, so repeating a request with the same seed reproduces + its action chunk. """ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None: @@ -123,7 +128,8 @@ def forward(self, req: DiffusionRequestBatch, **kwargs) -> DiffusionOutput: self.reset() policy_obs = _normalize_observation(robot_obs, language_key=self.policy.language_key) - result = self.policy.get_action(policy_obs) + generator = _noise_generator(req.sampling_params, self.device) + result = self.policy.get_action(policy_obs, {"generator": generator} if generator is not None else None) actions = result[0] if isinstance(result, tuple) else result if not isinstance(actions, Mapping): return DiffusionOutput(error=f"GR00T policy returned {type(actions).__name__}; expected dict actions.") @@ -132,6 +138,35 @@ def forward(self, req: DiffusionRequestBatch, **kwargs) -> DiffusionOutput: return DiffusionOutput(output={"actions": _to_float32_action_dict(actions)}) +def _noise_generator(sampling_params: OmniDiffusionSamplingParams, device: str) -> torch.Generator | None: + """Pick the generator that draws the action head's initial noise. + + The runner materialises ``sampling_params.seed`` into ``sampling_params.generator`` on + the model device before every forward, and ``OmniDiffusionRequest`` assigns a random + seed when the caller omits one, so in serving the generator is always present. The + seed fallback serves callers that bypass the runner (direct use, unit tests). ``None`` + leaves the action head on the global RNG, as upstream Isaac-GR00T runs. GR00T draws one + noise tensor per request, so a generator in any other form (for example a list of several + generators) is ignored with a warning and the seed fallback applies. + """ + generator = sampling_params.generator + if isinstance(generator, list) and len(generator) == 1: + generator = generator[0] + if isinstance(generator, torch.Generator): + return generator + if generator is not None: + received = f"a list of {len(generator)} generators" if isinstance(generator, list) else type(generator).__name__ + logger.warning( + "GR00T expects a single torch.Generator per request but got %s; ignoring it. The initial noise " + "falls back to sampling_params.seed=%s, or to the global RNG when that is None.", + received, + sampling_params.seed, + ) + if sampling_params.seed is not None: + return torch.Generator(device=device).manual_seed(int(sampling_params.seed)) + return None + + def _normalize_observation(robot_obs: Mapping[str, Any], *, language_key: str) -> dict[str, Any]: obs: dict[str, Any] = {} if "video" in robot_obs: diff --git a/vllm_omni/diffusion/models/gr00t/policy.py b/vllm_omni/diffusion/models/gr00t/policy.py index c16b344632f..118f57545b8 100644 --- a/vllm_omni/diffusion/models/gr00t/policy.py +++ b/vllm_omni/diffusion/models/gr00t/policy.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 pathlib import Path from typing import Any @@ -328,7 +328,8 @@ def _get_action( Args: observation: Batched observation dictionary - options: Optional parameters (currently unused) + options: Optional per-request model parameters, forwarded verbatim to + the action head (e.g. ``generator``) Returns: Tuple of (actions_dict, info_dict) @@ -351,7 +352,7 @@ def _get_action( # Step 4: Run model inference to predict actions with torch.inference_mode(): - model_pred = self.model.get_action(**collated_inputs) + model_pred = self.model.get_action(**collated_inputs, options=options) normalized_action = model_pred["action_pred"].float() # Step 5: Decode actions from normalized space back to physical units diff --git a/vllm_omni/entrypoints/openpi/serving.py b/vllm_omni/entrypoints/openpi/serving.py index 4f1d5808144..f2448dfa19d 100644 --- a/vllm_omni/entrypoints/openpi/serving.py +++ b/vllm_omni/entrypoints/openpi/serving.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 """Serving layer for robot policy inference via `/v1/realtime/robot/openpi`. @@ -12,7 +12,7 @@ from collections.abc import Mapping from dataclasses import dataclass from itertools import count -from typing import Any +from typing import Any, TypeAlias import numpy as np from omegaconf import OmegaConf @@ -20,7 +20,7 @@ logger = init_logger(__name__) -ActionOutput = np.ndarray | dict[str, np.ndarray] +ActionOutput: TypeAlias = np.ndarray | dict[str, np.ndarray] def _to_builtin_container(value: Any) -> Any: @@ -150,6 +150,9 @@ def _build_request(self, obs: dict, *, session_id: str, reset: bool) -> Any: from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniDiffusionSamplingParams + # An optional integer ``seed`` in the inference message becomes the engine request's + # ``sampling_params.seed``; omitted, ``OmniDiffusionRequest`` assigns a random one. + seed = obs.pop("seed", None) extra_args = { "reset": reset, "session_id": session_id, @@ -157,7 +160,10 @@ def _build_request(self, obs: dict, *, session_id: str, reset: bool) -> Any: } prompt = obs.get("prompt", "") - sampling_params = OmniDiffusionSamplingParams(extra_args=extra_args) + sampling_params = OmniDiffusionSamplingParams( + seed=int(seed) if seed is not None else None, + extra_args=extra_args, + ) return OmniDiffusionRequest( prompt=prompt, sampling_params=sampling_params, From fdaf4e31bd6638ddb472f93764aabc034c20bb6d Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:08:38 +0100 Subject: [PATCH 05/93] Add vLLM-Omni library info to Hugging Face Hub requests (#5381) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Hongsheng Liu Co-authored-by: Claude Opus 5 (1M context) --- tests/benchmarks/test_accuracy_bench_utils.py | 12 ++--- .../perf/scripts/run_diffusion_benchmark.py | 13 ++--- .../model_loader/test_diffusers_loader.py | 4 +- .../model_loader/test_hub_prefetch.py | 6 +-- .../cosmos3/test_cosmos3_sound_tokenizer.py | 7 ++- .../models/ltx2/test_ltx2_pipeline.py | 13 +++-- .../test_minimax_h3_quantization_quality.py | 8 +-- .../ltx/test_ltx_official_similarity.py | 10 ++-- .../features/rlhf_test/test_verl_omni_e2e.py | 6 +-- .../test_cosyvoice3_expansion.py | 6 +-- .../test_mammoth_moda2_expansion.py | 7 ++- .../test_moss_tts_realtime.py | 7 +-- .../helpers/minicpmo_4_5_duplex.py | 4 +- tests/engine/test_arg_utils.py | 14 +++-- tests/helpers/__init__.py | 5 +- .../minicpmo_4_5/test_code2wav_model_dir.py | 10 ++-- .../omnivoice/test_nonverbal_tags_tokenize.py | 7 ++- tools/pre_commit/check_forbidden_imports.py | 54 ++++--------------- .../data_modules/daily_omni_dataset.py | 8 +-- .../data_modules/seed_tts_dataset.py | 7 ++- .../benchmarks/data_modules/seed_tts_eval.py | 8 ++- vllm_omni/diffusion/lora/loader.py | 4 +- .../diffusion/model_loader/hub_prefetch.py | 6 +-- .../models/cosmos3/sound_tokenizer.py | 6 +-- .../models/dreamzero/pipeline_dreamzero.py | 4 +- .../models/helios/pipeline_helios.py | 4 +- .../pipeline_longcat_video_avatar.py | 6 +-- .../diffusion/models/ltx2/ltx2_components.py | 10 ++-- .../models/omnivoice/pipeline_omnivoice.py | 6 +-- .../diffusion/models/pi0/pipeline_pi0.py | 6 +-- .../sensenova_u1/pipeline_sensenova_u1.py | 5 +- vllm_omni/diffusion/models/utils.py | 4 +- .../models/wan2_2/pipeline_wan2_2.py | 8 +-- vllm_omni/engine/arg_utils.py | 4 +- vllm_omni/engine/stage_init_utils.py | 2 +- vllm_omni/entrypoints/openai/api_server.py | 4 +- .../openai/tts_adapters/cosyvoice3.py | 4 +- .../model_loader/weight_utils.py | 9 +++- .../model_executor/models/audex/checkpoint.py | 20 +++---- .../model_executor/models/bagel/bagel.py | 7 ++- .../models/cosyvoice3/cosyvoice3.py | 8 +-- .../models/covo_audio/covo_audio_code2wav.py | 7 ++- .../models/dots_tts/dots_tts_talker.py | 6 +-- .../models/dynin_omni/dynin_omni_common.py | 17 +++--- .../dynin_omni/dynin_omni_token2audio.py | 9 ++-- .../model_executor/models/glm_tts/glm_tts.py | 10 ++-- .../higgs_audio_v2/higgs_audio_v2_code2wav.py | 6 +-- .../higgs_audio_v2_tokenizer.py | 6 +-- .../models/indextts2/preprocess_utils.py | 13 +++-- .../models/indextts2/s2mel/modules/bigvgan.py | 8 ++- .../models/indextts2/tokenizer_v2_5.py | 6 +-- .../minicpmo_4_5/minicpmo_4_5_code2wav.py | 4 +- .../models/omnivoice/omnivoice.py | 6 +-- .../models/personaplex/personaplex_mimi.py | 7 +-- .../step_audio2/step_audio2_token2wav.py | 4 +- .../models/voxtral_tts/voxtral_tts.py | 7 ++- vllm_omni/transformers_utils/repo_utils.py | 20 +++++++ 57 files changed, 259 insertions(+), 220 deletions(-) create mode 100644 vllm_omni/transformers_utils/repo_utils.py diff --git a/tests/benchmarks/test_accuracy_bench_utils.py b/tests/benchmarks/test_accuracy_bench_utils.py index ab539234616..8d78da2b2da 100644 --- a/tests/benchmarks/test_accuracy_bench_utils.py +++ b/tests/benchmarks/test_accuracy_bench_utils.py @@ -1,9 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # ruff: noqa: E402, I001 import argparse import math import os import sys -import types from pathlib import Path import pytest @@ -106,11 +108,9 @@ def fake_snapshot_download(*, repo_id, repo_type, allow_patterns): captured["allow_patterns"] = allow_patterns return str(downloaded_root) - monkeypatch.setitem( - sys.modules, - "huggingface_hub", - types.SimpleNamespace(snapshot_download=fake_snapshot_download), - ) + from vllm_omni.transformers_utils import repo_utils + + monkeypatch.setattr(repo_utils.hf_api(), "snapshot_download", fake_snapshot_download) resolved = resolve_seed_tts_root( "zhaochenyang20/seed-tts-eval", diff --git a/tests/dfx/perf/scripts/run_diffusion_benchmark.py b/tests/dfx/perf/scripts/run_diffusion_benchmark.py index bb69391fdf1..ffcfc45b721 100644 --- a/tests/dfx/perf/scripts/run_diffusion_benchmark.py +++ b/tests/dfx/perf/scripts/run_diffusion_benchmark.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Performance benchmark CI runner for diffusion models. @@ -356,6 +359,8 @@ def _resolve_offline_model(model: str) -> str: """ import huggingface_hub + from vllm_omni.transformers_utils.repo_utils import hf_api + if not model or os.path.isdir(model): return model @@ -380,9 +385,7 @@ def _resolve_offline_model(model: str) -> str: if len(parts) >= 3: repo_id = "/".join(parts[:2]) subfolder = "/".join(parts[2:]) - from huggingface_hub import snapshot_download - - snapshot_root = snapshot_download( + snapshot_root = hf_api().snapshot_download( repo_id, allow_patterns=[f"{subfolder}/**"], local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, @@ -391,9 +394,7 @@ def _resolve_offline_model(model: str) -> str: if not huggingface_hub.constants.HF_HUB_OFFLINE: return model - from huggingface_hub import snapshot_download - - return snapshot_download(model, local_files_only=True) + return hf_api().snapshot_download(model, local_files_only=True) class DiffusionServer: diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 555cb19bb4d..565949ff616 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -13,7 +13,6 @@ import pytest import torch import torch.nn as nn -from huggingface_hub import snapshot_download from safetensors.torch import save_file from vllm.config.load import LoadConfig @@ -32,6 +31,7 @@ from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract from vllm_omni.diffusion.registry import initialize_model from vllm_omni.quantization.component_config import ComponentQuantizationConfig +from vllm_omni.transformers_utils.repo_utils import hf_api pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] @@ -41,7 +41,7 @@ @pytest.fixture(scope="module") def prefetch_helios_model(): """Downloads the tiny helios model prior to running a test.""" - snapshot_download(model_path) + hf_api().snapshot_download(model_path) @pytest.fixture(scope="function") diff --git a/tests/diffusion/model_loader/test_hub_prefetch.py b/tests/diffusion/model_loader/test_hub_prefetch.py index 39eb2b27711..94170edcc3f 100644 --- a/tests/diffusion/model_loader/test_hub_prefetch.py +++ b/tests/diffusion/model_loader/test_hub_prefetch.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 contextlib @@ -14,10 +14,10 @@ def test_prefetch_subfolders_propagates_revision(monkeypatch): calls = [] - def fake_snapshot_download(**kwargs): + def fake_snapshot_download(self, **kwargs): calls.append(kwargs) - monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download) + monkeypatch.setattr(huggingface_hub.HfApi, "snapshot_download", fake_snapshot_download) monkeypatch.setattr(hub_prefetch, "_repo_prefetch_lock", lambda _model: contextlib.nullcontext()) hub_prefetch.prefetch_subfolders( diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_sound_tokenizer.py b/tests/diffusion/models/cosmos3/test_cosmos3_sound_tokenizer.py index 7ab04cc212f..1aa653d4a37 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_sound_tokenizer.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_sound_tokenizer.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 __future__ import annotations @@ -81,9 +81,8 @@ def test_from_config_loads_local_diffusers_component(tmp_path, monkeypatch: pyte def test_from_config_downloads_component_from_hf_repo(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None: - import huggingface_hub - from vllm_omni.diffusion.models.cosmos3 import sound_tokenizer + from vllm_omni.transformers_utils import repo_utils cache_dir = tmp_path / "hf" _write_component(cache_dir, checkpoint_name=DIFFUSERS_SOUND_TOKENIZER_CHECKPOINT_NAME) @@ -95,7 +94,7 @@ def fake_snapshot_download(repo_id: str, *, revision: str | None, allow_patterns calls.append((repo_id, revision, allow_patterns)) return str(cache_dir) - monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download) + monkeypatch.setattr(repo_utils.hf_api(), "snapshot_download", fake_snapshot_download) sound_tokenizer.Cosmos3SoundTokenizer.from_config( SimpleNamespace( diff --git a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py index 3e630d8ddd6..a7f6b95dfb8 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py @@ -10,6 +10,7 @@ from types import SimpleNamespace from typing import Any +import huggingface_hub import numpy as np import pytest import torch @@ -302,11 +303,11 @@ def test_ltx_artifact_uses_source_revision_and_hub_fallback( filename = "ltx-sidecar.safetensors" calls = [] - def fake_download(**kwargs): + def fake_download(self, **kwargs): calls.append(kwargs) return "/cache/ltx-sidecar.safetensors" - monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download) + monkeypatch.setattr(huggingface_hub.HfApi, "hf_hub_download", fake_download) assert ( resolve_ltx_artifact( @@ -331,7 +332,9 @@ def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch): filename = "ltx-sidecar.safetensors" expected = tmp_path / filename expected.write_bytes(b"sidecar") - monkeypatch.setattr(ltx2_components, "hf_hub_download", lambda **_kwargs: pytest.fail("unexpected Hub lookup")) + monkeypatch.setattr( + huggingface_hub.HfApi, "hf_hub_download", lambda *_args, **_kwargs: pytest.fail("unexpected Hub lookup") + ) assert resolve_ltx_artifact( str(tmp_path), @@ -345,11 +348,11 @@ def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch): def test_ltx_artifact_local_model_missing_sidecar_falls_back_to_hub(tmp_path, monkeypatch): calls = [] - def fake_download(**kwargs): + def fake_download(self, **kwargs): calls.append(kwargs) return "/cache/ltx-sidecar.safetensors" - monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download) + monkeypatch.setattr(huggingface_hub.HfApi, "hf_hub_download", fake_download) assert ( resolve_ltx_artifact( diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py index 0fb043a1ec5..a45631f7dce 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py @@ -42,10 +42,10 @@ def _resolve_fl2va_model_ref() -> str: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api repo_root = Path( - snapshot_download( + hf_api().snapshot_download( repo_id=_MINIMAX_H3_REPO, revision=_MINIMAX_H3_REVISION, allow_patterns=["FL2VA/**"], @@ -253,13 +253,13 @@ def test_resolve_fl2va_model_ref(tmp_path, monkeypatch): fl2va_root.mkdir() (fl2va_root / "model_index.json").write_text("{}", encoding="utf-8") - def fake_snapshot_download(*, repo_id, revision, allow_patterns): + def fake_snapshot_download(self, *, repo_id, revision, allow_patterns): assert repo_id == _MINIMAX_H3_REPO assert revision == _MINIMAX_H3_REVISION assert allow_patterns == ["FL2VA/**"] return str(tmp_path) - monkeypatch.setattr("huggingface_hub.snapshot_download", fake_snapshot_download) + monkeypatch.setattr("huggingface_hub.HfApi.snapshot_download", fake_snapshot_download) assert _resolve_fl2va_model_ref() == str(fl2va_root) diff --git a/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py b/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py index 63ce7392371..c1b083495e9 100644 --- a/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py +++ b/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E accuracy guard against a pinned Lightricks LTX pipeline revision. @@ -22,11 +22,11 @@ import numpy as np import pytest import torch -from huggingface_hub import hf_hub_download, snapshot_download from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure from tests.e2e.accuracy.helpers import reset_artifact_dir from tests.helpers.mark import hardware_test +from vllm_omni.transformers_utils.repo_utils import hf_api OFFICIAL_REPOSITORY = "https://github.com/Lightricks/LTX-2.git" OFFICIAL_REVISION = "9377758131b1ffde4b7f766804590a6617bf2ab9" @@ -326,7 +326,7 @@ def _resolve_model(case: LTXAccuracyCase) -> Path: if revision is None and model_id == case.model_id: revision = case.model_revision return Path( - snapshot_download( + hf_api().snapshot_download( repo_id=model_id, revision=revision, allow_patterns=[ @@ -359,7 +359,7 @@ def _resolve_gemma_root(case: LTXAccuracyCase, model: Path) -> Path: if configured_model and Path(configured_model).is_dir(): return Path(configured_model) return Path( - snapshot_download( + hf_api().snapshot_download( repo_id=case.gemma_model_id, revision=case.gemma_model_revision, allow_patterns=[ @@ -383,7 +383,7 @@ def _resolve_artifact(artifact: LTXArtifact, model: Path | None = None) -> Path: if model_path.is_file(): return model_path return Path( - hf_hub_download( + hf_api().hf_hub_download( repo_id=artifact.repo_id, repo_type=None if artifact.repo_type == "model" else artifact.repo_type, filename=artifact.filename, diff --git a/tests/e2e/features/rlhf_test/test_verl_omni_e2e.py b/tests/e2e/features/rlhf_test/test_verl_omni_e2e.py index dab5feef1f8..8d8062092a1 100644 --- a/tests/e2e/features/rlhf_test/test_verl_omni_e2e.py +++ b/tests/e2e/features/rlhf_test/test_verl_omni_e2e.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E test that follows the EXACT same flow as ``verl-omni/tests/workers/rollout/rollout_vllm/test_vllm_omni_generate.py`` @@ -35,7 +35,6 @@ import pytest import ray import torch -from huggingface_hub import snapshot_download from omegaconf import OmegaConf from transformers import AutoTokenizer @@ -46,6 +45,7 @@ vLLMOmniHttpServerLocal, ) from tests.helpers.mark import hardware_test +from vllm_omni.transformers_utils.repo_utils import hf_api MODEL = "tiny-random/Qwen-Image" TOKENIZER_MODEL = "Qwen/Qwen2-1.5B-Instruct" @@ -55,7 +55,7 @@ def _resolve_model_path(repo_id: str) -> str: if os.path.isdir(repo_id): return repo_id - return snapshot_download(repo_id=repo_id) + return hf_api().snapshot_download(repo_id=repo_id) @lru_cache(maxsize=1) diff --git a/tests/e2e/offline_inference/test_cosyvoice3_expansion.py b/tests/e2e/offline_inference/test_cosyvoice3_expansion.py index 7e57433ef5c..dd7c8ec8509 100644 --- a/tests/e2e/offline_inference/test_cosyvoice3_expansion.py +++ b/tests/e2e/offline_inference/test_cosyvoice3_expansion.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Offline E2E smoke test for CosyVoice3 zero-shot reference inference. @@ -20,7 +20,6 @@ import numpy as np import pytest import soundfile as sf -from huggingface_hub import snapshot_download from vllm.sampling_params import SamplingParams from tests.helpers.mark import hardware_test @@ -29,6 +28,7 @@ from tests.helpers.stage_config import get_deploy_config_path from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config +from vllm_omni.transformers_utils.repo_utils import hf_api MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512" MODEL_DIR_ENV = "VLLM_OMNI_COSYVOICE3_MODEL_DIR" @@ -58,7 +58,7 @@ def _resolve_model_dir() -> Path: override = os.environ.get(MODEL_DIR_ENV) if override: return Path(override).expanduser().resolve() - return Path(snapshot_download(MODEL, allow_patterns=["*"])) + return Path(hf_api().snapshot_download(MODEL, allow_patterns=["*"])) def _reference_zero_shot_stage0_sampling(*, text: str) -> SamplingParams: diff --git a/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py b/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py index f570a6ea717..4f21dd1c485 100644 --- a/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py +++ b/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ End-to-end test for MammothModa2 text-to-image generation. @@ -19,12 +22,12 @@ import pytest import torch -from huggingface_hub import snapshot_download from vllm.sampling_params import SamplingParams 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.transformers_utils.repo_utils import hf_api # --------------------------------------------------------------------------- # Constants @@ -64,7 +67,7 @@ # Helpers # --------------------------------------------------------------------------- def _load_t2i_gen_config(repo_id: str) -> dict: - weights_dir = Path(snapshot_download(repo_id)) + weights_dir = Path(hf_api().snapshot_download(repo_id)) cfg_path = weights_dir / "t2i_generation_config.json" if not cfg_path.exists(): pytest.skip(f"t2i_generation_config.json not found at {cfg_path}") diff --git a/tests/e2e/offline_inference/test_moss_tts_realtime.py b/tests/e2e/offline_inference/test_moss_tts_realtime.py index 34d36118d99..81d97cd743c 100644 --- a/tests/e2e/offline_inference/test_moss_tts_realtime.py +++ b/tests/e2e/offline_inference/test_moss_tts_realtime.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """E2E offline inference tests for MOSS-TTS-Realtime (MossTTSRealtime, 1.7B). Uses the standard omni_runner + pytestmark pattern (one module-scoped engine @@ -160,12 +160,13 @@ def _build_request(ref_audio_path: str, text: str) -> dict: Frees the codec model before returning to avoid competing with the running vllm engine for GPU/CPU memory. """ - from huggingface_hub import snapshot_download from transformers import AutoModel, AutoTokenizer + from vllm_omni.transformers_utils.repo_utils import hf_api + # Step 1: locate the realtime processor module in the snapshot. try: - snap_dir = Path(snapshot_download(repo_id=MODEL)) + snap_dir = Path(hf_api().snapshot_download(repo_id=MODEL)) except Exception as exc: msg = f"Cannot locate snapshot for {MODEL}: {exc}" if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"): diff --git a/tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py b/tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py index 432124e76ff..66f8be616f6 100644 --- a/tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py +++ b/tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py @@ -12,7 +12,6 @@ from types import SimpleNamespace import pytest -from huggingface_hub import snapshot_download from tests.helpers.runtime import OmniServerParams, get_model_prefix from tests.helpers.stage_config import ( @@ -20,6 +19,7 @@ get_deploy_duplex_max_sessions, modify_stage_config, ) +from vllm_omni.transformers_utils.repo_utils import hf_api MODEL = "openbmb/MiniCPM-o-4_5" DEPLOY_CONFIG_REL = "minicpmo_4_5.yaml" @@ -111,7 +111,7 @@ def resolve_ref_audio() -> Path: model_prefix = get_model_prefix() model_root = Path(model_prefix) / MODEL if model_prefix else Path(MODEL) if not model_root.is_dir(): - model_root = Path(snapshot_download(MODEL, local_files_only=True)) + model_root = Path(hf_api().snapshot_download(MODEL, local_files_only=True)) ref_audio = model_root / REF_AUDIO_RELATIVE_PATH if not ref_audio.is_file(): raise FileNotFoundError(f"MiniCPM-o checkpoint ref audio is missing: {ref_audio}") diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index fb8231a75c6..0e88a00ef88 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Tests for Omni config utils. For stability, these tests should largely be invariant to the specific attributes of vLLM config except in cases where we @@ -63,7 +66,7 @@ class ConnectorRunner(OmniConnectorModelRunnerMixin): class ConnectorWorker: model_runner_cls = ConnectorRunner - worker_module.ConnectorWorker = ConnectorWorker + setattr(worker_module, "ConnectorWorker", ConnectorWorker) monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module) monkeypatch.setattr(OmniEngineArgs, "_ensure_omni_models_registered", lambda _self: None) @@ -95,7 +98,7 @@ def test_stage_without_connector_configuration_accepts_plain_runner(monkeypatch) class WorkerWithoutConnector: model_runner_cls = object - worker_module.WorkerWithoutConnector = WorkerWithoutConnector + setattr(worker_module, "WorkerWithoutConnector", WorkerWithoutConnector) monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module) args = OmniEngineArgs( @@ -113,7 +116,7 @@ def test_full_payload_capability_requires_selected_worker_connector(monkeypatch) class WorkerWithoutConnector: model_runner_cls = object - worker_module.WorkerWithoutConnector = WorkerWithoutConnector + setattr(worker_module, "WorkerWithoutConnector", WorkerWithoutConnector) monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module) with pytest.raises(ValueError, match="does not provide an Omni connector model runner"): @@ -130,7 +133,7 @@ def test_full_payload_capability_validates_platform_selected_worker(monkeypatch) class WorkerWithoutConnector: model_runner_cls = object - worker_module.WorkerWithoutConnector = WorkerWithoutConnector + setattr(worker_module, "WorkerWithoutConnector", WorkerWithoutConnector) monkeypatch.setitem(sys.modules, worker_module.__name__, worker_module) monkeypatch.setattr( current_omni_platform, @@ -293,7 +296,7 @@ def test_remote_tokenizer_subfolder_download_does_not_report_failure(tmp_path, m baseline_config = Mock() warning = mocker.patch("vllm_omni.engine.arg_utils.logger.warning") - monkeypatch.setattr("huggingface_hub.snapshot_download", lambda *args, **kwargs: str(tmp_path)) + monkeypatch.setattr("huggingface_hub.HfApi.snapshot_download", lambda *args, **kwargs: str(tmp_path)) monkeypatch.setattr(OmniEngineArgs, "_patch_empty_hf_config", lambda *args, **kwargs: None) monkeypatch.setattr(EngineArgs, "create_model_config", lambda _self: baseline_config) monkeypatch.setattr( @@ -339,6 +342,7 @@ def test_patch_missing_local_hf_config(tmp_path): def test_non_missing_local_hf_config_error_reaches_parent_loader(tmp_path, monkeypatch, config_entry): """Non-missing config errors must reach vLLM's normal loader.""" config_path = tmp_path / "config.json" + loader_error: Exception if config_entry == "malformed": config_path.write_text("{not valid json", encoding="utf-8") loader_error = json.JSONDecodeError("invalid config", "{not valid json", 1) diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py index a4e641a571a..9120e183f8d 100644 --- a/tests/helpers/__init__.py +++ b/tests/helpers/__init__.py @@ -29,12 +29,13 @@ def skip_if_gated_repo_inaccessible( and we skip cleanly. """ try: - from huggingface_hub import hf_hub_download from huggingface_hub.errors import GatedRepoError, RepositoryNotFoundError + + from vllm_omni.transformers_utils.repo_utils import hf_api except Exception: return try: - hf_hub_download(repo_id=repo_id, filename=filename, revision=revision) + hf_api().hf_hub_download(repo_id=repo_id, filename=filename, revision=revision) except GatedRepoError as exc: import pytest diff --git a/tests/model_executor/models/minicpmo_4_5/test_code2wav_model_dir.py b/tests/model_executor/models/minicpmo_4_5/test_code2wav_model_dir.py index 5aa78262ae8..8d5aeade657 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_code2wav_model_dir.py +++ b/tests/model_executor/models/minicpmo_4_5/test_code2wav_model_dir.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 """Unit tests for MiniCPM-o 4.5 code2wav model-dir resolution (#5442). In hub/CI deployments ``model_config.model`` is a repo id rather than a local @@ -27,7 +27,7 @@ def _no_hub(monkeypatch): def _fail(*args, **kwargs): # pragma: no cover - must not be reached raise AssertionError("snapshot_download must not be called here") - monkeypatch.setattr(huggingface_hub, "snapshot_download", _fail) + monkeypatch.setattr(huggingface_hub.HfApi, "snapshot_download", _fail) def test_local_directory_is_returned_unchanged(tmp_path, monkeypatch): @@ -38,13 +38,13 @@ def test_local_directory_is_returned_unchanged(tmp_path, monkeypatch): def test_repo_id_resolves_via_snapshot_download(tmp_path, monkeypatch): calls = {} - def _fake_snapshot_download(model_ref, revision=None, allow_patterns=None): + def _fake_snapshot_download(self, model_ref, revision=None, allow_patterns=None): calls["model_ref"] = model_ref calls["revision"] = revision calls["allow_patterns"] = allow_patterns return str(tmp_path / "snapshot") - monkeypatch.setattr(huggingface_hub, "snapshot_download", _fake_snapshot_download) + monkeypatch.setattr(huggingface_hub.HfApi, "snapshot_download", _fake_snapshot_download) resolved = _resolve_model_dir("openbmb/MiniCPM-o-4_5", revision="abc123") assert resolved == str(tmp_path / "snapshot") assert calls["model_ref"] == "openbmb/MiniCPM-o-4_5" @@ -56,7 +56,7 @@ def test_snapshot_download_failure_propagates(monkeypatch): def _raise(*args, **kwargs): raise FileNotFoundError("offline and not cached") - monkeypatch.setattr(huggingface_hub, "snapshot_download", _raise) + monkeypatch.setattr(huggingface_hub.HfApi, "snapshot_download", _raise) with pytest.raises(FileNotFoundError): _resolve_model_dir("openbmb/MiniCPM-o-4_5") diff --git a/tests/model_executor/models/omnivoice/test_nonverbal_tags_tokenize.py b/tests/model_executor/models/omnivoice/test_nonverbal_tags_tokenize.py index fbf5ac0b010..95524fdced3 100644 --- a/tests/model_executor/models/omnivoice/test_nonverbal_tags_tokenize.py +++ b/tests/model_executor/models/omnivoice/test_nonverbal_tags_tokenize.py @@ -1,12 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import pytest -from huggingface_hub import hf_hub_download from tokenizers import Tokenizer as HFTokenizer from vllm_omni.diffusion.models.omnivoice.pipeline_omnivoice import ( _tokenize_with_nonverbal_tags, ) +from vllm_omni.transformers_utils.repo_utils import hf_api pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -18,7 +21,7 @@ def _ids(tok: HFTokenizer, s: str) -> list[int]: class TestNonVerbalTags: @classmethod def setup_class(cls): - tokenizer_path = hf_hub_download(repo_id="k2-fsa/OmniVoice", filename="tokenizer.json") + tokenizer_path = hf_api().hf_hub_download(repo_id="k2-fsa/OmniVoice", filename="tokenizer.json") cls.tokenizer = HFTokenizer.from_file(tokenizer_path) def test_plain_text_fallback(self): diff --git a/tools/pre_commit/check_forbidden_imports.py b/tools/pre_commit/check_forbidden_imports.py index 22ba693d4f2..c5a96148ea2 100644 --- a/tools/pre_commit/check_forbidden_imports.py +++ b/tools/pre_commit/check_forbidden_imports.py @@ -17,7 +17,7 @@ _HF_NAMES = ( r"HfApi|HfFileSystem|hf_hub_download|snapshot_download" r"|list_repo_files|file_exists|try_to_load_from_cache" - r"|list_repo_refs|repo_exists" + r"|list_repo_refs|repo_exists|get_safetensors_metadata" ) # Non-library trees: examples, tests, and tooling may keep stdlib imports. @@ -169,56 +169,20 @@ class ForbiddenImport: r"^\s*from\s+huggingface_hub\s+import\b[^\n]*\b(?:" + _HF_NAMES + r")\b" ), tip=( - "Use the shared, vLLM-tagged helpers from " - "vllm.transformers_utils.repo_utils (e.g. hf_api(), hf_fs(), " - "list_repo_files, file_exists) instead of calling " - "huggingface_hub directly." + "Use 'hf_api()' from 'vllm_omni.transformers_utils.repo_utils' (or " + "add a tagged helper there) instead, so requests are tagged with " + "vLLM-Omni's library info." ), allowed_dirs=_NON_LIBRARY_DIRS, allowed_files={ - "vllm_omni/benchmarks/data_modules/daily_omni_dataset.py", - "vllm_omni/benchmarks/data_modules/seed_tts_dataset.py", - "vllm_omni/benchmarks/data_modules/seed_tts_eval.py", - "vllm_omni/diffusion/lora/loader.py", - "vllm_omni/diffusion/model_loader/hub_prefetch.py", - "vllm_omni/diffusion/models/cosmos3/sound_tokenizer.py", - "vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py", - "vllm_omni/diffusion/models/helios/pipeline_helios.py", - "vllm_omni/diffusion/models/longcat_video/pipeline_longcat_video_avatar.py", - "vllm_omni/diffusion/models/ltx2/ltx2_components.py", - "vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py", - "vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py", - "vllm_omni/diffusion/models/pi0/pipeline_pi0.py", - "vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py", - "vllm_omni/diffusion/models/soulx_singer/utils.py", - "vllm_omni/diffusion/models/utils.py", - "vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py", - "vllm_omni/engine/arg_utils.py", - "vllm_omni/entrypoints/openai/api_server.py", - "vllm_omni/entrypoints/openai/serving_speech.py", - "vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py", - "vllm_omni/model_executor/model_loader/weight_utils.py", - "vllm_omni/model_executor/models/audex/checkpoint.py", - "vllm_omni/model_executor/models/bagel/bagel.py", - "vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py", - "vllm_omni/model_executor/models/covo_audio/covo_audio_code2wav.py", - "vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py", - "vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.py", - "vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2audio.py", - "vllm_omni/model_executor/models/glm_tts/glm_tts.py", + # ``try_to_load_from_cache`` is a local cache lookup that issues no + # Hub request, so there is nothing for ``hf_api()`` to tag. "vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py", "vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py", "vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py", - "vllm_omni/model_executor/models/indextts2/preprocess_utils.py", - "vllm_omni/model_executor/models/indextts2/s2mel/modules/bigvgan.py", - "vllm_omni/model_executor/models/indextts2/tokenizer_v2_5.py", - "vllm_omni/model_executor/models/ming_tts/speaker_extractor.py", - "vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py", - "vllm_omni/model_executor/models/omnivoice/omnivoice.py", - "vllm_omni/model_executor/models/personaplex/personaplex_mimi.py", - "vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py", - "vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py", "vllm_omni/transformers_utils/configs/higgs_audio_v3.py", + # Defines the tagged helper itself. + "vllm_omni/transformers_utils/repo_utils.py", }, ), } @@ -330,6 +294,7 @@ def matches(rule: str, content: str) -> bool: ("from huggingface_hub import HfFileSystem", True), ("from huggingface_hub import list_repo_files", True), ("from huggingface_hub import try_to_load_from_cache", True), + ("from huggingface_hub import get_safetensors_metadata", True), (" from huggingface_hub import snapshot_download", True), ("from huggingface_hub import PyTorchModelHubMixin, hf_hub_download", True), ("from huggingface_hub import (snapshot_download)", True), @@ -346,6 +311,7 @@ def matches(rule: str, content: str) -> bool: ("from huggingface_hub.constants import HF_HUB_CACHE", False), ("from huggingface_hub.utils import EntryNotFoundError", False), ("from vllm.transformers_utils.repo_utils import hf_api", False), + ("from vllm_omni.transformers_utils.repo_utils import hf_api", False), ("from huggingface_hub import (\n PyTorchModelHubMixin,\n)", False), ("# from huggingface_hub import snapshot_download", False), ] diff --git a/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py b/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py index abb02b1d500..3bf852e8c54 100644 --- a/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py +++ b/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py @@ -170,7 +170,7 @@ def ensure_daily_omni_hub_videos_dir(repo_id: str) -> Path: raise ValueError("repo_id is required to download Daily-Omni Videos.tar") try: - from huggingface_hub import hf_hub_download + from vllm_omni.transformers_utils.repo_utils import hf_api except ImportError as e: raise ImportError( "Daily-Omni Hub media download requires huggingface_hub. " @@ -180,7 +180,7 @@ def ensure_daily_omni_hub_videos_dir(repo_id: str) -> Path: tar_path: Path | None = None for fname in ("Videos.tar", "videos.tar"): try: - tar_path = Path(hf_hub_download(repo_id=rid, filename=fname, repo_type="dataset")) + tar_path = Path(hf_api().hf_hub_download(repo_id=rid, filename=fname, repo_type="dataset")) break except Exception: continue @@ -237,14 +237,14 @@ def ensure_daily_omni_hub_root(repo_id: str) -> Path: raise ValueError("repo_id is required to download Daily-Omni from Hugging Face") try: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api except ImportError as e: raise ImportError( "Install huggingface_hub to load Daily-Omni from the Hub, or pass a local " "--daily-omni-qa-json / --daily-omni-video-dir mirror." ) from e - cache = snapshot_download( + cache = hf_api().snapshot_download( repo_id=rid, repo_type="dataset", allow_patterns=_DAILY_OMNI_QA_ALLOW_PATTERNS, diff --git a/vllm_omni/benchmarks/data_modules/seed_tts_dataset.py b/vllm_omni/benchmarks/data_modules/seed_tts_dataset.py index 23e8b8ade63..996f4f24865 100644 --- a/vllm_omni/benchmarks/data_modules/seed_tts_dataset.py +++ b/vllm_omni/benchmarks/data_modules/seed_tts_dataset.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Seed-TTS zero-shot evaluation-style prompts for ``vllm bench serve``. Loads rows from the `meta.lst` format used in `BytedanceSpeech/seed-tts-eval`_ (or any @@ -121,7 +124,7 @@ def resolve_seed_tts_root(dataset_path: str | None, *, explicit_root: str | None repo_id = dataset_path.strip() try: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api except ImportError as e: raise ImportError( "Install huggingface_hub to download Seed-TTS from the Hub, or clone the dataset " @@ -130,7 +133,7 @@ def resolve_seed_tts_root(dataset_path: str | None, *, explicit_root: str | None # Download only the requested locale subtree instead of the whole dataset # repo. This avoids large, flaky nightly downloads when we only need e.g. # ``en/meta.lst`` + ``en/prompt-wavs/**``. - cache = snapshot_download( + cache = hf_api().snapshot_download( repo_id=repo_id, repo_type="dataset", allow_patterns=[f"{locale}/**"], diff --git a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py index ddfbbd63058..9dc99548aeb 100644 --- a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py +++ b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Seed-TTS WER aligned with Bytedance ``seed-tts-eval`` / ``run_wer.py``. Matches the published protocol (see Hugging Face dataset card and @@ -329,7 +332,8 @@ def _ensure_utmos_jit_model() -> Any | None: return _utmos_jit_model try: import torch - from huggingface_hub import hf_hub_download + + from vllm_omni.transformers_utils.repo_utils import hf_api repo = os.environ.get("SEED_TTS_UTMOS_HF_REPO", "balacoon/utmos").strip() or "balacoon/utmos" fname = os.environ.get("SEED_TTS_UTMOS_JIT_FILE", "utmos.jit").strip() or "utmos.jit" @@ -338,7 +342,7 @@ def _ensure_utmos_jit_model() -> Any | None: repo, fname, ) - path = hf_hub_download(repo_id=repo, filename=fname, repo_type="model") + path = hf_api().hf_hub_download(repo_id=repo, filename=fname, repo_type="model") # TODO The model weights in UTMOS must be loaded in cuda:0; otherwise, the model execution will fail. want = "cuda:0" diff --git a/vllm_omni/diffusion/lora/loader.py b/vllm_omni/diffusion/lora/loader.py index 3f74b4853a2..a13363c7a77 100644 --- a/vllm_omni/diffusion/lora/loader.py +++ b/vllm_omni/diffusion/lora/loader.py @@ -10,11 +10,11 @@ _convert_non_diffusers_qwen_lora_to_diffusers, _convert_non_diffusers_wan_lora_to_diffusers, ) -from huggingface_hub import hf_hub_download from safetensors.torch import load_file from vllm.logger import init_logger from vllm_omni.diffusion.utils.tf_utils import get_transformer_from_pipeline +from vllm_omni.transformers_utils.repo_utils import hf_api logger = init_logger(__name__) @@ -131,7 +131,7 @@ def _load_lora_state_dict( # finally, we try to load it from the internet try: - model_file = hf_hub_download( + model_file = hf_api().hf_hub_download( pretrained_model_name_or_path, filename=weights_name, subfolder=subfolder, diff --git a/vllm_omni/diffusion/model_loader/hub_prefetch.py b/vllm_omni/diffusion/model_loader/hub_prefetch.py index 3594aa4e1c0..a1e893b3be6 100644 --- a/vllm_omni/diffusion/model_loader/hub_prefetch.py +++ b/vllm_omni/diffusion/model_loader/hub_prefetch.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 """Best-effort HuggingFace Hub prefetch for multi-subfolder pipelines. @@ -297,7 +297,7 @@ def prefetch_subfolders( logger.info("Prefetching %s subfolders: %s", model, subfolders) try: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api except ImportError: # pragma: no cover - huggingface_hub is a hard dep logger.debug("huggingface_hub unavailable; skipping prefetch of %s", model) return @@ -334,7 +334,7 @@ def prefetch_subfolders( for attempt in range(1, _PREFETCH_MAX_ATTEMPTS + 1): try: with _repo_prefetch_lock(model): - snapshot_download( + hf_api().snapshot_download( repo_id=model, revision=revision, allow_patterns=allow_patterns, diff --git a/vllm_omni/diffusion/models/cosmos3/sound_tokenizer.py b/vllm_omni/diffusion/models/cosmos3/sound_tokenizer.py index 5d6ff8dd968..8c1337ec90d 100644 --- a/vllm_omni/diffusion/models/cosmos3/sound_tokenizer.py +++ b/vllm_omni/diffusion/models/cosmos3/sound_tokenizer.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 """Cosmos3 sound tokenizer integration. The tokenizer decodes model-generated sound latents for video+audio output. It @@ -338,9 +338,9 @@ def from_config(cls, od_config: OmniDiffusionConfig) -> Cosmos3SoundTokenizer: model_root = str(model_path) if model_path and os.path.isdir(model_path) else None if model_root is None and model_path and not explicit_avae_path: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - model_root = snapshot_download( + model_root = hf_api().snapshot_download( repo_id=str(model_path), revision=getattr(od_config, "revision", None), allow_patterns=[ diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py index bf19f0a0fc0..09bd536cd7f 100644 --- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py +++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py @@ -21,7 +21,6 @@ import numpy as np import torch import torch.nn as nn -from huggingface_hub import hf_hub_download from transformers import AutoTokenizer, UMT5Config, UMT5EncoderModel from vllm.model_executor.model_loader.weight_utils import default_weight_loader @@ -67,6 +66,7 @@ resolve_session_state_config, ) from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.transformers_utils.repo_utils import hf_api logger = logging.getLogger(__name__) MAX_DREAMZERO_SESSIONS = 64 @@ -581,7 +581,7 @@ def _load_repo_json(model_path: str, relative_path: str, local_files_only: bool) return json.load(f) try: - json_path = hf_hub_download(model_path, relative_path) + json_path = hf_api().hf_hub_download(model_path, relative_path) with open(json_path) as f: return json.load(f) except Exception: diff --git a/vllm_omni/diffusion/models/helios/pipeline_helios.py b/vllm_omni/diffusion/models/helios/pipeline_helios.py index 820f9a1d98a..610fb2b8fbe 100644 --- a/vllm_omni/diffusion/models/helios/pipeline_helios.py +++ b/vllm_omni/diffusion/models/helios/pipeline_helios.py @@ -76,9 +76,9 @@ def load_json_config(model_path: str, subfolder: str, filename: str, local_files return json.load(f) else: try: - from huggingface_hub import hf_hub_download + from vllm_omni.transformers_utils.repo_utils import hf_api - config_path = hf_hub_download( + config_path = hf_api().hf_hub_download( repo_id=model_path, filename=f"{subfolder}/{filename}", ) diff --git a/vllm_omni/diffusion/models/longcat_video/pipeline_longcat_video_avatar.py b/vllm_omni/diffusion/models/longcat_video/pipeline_longcat_video_avatar.py index 33b0e2a458a..6cb2dd7fb50 100644 --- a/vllm_omni/diffusion/models/longcat_video/pipeline_longcat_video_avatar.py +++ b/vllm_omni/diffusion/models/longcat_video/pipeline_longcat_video_avatar.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 __future__ import annotations @@ -376,9 +376,9 @@ def _ensure_local_dir(model: str | os.PathLike[str], allow_patterns: list[str] | model_path = Path(model) if model_path.exists(): return model_path - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - return Path(snapshot_download(str(model), allow_patterns=allow_patterns)) + return Path(hf_api().snapshot_download(str(model), allow_patterns=allow_patterns)) def _ensure_avatar_model_index(model_dir: Path, use_int8: bool) -> None: diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py index c00fdd8355a..72e63a22f3b 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_components.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_components.py @@ -20,7 +20,6 @@ from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder from diffusers.video_processor import VideoProcessor -from huggingface_hub import hf_hub_download from safetensors import safe_open from safetensors.torch import load_file from transformers import AutoModelForImageTextToText, AutoTokenizer, Gemma3ForConditionalGeneration @@ -32,6 +31,7 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery +from vllm_omni.transformers_utils.repo_utils import hf_api if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig @@ -279,7 +279,7 @@ def resolve_ltx_artifact( # independently pinned artifact revision. revision = model_revision if model == repo_id else artifact_revision try: - return hf_hub_download( + return hf_api().hf_hub_download( repo_id=repo_id, filename=filename, revision=revision, @@ -351,7 +351,7 @@ def _load_ltx_metadata_json(model: str, filename: str, revision: str | None = No return {} else: try: - path = hf_hub_download(repo_id=model, filename=filename, revision=revision) + path = hf_api().hf_hub_download(repo_id=model, filename=filename, revision=revision) except Exception: return {} try: @@ -529,7 +529,7 @@ def _detect_vocoder_output_sample_rate(model: str, revision: str | None = None) vocoder_config_path = os.path.join(model, "vocoder", "config.json") if not os.path.exists(vocoder_config_path): try: - vocoder_config_path = hf_hub_download(model, "vocoder/config.json", revision=revision) + vocoder_config_path = hf_api().hf_hub_download(model, "vocoder/config.json", revision=revision) except Exception: return None try: @@ -839,7 +839,7 @@ def load_transformer_config( if not os.path.exists(config_path): raise FileNotFoundError(f"LTX transformer config not found: {config_path}") else: - config_path = hf_hub_download( + config_path = hf_api().hf_hub_download( repo_id=model_path, filename=f"{subfolder}/config.json", revision=revision, diff --git a/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py b/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py index 62ac027ca8b..cfc40fcb6bc 100644 --- a/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py +++ b/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.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 """ OmniVoice TTS Pipeline for vLLM-Omni diffusion engine. @@ -141,9 +141,9 @@ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): # Resolve model path (HF hub ID → local cache) if not os.path.isdir(self.model_path): - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - self.model_path = snapshot_download(self.model_path) + self.model_path = hf_api().snapshot_download(self.model_path) # Load OmniVoice config config_path = os.path.join(self.model_path, "config.json") diff --git a/vllm_omni/diffusion/models/pi0/pipeline_pi0.py b/vllm_omni/diffusion/models/pi0/pipeline_pi0.py index 1eeca29692c..a03801c8d4e 100644 --- a/vllm_omni/diffusion/models/pi0/pipeline_pi0.py +++ b/vllm_omni/diffusion/models/pi0/pipeline_pi0.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 """π0 (Pi-Zero) VLA pipeline for vllm-omni. Entry point for ``DiffusionEngine.step() → pipeline.forward(req)``. Mirrors the @@ -94,9 +94,9 @@ def _resolve_model_dir(model: str | None) -> str | None: return None if os.path.isdir(model): return model - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - return snapshot_download( + return hf_api().snapshot_download( repo_id=model, allow_patterns=["*.json", "*.safetensors", "*.model", "tokenizer*"], ) diff --git a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py index de636fb6dca..2be1a9c7a39 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py +++ b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py @@ -27,7 +27,6 @@ import torch import torch.nn as nn import torchvision.transforms as T -from huggingface_hub import snapshot_download from PIL import Image from transformers import AutoTokenizer from vllm.logger import init_logger @@ -109,7 +108,9 @@ def _resolve_model_path(model_path: str) -> str: """Resolve a HuggingFace model ID or local path to a local directory.""" if os.path.isdir(model_path): return model_path - return snapshot_download(model_path) + from vllm_omni.transformers_utils.repo_utils import hf_api + + return hf_api().snapshot_download(model_path) def _round_by_factor(number: float, factor: int) -> int: diff --git a/vllm_omni/diffusion/models/utils.py b/vllm_omni/diffusion/models/utils.py index 2d783d90712..429c34b9d5c 100644 --- a/vllm_omni/diffusion/models/utils.py +++ b/vllm_omni/diffusion/models/utils.py @@ -168,8 +168,8 @@ def _load_json(model_path: str, filename: str, local_files_only: bool = True) -> with open(path) as f: return json.load(f) else: - from huggingface_hub import hf_hub_download + from vllm_omni.transformers_utils.repo_utils import hf_api - cached = hf_hub_download(repo_id=model_path, filename=filename) + cached = hf_api().hf_hub_download(repo_id=model_path, filename=filename) with open(cached) as f: return json.load(f) diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index 2383fc6fd09..c9ef325c4ee 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -156,9 +156,9 @@ def load_transformer_config(model_path: str, subfolder: str = "transformer", loc else: # Try to download config from HF Hub try: - from huggingface_hub import hf_hub_download + from vllm_omni.transformers_utils.repo_utils import hf_api - config_path = hf_hub_download( + config_path = hf_api().hf_hub_download( repo_id=model_path, filename=f"{subfolder}/config.json", ) @@ -364,9 +364,9 @@ def __init__( else: # For remote models, download and read model_index.json try: - from huggingface_hub import hf_hub_download + from vllm_omni.transformers_utils.repo_utils import hf_api - model_index_path = hf_hub_download(repo_id=model, filename="model_index.json") + model_index_path = hf_api().hf_hub_download(repo_id=model, filename="model_index.json") with open(model_index_path) as f: model_index = json.load(f) self.expand_timesteps = model_index.get("expand_timesteps", False) diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index f9ee40cbbff..bd6a7156215 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -381,9 +381,9 @@ def create_model_config(self) -> OmniModelConfig: if tokenizer_subfolder: # Download just the tokenizer files from the subfolder try: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - local_dir = snapshot_download( + local_dir = hf_api().snapshot_download( model_path, allow_patterns=[ f"{tokenizer_subfolder}/tokenizer*", diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 0092e40cc53..1e57036a0d1 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -31,7 +31,6 @@ from vllm.renderers import BaseRenderer from vllm.sampling_params import SamplingParams from vllm.tokenizers import cached_tokenizer_from_config -from vllm.transformers_utils.repo_utils import hf_api from vllm.transformers_utils.runai_utils import is_runai_obj_uri from vllm.usage.usage_lib import UsageContext from vllm.v1.engine.input_processor import InputProcessor @@ -59,6 +58,7 @@ from vllm_omni.outputs.output_processor import MultimodalOutputProcessor from vllm_omni.platforms import current_omni_platform from vllm_omni.quantization.inc_config import OmniINCConfig +from vllm_omni.transformers_utils.repo_utils import hf_api logger = init_logger(__name__) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index f10edfc7f31..f97207bbc9c 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -201,9 +201,9 @@ def _load_model_chat_template_json(model: str) -> str | None: if template_path is None: try: - from huggingface_hub import hf_hub_download + from vllm_omni.transformers_utils.repo_utils import hf_api - template_path = hf_hub_download( + template_path = hf_api().hf_hub_download( repo_id=model, filename="chat_template.json", local_files_only=True, diff --git a/vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py b/vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py index 25993ed6e3e..1196c5c709d 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py @@ -116,9 +116,9 @@ def apply_sampling_overrides( if tokenizer is None: model_path = server.engine_client.model_config.model if not os.path.isdir(model_path): - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - model_path = snapshot_download(model_path) + model_path = hf_api().snapshot_download(model_path) tokenizer = get_qwen_tokenizer( token_path=os.path.join(model_path, hf_cfg.qwen_pretrain_path), skip_special_tokens=hf_cfg.skip_special_tokens, diff --git a/vllm_omni/model_executor/model_loader/weight_utils.py b/vllm_omni/model_executor/model_loader/weight_utils.py index 59c2cd86022..559a90b5062 100644 --- a/vllm_omni/model_executor/model_loader/weight_utils.py +++ b/vllm_omni/model_executor/model_loader/weight_utils.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import os import time from pathlib import Path @@ -10,7 +13,11 @@ if envs.VLLM_USE_MODELSCOPE: from modelscope.hub.snapshot_download import snapshot_download else: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api + + def snapshot_download(*args, **kwargs): + return hf_api().snapshot_download(*args, **kwargs) + logger = init_logger(__name__) diff --git a/vllm_omni/model_executor/models/audex/checkpoint.py b/vllm_omni/model_executor/models/audex/checkpoint.py index 38b0531e660..ae8230e05d6 100644 --- a/vllm_omni/model_executor/models/audex/checkpoint.py +++ b/vllm_omni/model_executor/models/audex/checkpoint.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 """Audex checkpoint layout preparation. The nvidia/Nemotron-Labs-Audex-2B repo deduplicates weights: the @@ -57,12 +57,12 @@ def _download_audiogen_index(model: str) -> str: metadata round trip per stage-engine boot; a cache miss falls back to the network. """ - from huggingface_hub import hf_hub_download + from vllm_omni.transformers_utils.repo_utils import hf_api try: - return hf_hub_download(model, _AUDIOGEN_INDEX_FILE, local_files_only=True) + return hf_api().hf_hub_download(model, _AUDIOGEN_INDEX_FILE, local_files_only=True) except Exception: - return hf_hub_download(model, _AUDIOGEN_INDEX_FILE) + return hf_api().hf_hub_download(model, _AUDIOGEN_INDEX_FILE) def _dedup_shard_patterns(model: str) -> list[str]: @@ -111,13 +111,13 @@ def ensure_audex_snapshot(model: str, profile: str = "tts") -> str: if profile in ("tts", "tta"): patterns += _dedup_shard_patterns(model) - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api try: - return snapshot_download(model, allow_patterns=patterns) + return hf_api().snapshot_download(model, allow_patterns=patterns) except Exception as download_exc: try: - return snapshot_download(model, allow_patterns=patterns, local_files_only=True) + return hf_api().snapshot_download(model, allow_patterns=patterns, local_files_only=True) except Exception: raise RuntimeError( f"Could not resolve the Audex repo {model!r} (profile {profile!r}): the download " @@ -135,13 +135,13 @@ def ensure_xcodec1_snapshot(model: str | None) -> str: if os.path.isdir(model): return model - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api try: - return snapshot_download(model) + return hf_api().snapshot_download(model) except Exception as download_exc: try: - return snapshot_download(model, local_files_only=True) + return hf_api().snapshot_download(model, local_files_only=True) except Exception: raise RuntimeError( f"Could not resolve the XCodec1 checkpoint {model!r}: the download failed and " diff --git a/vllm_omni/model_executor/models/bagel/bagel.py b/vllm_omni/model_executor/models/bagel/bagel.py index a3dfd7bf151..495e0037f29 100644 --- a/vllm_omni/model_executor/models/bagel/bagel.py +++ b/vllm_omni/model_executor/models/bagel/bagel.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from collections.abc import Iterable, Mapping, Sequence from math import isqrt from typing import Any @@ -122,9 +125,9 @@ def _patch_max_latent_size(self, config): if p.is_dir(): index_path = p / "model.safetensors.index.json" else: - from huggingface_hub import hf_hub_download + from vllm_omni.transformers_utils.repo_utils import hf_api - index_path = Path(hf_hub_download(model_name, "model.safetensors.index.json")) + index_path = Path(hf_api().hf_hub_download(model_name, "model.safetensors.index.json")) if not index_path.exists(): return diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index 81bae6d6c2b..84652cadabf 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -10,7 +10,6 @@ import onnxruntime import torch import torch.nn as nn -from huggingface_hub import snapshot_download from scipy.signal import resample_poly from transformers import Qwen2Config from transformers.feature_extraction_utils import BatchFeature @@ -52,6 +51,7 @@ from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config +from vllm_omni.transformers_utils.repo_utils import hf_api from vllm_omni.utils.speaker_cache import get_speaker_cache logger = init_logger(__name__) @@ -137,7 +137,7 @@ def _build_runtime_components(self, model_dir: str, config: CosyVoice3Config) -> """Build the per-model runtime components once (cached process-wide).""" # If model_dir is an HF repo ID (not a local path), resolve to cache. if not os.path.isdir(model_dir): - model_dir = snapshot_download(model_dir) + model_dir = hf_api().snapshot_download(model_dir) tokenizer = get_qwen_tokenizer( token_path=os.path.join(model_dir, config.qwen_pretrain_path), @@ -463,7 +463,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.model_stage = vllm_config.model_config.model_stage model_dir = vllm_config.model_config.model if not os.path.isdir(model_dir): - model_dir = snapshot_download(model_dir) + model_dir = hf_api().snapshot_download(model_dir) self.model_dir = model_dir self.model = None if self.model_stage == "cosyvoice3_talker": @@ -918,7 +918,7 @@ def _resolve_flow_estimator_onnx(self) -> str | None: repo = getattr(self.config, "flow_estimator_onnx_repo", None) if repo: try: - fetched_dir = snapshot_download(repo, allow_patterns=[fp16_name]) + fetched_dir = hf_api().snapshot_download(repo, allow_patterns=[fp16_name]) fetched = os.path.join(fetched_dir, fp16_name) if os.path.exists(fetched): return fetched diff --git a/vllm_omni/model_executor/models/covo_audio/covo_audio_code2wav.py b/vllm_omni/model_executor/models/covo_audio/covo_audio_code2wav.py index 4b0b2fe4bea..15cd0875e64 100644 --- a/vllm_omni/model_executor/models/covo_audio/covo_audio_code2wav.py +++ b/vllm_omni/model_executor/models/covo_audio/covo_audio_code2wav.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # Copyright 2026 Tencent. import os from collections.abc import Iterable @@ -25,9 +28,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): if os.path.isdir(model_name): model_path = model_name else: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - model_path = snapshot_download(model_name) + model_path = hf_api().snapshot_download(model_name) token2wav_path = os.path.join(model_path, "token2wav") code2wav_config = CovoAudioCode2WavConfig() diff --git a/vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py b/vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py index ff2f460d637..862af8e17cf 100644 --- a/vllm_omni/model_executor/models/dots_tts/dots_tts_talker.py +++ b/vllm_omni/model_executor/models/dots_tts/dots_tts_talker.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 """dots.tts talker — vLLM-native AR base LM + audio side path. Mirrors upstream rednote-hilab/dots.tts (pinned @ a393d2e): @@ -516,9 +516,9 @@ def _resolve_latent_stats_path(model_arg: str) -> str | None: if os.path.exists(local): return local try: - from huggingface_hub import hf_hub_download + from vllm_omni.transformers_utils.repo_utils import hf_api - return hf_hub_download( + return hf_api().hf_hub_download( repo_id=model_arg, filename="latent_stats.pt", local_files_only=True, diff --git a/vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.py b/vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.py index 5c67decaa56..125822c1ef4 100644 --- a/vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.py +++ b/vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.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 hashlib @@ -20,9 +23,9 @@ logger = init_logger(__name__) try: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api except Exception: # pragma: no cover - snapshot_download = None + hf_api = None class DetokTarget(IntEnum): @@ -397,13 +400,13 @@ def _find_dynin_config_under_root(root: Path) -> Path | None: @lru_cache(maxsize=16) def _resolve_dynin_config_from_hf_repo(repo_id: str) -> str | None: - if not _looks_like_hf_repo_id(repo_id) or snapshot_download is None: + if not _looks_like_hf_repo_id(repo_id) or hf_api is None: return None try: snapshot_dir = ( Path( - snapshot_download( + hf_api().snapshot_download( repo_id=repo_id, repo_type="model", allow_patterns=list(_DYNIN_CONFIG_CANDIDATE_RELPATHS), @@ -719,7 +722,7 @@ def _resolve_remote_snapshot_dir( if source_path.is_dir(): return str(source_path.resolve()) - if snapshot_download is None: + if hf_api is None: raise RuntimeError("huggingface_hub is required to load remote code.") kwargs: dict[str, Any] = { @@ -732,10 +735,10 @@ def _resolve_remote_snapshot_dir( kwargs["revision"] = revision try: - return str(snapshot_download(**kwargs)) + return str(hf_api().snapshot_download(**kwargs)) except TypeError: kwargs.pop("local_files_only", None) - return str(snapshot_download(**kwargs)) + return str(hf_api().snapshot_download(**kwargs)) def _ensure_remote_package(snapshot_dir: str) -> str: diff --git a/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2audio.py b/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2audio.py index 80bc0a78b47..dbbb635245f 100644 --- a/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2audio.py +++ b/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2audio.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 os @@ -80,7 +83,7 @@ def _ensure_remote_s2u_vendor_root( return str(existing_path) try: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api except Exception as e: logger.warning("huggingface_hub unavailable; cannot fetch s2u_vendor from %s: %s", repo_id, e) return None @@ -91,7 +94,7 @@ def _ensure_remote_s2u_vendor_root( for revision in revisions: try: - snapshot_dir = snapshot_download( + snapshot_dir = hf_api().snapshot_download( repo_id=repo_id, revision=revision, allow_patterns=["s2u_vendor/**"], @@ -99,7 +102,7 @@ def _ensure_remote_s2u_vendor_root( ) except TypeError: try: - snapshot_dir = snapshot_download( + snapshot_dir = hf_api().snapshot_download( repo_id=repo_id, revision=revision, allow_patterns=["s2u_vendor/**"], diff --git a/vllm_omni/model_executor/models/glm_tts/glm_tts.py b/vllm_omni/model_executor/models/glm_tts/glm_tts.py index 664a794c4c0..fb3517036d0 100644 --- a/vllm_omni/model_executor/models/glm_tts/glm_tts.py +++ b/vllm_omni/model_executor/models/glm_tts/glm_tts.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 """GLM-TTS AR Model (Stage 0): Text → Speech Tokens. Based on Llama architecture, generates speech token sequences from input text. @@ -113,9 +113,9 @@ def resolve_glm_tts_tokenizer_path(model_name_or_path: Any) -> str: return candidate return model_path - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - local_dir = snapshot_download( + local_dir = hf_api().snapshot_download( model_path, allow_patterns=[ f"{_GLM_TTS_TOKENIZER_SUBDIR}/tokenizer*", @@ -169,9 +169,9 @@ def _has_files(root: str) -> bool: except Exception: pass - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - return snapshot_download(model_name_or_path) + return hf_api().snapshot_download(model_name_or_path) def _first_glm_tts_value(value: Any) -> Any: diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py index 91609e1a4ee..c6a54a39350 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py +++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.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 """Stage 1 (codec decoder) for higgs-audio v2. Two surfaces exposed: @@ -231,9 +231,9 @@ def _resolve_hf_id_to_local(repo_id: str) -> str: if os.path.isfile(os.path.join(candidate, "config.json")): return candidate - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - return snapshot_download(repo_id=repo_id) + return hf_api().snapshot_download(repo_id=repo_id) # ------------------------------------------------------ direct decode API @torch.inference_mode() diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py index c3f5024b3bc..587bcda6718 100755 --- a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py +++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.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 """TTS prompt builder + scope validators for higgs-audio v2. vllm-omni's higgs path supports two request shapes: @@ -248,9 +248,9 @@ def _load_audio_tokenizer(): audio_tokenizer_dir = _resolve_audio_tokenizer_dir() if audio_tokenizer_dir is None: - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - repo_path = snapshot_download( + repo_path = hf_api().snapshot_download( _K2_OMNIVOICE_REPO, allow_patterns=[f"{_K2_OMNIVOICE_SUBDIR}/*"], ) diff --git a/vllm_omni/model_executor/models/indextts2/preprocess_utils.py b/vllm_omni/model_executor/models/indextts2/preprocess_utils.py index 88e9a063952..82729160270 100644 --- a/vllm_omni/model_executor/models/indextts2/preprocess_utils.py +++ b/vllm_omni/model_executor/models/indextts2/preprocess_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 """External model loading, audio I/O, and emotion conditioning for IndexTTS2.""" from __future__ import annotations @@ -171,9 +171,10 @@ def load_semantic_codec( codec.load_state_dict(state, strict=False) else: import safetensors.torch - from huggingface_hub import hf_hub_download - ckpt_path = hf_hub_download("amphion/MaskGCT", filename="semantic_codec/model.safetensors") + from vllm_omni.transformers_utils.repo_utils import hf_api + + ckpt_path = hf_api().hf_hub_download("amphion/MaskGCT", filename="semantic_codec/model.safetensors") safetensors.torch.load_model(codec, ckpt_path) _freeze(codec.to(device=device, dtype=torch.float32)) _semantic_codec_cache[cache_key] = codec @@ -198,9 +199,11 @@ def load_campplus(model_path: str, device: torch.device): if ckpt_path is not None: break if ckpt_path is None: - from huggingface_hub import hf_hub_download + if os.path.isdir(model_path): + raise FileNotFoundError(f"IndexTTS CAMPPlus checkpoint is missing from local bundle {model_path!r}") + from vllm_omni.transformers_utils.repo_utils import hf_api - ckpt_path = hf_hub_download("funasr/campplus", filename="campplus_cn_common.bin") + ckpt_path = hf_api().hf_hub_download("funasr/campplus", filename="campplus_cn_common.bin") state = torch.load(ckpt_path, map_location="cpu", weights_only=True) campplus.load_state_dict(state, strict=False) _freeze(campplus.to(device=device, dtype=torch.float32)) diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/bigvgan.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/bigvgan.py index 5876bb14742..048d1de2dcf 100644 --- a/vllm_omni/model_executor/models/indextts2/s2mel/modules/bigvgan.py +++ b/vllm_omni/model_executor/models/indextts2/s2mel/modules/bigvgan.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # Copyright (c) 2024 NVIDIA CORPORATION. # Licensed under the MIT license. @@ -9,13 +12,14 @@ import torch import torch.nn as nn -from huggingface_hub import PyTorchModelHubMixin, hf_hub_download +from huggingface_hub import PyTorchModelHubMixin from torch.nn import Conv1d, ConvTranspose1d from torch.nn.utils import remove_weight_norm, weight_norm from vllm_omni.model_executor.models.common.alias_free_activation import AliasFreeActivation1d from vllm_omni.model_executor.models.common.snake_activation import Snake, SnakeBeta from vllm_omni.model_executor.models.indextts2.s2mel.modules.commons import AttrDict +from vllm_omni.transformers_utils.repo_utils import hf_api # --------------------------------------------------------------------------- # Helpers (inlined from env.py / utils.py) @@ -328,7 +332,7 @@ def _from_pretrained(cls, *, model_id: str, map_location: str = "cpu", **kwargs) def _resolve(filename): if os.path.isdir(model_id): return os.path.join(model_id, filename) - return hf_hub_download( + return hf_api().hf_hub_download( repo_id=model_id, filename=filename, **{ diff --git a/vllm_omni/model_executor/models/indextts2/tokenizer_v2_5.py b/vllm_omni/model_executor/models/indextts2/tokenizer_v2_5.py index 6eae97703a7..ba27cad6243 100644 --- a/vllm_omni/model_executor/models/indextts2/tokenizer_v2_5.py +++ b/vllm_omni/model_executor/models/indextts2/tokenizer_v2_5.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 """IndexTTS 2.5 multilingual tiktoken vocabulary. The vocabulary layout mirrors the official IndexTTS 2.5 tokenizer: 58,836 @@ -234,9 +234,9 @@ def resolve_indextts25_tokenizer_file( return str(model_path) try: - from huggingface_hub import hf_hub_download + from vllm_omni.transformers_utils.repo_utils import hf_api - return hf_hub_download(repo_id=model_dir, filename=tokenizer_file) + return hf_api().hf_hub_download(repo_id=model_dir, filename=tokenizer_file) except Exception as exc: raise FileNotFoundError( f"Could not resolve IndexTTS 2.5 tokenizer {tokenizer_file!r} from {model_dir!r}" diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py index 8b3926cc0ca..f701e8b2699 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py @@ -39,9 +39,9 @@ def _resolve_model_dir(model_ref: str, revision: str | None = None) -> str: """ if Path(model_ref).is_dir(): return model_ref - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - return snapshot_download(model_ref, revision=revision, allow_patterns=["assets/*"]) + return hf_api().snapshot_download(model_ref, revision=revision, allow_patterns=["assets/*"]) def _batch_error(reason: str, **details: Any) -> RuntimeError: diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice.py b/vllm_omni/model_executor/models/omnivoice/omnivoice.py index c167aa1c0a1..26a6f6b5935 100644 --- a/vllm_omni/model_executor/models/omnivoice/omnivoice.py +++ b/vllm_omni/model_executor/models/omnivoice/omnivoice.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 """ OmniVoice model for vLLM-Omni two-stage TTS pipeline. @@ -487,9 +487,9 @@ def _resolve_model_dir(self) -> str: if os.path.isdir(model_dir): return model_dir # HF hub model ID — resolve to local cache - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - return snapshot_download(model_dir) + return hf_api().snapshot_download(model_dir) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: try: diff --git a/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py b/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py index 48df20fab1e..3b8813faf70 100644 --- a/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py +++ b/vllm_omni/model_executor/models/personaplex/personaplex_mimi.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Streaming Mimi codec for PersonaPlex (moshi-free). Frame-clocked duplex needs a codec that encodes/decodes exactly one 80 ms frame @@ -288,10 +288,11 @@ class PersonaPlexMimiCodec(nn.Module): def __init__(self, hf_repo: str = DEFAULT_HF_REPO, checkpoint: str | None = None, device: str = "cuda") -> None: super().__init__() - from huggingface_hub import hf_hub_download from safetensors.torch import load_file from transformers import MimiConfig, MimiModel + from vllm_omni.transformers_utils.repo_utils import hf_api + self.device = torch.device(device) # The PersonaPlex repo ships the reference mimi checkpoint in the moshi @@ -300,7 +301,7 @@ def __init__(self, hf_repo: str = DEFAULT_HF_REPO, checkpoint: str | None = None # transformer stacks remain absent because the streaming implementations # below replace them with the checkpoint's fused QKV layout. if checkpoint is None: - checkpoint = hf_hub_download( + checkpoint = hf_api().hf_hub_download( "nvidia/personaplex-7b-v1", "tokenizer-e351c8d8-checkpoint125.safetensors", ) diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py index e9ca78b21eb..8e3c56b327d 100644 --- a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py @@ -461,9 +461,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): if model_path is None: # Resolve HF repo names to local cache path if not os.path.isdir(model_name_or_path): - from huggingface_hub import snapshot_download + from vllm_omni.transformers_utils.repo_utils import hf_api - model_name_or_path = snapshot_download(model_name_or_path) + model_name_or_path = hf_api().snapshot_download(model_name_or_path) model_path = f"{model_name_or_path}/token2wav" float16 = getattr(self.config, "token2wav_float16", False) diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py index 69cfa6d2667..c61aeaaadbb 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from collections.abc import Iterable, Mapping from dataclasses import replace from functools import cached_property @@ -6,7 +9,6 @@ import regex as re import torch import torch.nn as nn -from huggingface_hub import hf_hub_download from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.models.interfaces import SupportsMultiModal @@ -28,6 +30,7 @@ VoxtralTTSMultiModalProcessor, VoxtralTTSProcessingInfo, ) +from vllm_omni.transformers_utils.repo_utils import hf_api logger = init_logger(__name__) @@ -133,7 +136,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.voice_to_embedding = {} for sid in speaker_id: if self.is_hf_model: - path = hf_hub_download(repo_id=self.repo_id, filename=f"voice_embedding/{sid}.pt") + path = hf_api().hf_hub_download(repo_id=self.repo_id, filename=f"voice_embedding/{sid}.pt") else: path = Path(self.repo_id) / "voice_embedding" / f"{sid}.pt" if Path(path).exists(): diff --git a/vllm_omni/transformers_utils/repo_utils.py b/vllm_omni/transformers_utils/repo_utils.py new file mode 100644 index 00000000000..e780c4b53fc --- /dev/null +++ b/vllm_omni/transformers_utils/repo_utils.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Utilities for model repo interaction.""" + +from huggingface_hub import HfApi + +from vllm_omni.version import __version__ as VLLM_OMNI_VERSION + +_hf_api: HfApi | None = None + + +def hf_api() -> HfApi: + """Return a shared HfApi instance tagged with vLLM-Omni's library info.""" + global _hf_api + if _hf_api is None: + _hf_api = HfApi( + library_name="vllm-omni", + library_version=VLLM_OMNI_VERSION, + ) + return _hf_api From 5e4baea64d0f4cd1b6ce147e4f7a28f0f41d995f Mon Sep 17 00:00:00 2001 From: Lei Ke <1141466880@qq.com> Date: Wed, 9 Sep 2026 00:08:46 +0800 Subject: [PATCH 06/93] [Bugfix][NPU] Limit MiniMax H3 modulation grid size (#6794) Signed-off-by: KrystalRay Co-authored-by: KrystalRay --- .buildkite/npu/test-npu-ready.yml | 1 + .../minimax_h3/test_minimax_h3_modulation.py | 142 +++++++++++++++++- .../attention/ops/minimax_h3_modulation.py | 48 +++++- 3 files changed, 182 insertions(+), 9 deletions(-) diff --git a/.buildkite/npu/test-npu-ready.yml b/.buildkite/npu/test-npu-ready.yml index 1f25ebac161..01fb328790a 100644 --- a/.buildkite/npu/test-npu-ready.yml +++ b/.buildkite/npu/test-npu-ready.yml @@ -20,6 +20,7 @@ steps: tests/diffusion/layers/test_activation.py tests/diffusion/layers/test_fused_qk_norm_rope_npu.py tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py + tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py -m 'core_model and npu and A3' --run-level 'core_model' diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py index dbeeaa66a68..e10fbe56a9a 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py @@ -5,14 +5,154 @@ import torch from vllm.triton_utils import HAS_TRITON +from tests.helpers.mark import hardware_marks from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import ( + _MAX_1D_GRID_SIZE, + _iter_row_chunks, + _launch_row_chunks, + indexed_gate, indexed_gate_rms_norm_scale_shift, + indexed_scale_shift_, rms_norm_indexed_scale_shift, ) +from vllm_omni.platforms import current_omni_platform -pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda] +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] +@pytest.mark.cpu +@pytest.mark.parametrize( + ("rows", "expected"), + [ + (0, []), + (_MAX_1D_GRID_SIZE, [(0, _MAX_1D_GRID_SIZE)]), + (_MAX_1D_GRID_SIZE + 1, [(0, _MAX_1D_GRID_SIZE), (_MAX_1D_GRID_SIZE, 1)]), + (69888, [(0, _MAX_1D_GRID_SIZE), (_MAX_1D_GRID_SIZE, 4353)]), + ], +) +def test_row_chunks_respect_ascend_grid_limit(rows, expected): + assert list(_iter_row_chunks(rows)) == expected + + +@pytest.mark.cpu +@pytest.mark.parametrize( + ("device_type", "expected_grids_and_offsets"), + [ + ("npu", [(_MAX_1D_GRID_SIZE, 0), (4353, _MAX_1D_GRID_SIZE)]), + ("cuda", [(69888, 0)]), + ], +) +def test_launch_row_chunks_is_npu_only(device_type, expected_grids_and_offsets): + class RecordingKernel: + def __init__(self): + self.calls = [] + + def __getitem__(self, grid): + def launch(*args, **kwargs): + self.calls.append((grid, args, kwargs)) + + return launch + + kernel = RecordingKernel() + first_arg = object() + second_arg = object() + + _launch_row_chunks(kernel, 69888, device_type, first_arg, second_arg, block_n=16) + + assert kernel.calls == [ + ((grid,), (first_arg, second_arg, row_offset), {"block_n": 16}) + for grid, row_offset in expected_grids_and_offsets + ] + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="requires Ascend NPU") +@pytest.mark.parametrize( + "device", + [pytest.param("npu", marks=hardware_marks(res={"npu": "A3"}, num_cards=1))], +) +def test_modulation_wrappers_match_reference_above_grid_limit(device): + rows = 69888 + hidden_size = 16 + num_conditions = 7 + eps = 1e-6 + dtype = torch.bfloat16 + + torch.manual_seed(0) + x = torch.randn(rows, hidden_size, dtype=dtype) + residual = torch.randn_like(x) + branch = torch.randn_like(x) + weight = torch.randn(hidden_size, dtype=dtype) + shift = torch.randn(num_conditions, hidden_size, dtype=dtype) + scale = torch.randn_like(shift) + gate = torch.randn_like(shift) + indices = torch.arange(rows) % num_conditions + + indexed_shift = shift.index_select(0, indices) + indexed_scale = scale.index_select(0, indices) + indexed_gate_values = gate.index_select(0, indices) + + x_float = x.float() + branch_float = branch.float() + indexed_shift_float = indexed_shift.float() + indexed_scale_float = indexed_scale.float() + indexed_gate_float = indexed_gate_values.float() + + expected_scale_shift = (x_float * (1.0 + indexed_scale_float) + indexed_shift_float).to(dtype) + expected_gate = (x_float + indexed_gate_float * branch_float).to(dtype) + + normalized = x_float + normalized = normalized * torch.rsqrt(normalized.pow(2).mean(-1, keepdim=True) + eps) + normalized = normalized * weight.float() + expected_rms = (normalized * (1.0 + indexed_scale_float) + indexed_shift_float).to(dtype) + + updated_residual = residual.float() + indexed_gate_float * branch_float + expected_residual = updated_residual.to(dtype) + normalized_residual = updated_residual + normalized_residual = normalized_residual * torch.rsqrt(normalized_residual.pow(2).mean(-1, keepdim=True) + eps) + normalized_residual = normalized_residual * weight.float() + expected_modulated = (normalized_residual * (1.0 + indexed_scale_float) + indexed_shift_float).to(dtype) + + npu_x = x.to(device) + npu_residual = residual.to(device) + npu_branch = branch.to(device) + npu_weight = weight.to(device) + npu_shift = shift.to(device) + npu_scale = scale.to(device) + npu_gate = gate.to(device) + npu_indices = indices.to(device) + + actual_scale_shift = indexed_scale_shift_(npu_x.clone(), npu_shift, npu_scale, npu_indices) + actual_gate = indexed_gate(npu_x, npu_gate, npu_branch, npu_indices) + actual_rms = rms_norm_indexed_scale_shift( + npu_x, + npu_weight, + npu_shift, + npu_scale, + npu_indices, + eps, + ) + actual_residual, actual_modulated = indexed_gate_rms_norm_scale_shift( + npu_residual, + npu_gate, + npu_branch, + npu_weight, + npu_shift, + npu_scale, + npu_indices, + eps, + ) + + for actual, expected in ( + (actual_scale_shift, expected_scale_shift), + (actual_gate, expected_gate), + (actual_rms, expected_rms), + (actual_residual, expected_residual), + (actual_modulated, expected_modulated), + ): + torch.testing.assert_close(actual.cpu(), expected, atol=2e-2, rtol=2e-2) + + +@pytest.mark.cuda @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.skipif(not HAS_TRITON, reason="Triton required") def test_fused_modulation_preserves_bf16_residual_boundary() -> None: diff --git a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py b/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py index f758338312f..20327b7e945 100644 --- a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py +++ b/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py @@ -7,6 +7,22 @@ import torch from vllm.triton_utils import tl, triton +# Ascend CANN limits the launch coreDim to uint16. +_MAX_1D_GRID_SIZE = 65535 + + +def _iter_row_chunks(rows: int): + for row_offset in range(0, rows, _MAX_1D_GRID_SIZE): + yield row_offset, min(rows - row_offset, _MAX_1D_GRID_SIZE) + + +def _launch_row_chunks(kernel, rows: int, device_type: str, *args, **kwargs) -> None: + if device_type != "npu": + kernel[(rows,)](*args, 0, **kwargs) + return + for row_offset, chunk_rows in _iter_row_chunks(rows): + kernel[(chunk_rows,)](*args, row_offset, **kwargs) + @triton.jit def _indexed_scale_shift_kernel( @@ -20,9 +36,10 @@ def _indexed_scale_shift_kernel( stride_shift_row, stride_scale_row, stride_indices, + row_offset, block_n: tl.constexpr, ): - row = tl.program_id(0) + row = tl.program_id(0) + row_offset columns = tl.arange(0, block_n) mask = columns < hidden_size index = tl.load(indices_ptr + row * stride_indices) @@ -51,9 +68,10 @@ def _indexed_gate_kernel( stride_gate_row, stride_other_row, stride_indices, + row_offset, block_n: tl.constexpr, ): - row = tl.program_id(0) + row = tl.program_id(0) + row_offset columns = tl.arange(0, block_n) mask = columns < hidden_size index = tl.load(indices_ptr + row * stride_indices) @@ -83,9 +101,10 @@ def _rms_norm_indexed_scale_shift_kernel( stride_shift_row, stride_scale_row, stride_indices, + row_offset, block_n: tl.constexpr, ): - row = tl.program_id(0) + row = tl.program_id(0) + row_offset columns = tl.arange(0, block_n) mask = columns < hidden_size index = tl.load(indices_ptr + row * stride_indices) @@ -123,9 +142,10 @@ def _indexed_gate_rms_norm_scale_shift_kernel( stride_shift_row, stride_scale_row, stride_indices, + row_offset, block_n: tl.constexpr, ): - row = tl.program_id(0) + row = tl.program_id(0) + row_offset columns = tl.arange(0, block_n) mask = columns < hidden_size index = tl.load(indices_ptr + row * stride_indices) @@ -162,7 +182,10 @@ def indexed_scale_shift_( rows, hidden_size = x.shape if rows == 0: return x - _indexed_scale_shift_kernel[(rows,)]( + _launch_row_chunks( + _indexed_scale_shift_kernel, + rows, + x.device.type, x, x, shift, @@ -192,7 +215,10 @@ def indexed_gate( rows, hidden_size = x.shape if rows == 0: return output - _indexed_gate_kernel[(rows,)]( + _launch_row_chunks( + _indexed_gate_kernel, + rows, + x.device.type, output, x, gate, @@ -229,7 +255,10 @@ def rms_norm_indexed_scale_shift( output = torch.empty_like(x) rows, hidden_size = x.shape if rows: - _rms_norm_indexed_scale_shift_kernel[(rows,)]( + _launch_row_chunks( + _rms_norm_indexed_scale_shift_kernel, + rows, + x.device.type, output, x, weight, @@ -274,7 +303,10 @@ def indexed_gate_rms_norm_scale_shift( modulated_out = torch.empty_like(residual) rows, hidden_size = residual.shape if rows: - _indexed_gate_rms_norm_scale_shift_kernel[(rows,)]( + _launch_row_chunks( + _indexed_gate_rms_norm_scale_shift_kernel, + rows, + residual.device.type, residual_out, modulated_out, residual, From b3dd45874a750f7edfa39bb02262804228e5ff7b Mon Sep 17 00:00:00 2001 From: Tianyao Wu <54675599+twu3202@users.noreply.github.com> Date: Wed, 9 Sep 2026 03:27:38 +0800 Subject: [PATCH 07/93] [Bugfix] Build the forced-aligner prompt without a chat template (word timestamps one bin late) (#7240) Signed-off-by: Tianyao Wu --- .../utils/test_qwen3_force_align_processor.py | 16 +++++++++++++- .../deploy/qwen3_tts_forced_aligner.yaml | 5 +++-- .../utils/qwen3_force_align_processor.py | 22 +++++++++---------- 3 files changed, 29 insertions(+), 14 deletions(-) diff --git a/tests/utils/test_qwen3_force_align_processor.py b/tests/utils/test_qwen3_force_align_processor.py index 4a928d0e509..f846ac7a807 100644 --- a/tests/utils/test_qwen3_force_align_processor.py +++ b/tests/utils/test_qwen3_force_align_processor.py @@ -10,7 +10,21 @@ def test_build_prompt_has_boundary_timestamp_markers(): assert prompt.count("") == 4 assert "helloworld" in prompt - assert prompt.endswith("<|im_start|>assistant\n") + # Official aligner format: audio placeholder + words, no chat template + # (a leading "<|im_start|>user\n" shifts the predicted markers one bin late). + assert prompt == f"{processor.AUDIO_PLACEHOLDER}helloworld" + assert "<|im_start|>" not in prompt + assert processor.build_prompt([]) == f"{processor.AUDIO_PLACEHOLDER}" + + +def test_build_prompt_matches_official_encode_timestamp(): + pytest.importorskip("qwen_asr") + from qwen_asr.inference.qwen3_forced_aligner import Qwen3ForceAlignProcessor + + text = "It's 3 o'clock, 你好 world." + word_list, official_prompt = Qwen3ForceAlignProcessor().encode_timestamp(text, "english") + + assert processor.build_prompt(list(word_list)) == official_prompt @pytest.mark.parametrize( diff --git a/vllm_omni/deploy/qwen3_tts_forced_aligner.yaml b/vllm_omni/deploy/qwen3_tts_forced_aligner.yaml index 2aecdbf9b3d..b600ff1e15a 100644 --- a/vllm_omni/deploy/qwen3_tts_forced_aligner.yaml +++ b/vllm_omni/deploy/qwen3_tts_forced_aligner.yaml @@ -5,8 +5,9 @@ # --forced-aligner sets just the model path. # # The stage runs the model with runner="pooling" + pooling_task="token_classify". -# Prompt tokens, chat template and word segmentation are Qwen-specific and live -# in qwen3_force_align_processor (not configurable here). +# Prompt tokens and word segmentation are Qwen-specific and live in +# qwen3_force_align_processor (not configurable here); the prompt is the same +# string qwen_asr's encode_timestamp builds, with no chat template. forced_aligner: model: Qwen/Qwen3-ForcedAligner-0.6B diff --git a/vllm_omni/utils/qwen3_force_align_processor.py b/vllm_omni/utils/qwen3_force_align_processor.py index e70d270dbb9..76c5b9084e7 100644 --- a/vllm_omni/utils/qwen3_force_align_processor.py +++ b/vllm_omni/utils/qwen3_force_align_processor.py @@ -65,18 +65,18 @@ def build_prompt(words: list[str]) -> str: - """Wrap segmented words in the Qwen3 aligner prompt template. - - Each word gets two trailing ```` markers (start + end); the - model classifies each marker into an audio time bin. + """Build the Qwen3 aligner prompt exactly as qwen_asr's + ``Qwen3ForceAlignProcessor.encode_timestamp`` does: the audio placeholder + followed by the words, each with two trailing ```` markers + (start + end) that the model classifies into audio time bins. + + No chat template. The official aligner feeds this string straight to the + tokenizer; wrapping it in ``<|im_start|>user ... <|im_start|>assistant`` + puts three tokens in front of the audio and shifts almost every predicted + marker one 80 ms bin later than the official output. """ - if not words: - # Pad with one timestamp so the decoder always has something to - # read; an empty result still surfaces as "[]" upstream. - body = TIMESTAMP_TOKEN - else: - body = f"{TIMESTAMP_TOKEN}{TIMESTAMP_TOKEN}".join(words) + f"{TIMESTAMP_TOKEN}{TIMESTAMP_TOKEN}" - return f"<|im_start|>user\n{AUDIO_PLACEHOLDER}{body}<|im_end|>\n<|im_start|>assistant\n" + body = f"{TIMESTAMP_TOKEN}{TIMESTAMP_TOKEN}".join(words) + f"{TIMESTAMP_TOKEN}{TIMESTAMP_TOKEN}" + return f"{AUDIO_PLACEHOLDER}{body}" # --- word segmentation (port of Qwen3ForceAlignProcessor) ---------------- From b6f4df320dc6ec187d11b3df2ecff14bd0fe5849 Mon Sep 17 00:00:00 2001 From: Anjie Hou <149605198+specture724@users.noreply.github.com> Date: Wed, 9 Sep 2026 08:17:38 +0800 Subject: [PATCH 08/93] [Refactor][Diffusion] Resolve offload topology through one plan resolver (#7209) Signed-off-by: specture724 Co-authored-by: Claude Opus 5 --- docs/design/feature/offloader/README.md | 35 +- .../test_distributed_layerwise_backend.py | 28 +- .../offloader/test_layerwise_backend.py | 3 +- .../diffusion/offloader/test_plan_resolver.py | 458 ++++++++++++++++++ vllm_omni/diffusion/offloader/__init__.py | 10 + .../diffusion/offloader/component_utils.py | 96 ++-- .../distributed_layerwise_backend.py | 142 ++---- .../diffusion/offloader/layerwise_backend.py | 89 ++-- vllm_omni/diffusion/offloader/offload_plan.py | 10 +- .../diffusion/offloader/plan_resolver.py | 330 +++++++++++++ .../diffusion/offloader/sequential_backend.py | 70 ++- 11 files changed, 978 insertions(+), 293 deletions(-) create mode 100644 tests/diffusion/offloader/test_plan_resolver.py create mode 100644 vllm_omni/diffusion/offloader/plan_resolver.py diff --git a/docs/design/feature/offloader/README.md b/docs/design/feature/offloader/README.md index 7b767f41f1f..96cba8f600c 100644 --- a/docs/design/feature/offloader/README.md +++ b/docs/design/feature/offloader/README.md @@ -53,6 +53,11 @@ must use distinct hook names and remove only hooks they own. ## Discovery and topology +`resolve_offload_plan(pipeline, config)` is the single topology entry point. +Every backend calls it once when `enable()` starts and then consumes the +resolved components, block stacks, resident heads, and staged-residency flags +it returns. Backends do not re-read model declarations. + Pipeline component discovery prefers `SupportsComponentDiscovery` declarations: - `_dit_modules`; @@ -61,16 +66,26 @@ Pipeline component discovery prefers `SupportsComponentDiscovery` declarations: - `_resident_modules`. Dotted paths are supported. Legacy pipelines may use the fallback scan of -well-known attribute names, but new integrations should declare components -explicitly. - -Both layerwise backends first consume pipeline `OffloadPlan` metadata. A -plan's `block_attrs` maps each DiT path to its ordered block containers, while -`encoder_block_attrs` declares streamable encoder stacks. DiTs absent from the -plan fall back to `_layerwise_offload_blocks_attrs` (including the deprecated -singular-name compatibility path). Discovery metadata describes structure -only; the backend remains responsible for transfer, synchronization, and -storage ownership. +well-known attribute names; it now warns once per pipeline class, and new +integrations must declare components explicitly. + +Block topology comes from the pipeline `OffloadPlan` first: `block_attrs` maps +each DiT path to its ordered block containers, `encoder_block_attrs` declares +streamable encoder stacks, `on_demand_component_paths` marks pipeline-managed +residency, `resident_dit_paths` marks the DiTs that may hold resident layers, +and `encoder_dlo_weight_replication` marks the encoders whose loader-produced +weights are safe for AllGather. DiTs absent from the plan fall back to +`_layerwise_offload_blocks_attrs` (including the deprecated singular-name +compatibility path). + +The resolver is pure: it moves no tensor, installs no hook, writes no module +attribute, and reads no process group — multi-rank facts come from +`OffloadConfig`. It owns topology validation, so a configuration the model +cannot serve fails before the first component is placed. Explicit component +selection turns every declaration problem into an error, while the +compatibility topology (no component selector) keeps the historical +warn-and-skip behavior. Metadata still describes structure only; the backend +remains responsible for transfer, synchronization, and storage ownership. ## Cross-strategy invariants diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index 05d7d93b583..c246394342e 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -52,6 +52,7 @@ OffloadPlan, get_offload_plan, ) +from vllm_omni.diffusion.offloader.plan_resolver import resolve_offload_plan from vllm_omni.diffusion.offloader.startup import OffloadStartupState, attach_offload_startup_state from vllm_omni.host_weight_runtime import MappedHostRegion from vllm_omni.platforms import current_omni_platform @@ -2813,18 +2814,21 @@ def test_encoder_allgather_rejects_undeclared_replication(self, patched_offload_ def test_encoder_allgather_rejects_stub_rank_before_block_discovery(self): """Every rank must reject an unsafe encoder group, including stub ranks.""" - backend = DistributedLayerwiseOffloadBackend( - OffloadConfig( - strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, - pin_cpu_memory=False, - dp_size=2, - components=frozenset({"text_encoder"}), - dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"}, - ), - torch.device("cpu"), + + class StubEncoderPipeline(nn.Module): + _offload_plan = OffloadPlan(encoder_block_attrs={"text_encoder": ("missing.blocks",)}) + + def __init__(self): + super().__init__() + self.text_encoder = nn.Module() + + config = OffloadConfig( + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + pin_cpu_memory=False, + dp_size=2, + components=frozenset({"text_encoder"}), + dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"}, ) - backend.dp_group = object() - stub_plan = OffloadPlan(encoder_block_attrs={"text_encoder": ("missing.blocks",)}) with pytest.raises(ValueError, match="not declared replicated"): - backend._try_layerwise_offload_encoder(nn.Module(), "text_encoder", stub_plan) + resolve_offload_plan(StubEncoderPipeline(), config) diff --git a/tests/diffusion/offloader/test_layerwise_backend.py b/tests/diffusion/offloader/test_layerwise_backend.py index 2dd8804cfbc..099e93f7902 100644 --- a/tests/diffusion/offloader/test_layerwise_backend.py +++ b/tests/diffusion/offloader/test_layerwise_backend.py @@ -488,7 +488,8 @@ def __init__(self): with pytest.raises(ValueError, match="Selected text encoder 'text_encoder_2' requires"): backend.enable(pipeline) - assert not pipeline.text_encoder._omni_layerwise_enabled + # The plan resolver rejects the topology before any encoder is hooked. + assert not getattr(pipeline.text_encoder, "_omni_layerwise_enabled", False) def test_default_selection_preserves_unplanned_auxiliaries(self, patched_offload_runtime): pipeline = _LegacyComponentPipeline() diff --git a/tests/diffusion/offloader/test_plan_resolver.py b/tests/diffusion/offloader/test_plan_resolver.py new file mode 100644 index 00000000000..483ce4060cd --- /dev/null +++ b/tests/diffusion/offloader/test_plan_resolver.py @@ -0,0 +1,458 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Unit tests for the shared offload plan resolver.""" + +from typing import ClassVar + +import pytest +import torch +from torch import nn + +from tests.diffusion.offloader.helpers import ( + _DummyBlock, + _PlainEncoder, + _SingleBlockModel, + _StagedEncoder, + _StagedVAE, + patch_offload_runtime, +) +from vllm_omni.diffusion.offloader import layerwise_backend as layerwise_backend_module +from vllm_omni.diffusion.offloader import plan_resolver as plan_resolver_module +from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy +from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend +from vllm_omni.diffusion.offloader.offload_plan import OffloadPlan +from vllm_omni.diffusion.offloader.plan_resolver import resolve_offload_plan + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + +CPU = torch.device("cpu") + + +@pytest.fixture +def patched_offload_runtime(monkeypatch): + patch_offload_runtime(monkeypatch, layerwise_backend_module.current_omni_platform) + + +def _config( + components: set[str] | None = None, + strategy: OffloadStrategy = OffloadStrategy.LAYER_WISE, + **kwargs, +) -> OffloadConfig: + return OffloadConfig( + strategy=strategy, + pin_cpu_memory=False, + components=None if components is None else frozenset(components), + **kwargs, + ) + + +def _stack_of(component): + assert len(component.stacks) == 1 + return component.stacks[0] + + +class _WanStylePipeline(nn.Module): + """Streamable DiT plus a plan-declared, replicated text encoder.""" + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] + _offload_plan = OffloadPlan( + encoder_component_types={"text_encoder": "text_encoder"}, + encoder_block_attrs={"text_encoder": ("encoder.block",)}, + encoder_dlo_weight_replication=frozenset({"text_encoder"}), + ) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=4) + self.text_encoder = _PlainEncoder() + self.vae = _StagedVAE() + + +class _StagedComponentPipeline(nn.Module): + """MiniMax-H3 shape: two encoder stacks plus pipeline-staged components.""" + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] + _offload_plan = OffloadPlan( + encoder_component_types={"text_encoder": "text_encoder"}, + encoder_block_attrs={"text_encoder": ("vision.blocks", "text_model.layers")}, + on_demand_component_paths=frozenset({"text_encoder", "vae"}), + resident_dit_paths=frozenset({"transformer"}), + ) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=4) + self.text_encoder = _StagedEncoder() + self.vae = _StagedVAE() + + +class _LegacyScanPipeline(nn.Module): + """No SupportsComponentDiscovery declaration and no OffloadPlan.""" + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=3) + self.text_encoder = _PlainEncoder() + + +class TestProducerPrecedence: + def test_plan_block_attrs_win_over_class_attributes(self): + class TwoStackDiT(nn.Module): + _layerwise_offload_blocks_attrs = ["blocks"] + + def __init__(self): + super().__init__() + self.blocks = nn.ModuleList([_DummyBlock() for _ in range(2)]) + self.alt_blocks = nn.ModuleList([_DummyBlock() for _ in range(3)]) + + class Pipeline(nn.Module): + _offload_plan = OffloadPlan(block_attrs={"transformer": ("alt_blocks",)}) + + def __init__(self): + super().__init__() + self.transformer = TwoStackDiT() + + pipeline = Pipeline() + resolved = resolve_offload_plan(pipeline, _config({"dit"})) + + stack = _stack_of(resolved.dits[0]) + assert stack.attrs == ("alt_blocks",) + assert list(stack.blocks) == list(pipeline.transformer.alt_blocks) + + def test_class_attributes_are_used_without_a_plan(self): + pipeline = _LegacyScanPipeline() + resolved = resolve_offload_plan(pipeline, _config({"dit"})) + + assert list(_stack_of(resolved.dits[0]).blocks) == list(pipeline.transformer.blocks) + + def test_declared_component_type_wins_over_leaf_name(self): + class Pipeline(nn.Module): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["mllm"] + _vae_modules: ClassVar[list[str]] = [] + _resident_modules: ClassVar[list[str]] = [] + _offload_plan = OffloadPlan( + encoder_component_types={"mllm": "text_encoder"}, + encoder_block_attrs={"mllm": ("encoder.block",)}, + ) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + self.mllm = _PlainEncoder() + + resolved = resolve_offload_plan(Pipeline(), _config({"text_encoder"})) + + encoder = resolved.encoders[0] + assert encoder.path == "mllm" + assert encoder.selected + assert len(_stack_of(encoder).blocks) == 2 + + def test_component_declaration_wins_over_attribute_scan(self): + class Pipeline(nn.Module): + _dit_modules: ClassVar[list[str]] = ["custom_dit"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = [] + _resident_modules: ClassVar[list[str]] = [] + + def __init__(self): + super().__init__() + self.custom_dit = _SingleBlockModel(num_blocks=2) + # Would be picked up by the deprecated attribute scan. + self.transformer = _SingleBlockModel(num_blocks=2) + + resolved = resolve_offload_plan(Pipeline(), _config({"dit"})) + + assert [component.path for component in resolved.dits] == ["custom_dit"] + + def test_attribute_scan_warns_once_per_pipeline_class(self, monkeypatch): + warnings: list[str] = [] + + class _Recorder: + def warning(self, message, *args): + warnings.append(message % args if args else message) + + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + monkeypatch.setattr(plan_resolver_module, "logger", _Recorder()) + + class UndeclaredPipeline(nn.Module): + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + + resolve_offload_plan(UndeclaredPipeline(), _config({"dit"})) + resolve_offload_plan(UndeclaredPipeline(), _config({"dit"})) + + scan_warnings = [message for message in warnings if "SupportsComponentDiscovery" in message] + assert len(scan_warnings) == 1 + assert "deprecated" in scan_warnings[0] + + +class TestResolvedTopology: + def test_encoder_plan_resolves_one_stack_per_declared_path(self): + pipeline = _WanStylePipeline() + resolved = resolve_offload_plan(pipeline, _config({"dit", "text_encoder"})) + + encoder = resolved.encoders[0] + assert encoder.selected and not encoder.on_demand + assert list(_stack_of(encoder).blocks) == list(pipeline.text_encoder.encoder.block) + assert [component.path for component in resolved.vaes] == ["vae"] + + def test_staged_plan_resolves_every_declared_stack(self): + pipeline = _StagedComponentPipeline() + resolved = resolve_offload_plan(pipeline, _config({"text_encoder"})) + + encoder = resolved.encoders[0] + assert encoder.on_demand + assert [len(stack.blocks) for stack in encoder.stacks] == [2, 2] + # An explicit selection never stages a VAE through the legacy plan. + assert not resolved.vaes[0].on_demand + # DiT is unselected, so it keeps no streaming topology. + assert resolved.dits[0].selected is False + assert resolved.dits[0].stacks == () + + def test_omitted_selector_keeps_the_legacy_staged_vae(self): + resolved = resolve_offload_plan(_StagedComponentPipeline(), _config()) + + assert resolved.vaes[0].on_demand + assert resolved.encoders[0].on_demand + assert resolved.dits[0].selected + + def test_on_demand_only_plan_resolves_without_stacks(self): + class Pipeline(nn.Module): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] + _offload_plan = OffloadPlan(on_demand_component_paths=frozenset({"text_encoder", "vae"})) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + self.text_encoder = _StagedEncoder() + self.vae = _StagedVAE() + + resolved = resolve_offload_plan(Pipeline(), _config({"text_encoder"})) + + assert resolved.encoders[0].on_demand + assert resolved.encoders[0].stacks == () + + def test_module_strategy_resolves_selection_without_stacks(self): + resolved = resolve_offload_plan( + _WanStylePipeline(), + _config({"dit", "text_encoder"}, strategy=OffloadStrategy.MODEL_LEVEL), + ) + + assert [component.path for component in resolved.components] == ["transformer", "text_encoder", "vae"] + assert all(component.stacks == () for component in resolved.components) + assert resolved.encoders[0].selected + + def test_module_strategy_leaves_staged_residency_to_the_pipeline(self): + resolved = resolve_offload_plan( + _StagedComponentPipeline(), + _config({"dit", "text_encoder"}, strategy=OffloadStrategy.MODEL_LEVEL), + ) + + assert not any(component.on_demand for component in resolved.components) + + +class TestResidentLayers: + def _dlo_config(self, resident_layers: int, components: set[str] | None = None) -> OffloadConfig: + return _config( + components, + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + dlo_transfers={"dit": "rank-local", "text_encoder": "rank-local"}, + dlo_resident_layers=resident_layers, + ) + + def test_resident_head_splits_the_declared_dit(self): + pipeline = _StagedComponentPipeline() + resolved = resolve_offload_plan(pipeline, self._dlo_config(2, {"dit"})) + + stack = _stack_of(resolved.dits[0]) + assert stack.resident_head == 2 + assert list(stack.resident) == list(pipeline.transformer.blocks)[:2] + assert list(stack.streaming) == list(pipeline.transformer.blocks)[2:] + + def test_resident_count_beyond_the_stack_keeps_every_block(self): + resolved = resolve_offload_plan(_StagedComponentPipeline(), self._dlo_config(9, {"dit"})) + + stack = _stack_of(resolved.dits[0]) + assert stack.resident_head == 4 + assert stack.streaming == () + + def test_resident_head_leaving_one_streaming_block_is_rejected(self): + with pytest.raises(ValueError, match="leaves only one streaming block"): + resolve_offload_plan(_StagedComponentPipeline(), self._dlo_config(3, {"dit"})) + + def test_legacy_single_streaming_block_stays_resident(self): + resolved = resolve_offload_plan(_StagedComponentPipeline(), self._dlo_config(3)) + + stack = _stack_of(resolved.dits[0]) + assert stack.resident_head == 4 + assert stack.streaming == () + + def test_resident_layers_without_a_declared_path_is_rejected(self): + with pytest.raises(ValueError, match="no matching resident_dit_paths"): + resolve_offload_plan(_WanStylePipeline(), self._dlo_config(2, {"dit"})) + + +class TestValidation: + def test_selected_dit_requires_a_discovered_module(self): + pipeline = nn.Module() + pipeline.text_encoder = _PlainEncoder() + + with pytest.raises(ValueError, match="No DiT/transformer modules found"): + resolve_offload_plan(pipeline, _config({"dit"})) + + def test_model_level_selection_requires_both_swap_sides(self): + pipeline = nn.Module() + pipeline.transformer = _SingleBlockModel(num_blocks=2) + + with pytest.raises(ValueError, match="requires an encoder execution stage"): + resolve_offload_plan(pipeline, _config({"dit"}, strategy=OffloadStrategy.MODEL_LEVEL)) + + def test_selected_text_encoder_requires_a_matching_module(self): + pipeline = nn.Module() + pipeline.transformer = _SingleBlockModel(num_blocks=2) + + with pytest.raises(ValueError, match="No text encoder modules found"): + resolve_offload_plan(pipeline, _config({"text_encoder"})) + + def test_selected_encoder_requires_a_streamable_or_staged_plan(self): + pipeline = nn.Module() + pipeline.transformer = _SingleBlockModel(num_blocks=2) + pipeline.text_encoder = _PlainEncoder() + + with pytest.raises(ValueError, match="Selected text encoder 'text_encoder' requires"): + resolve_offload_plan(pipeline, _config({"text_encoder"})) + + def test_staged_component_requires_the_explicit_lifecycle(self): + class Pipeline(nn.Module): + _offload_plan = OffloadPlan(on_demand_component_paths=frozenset({"text_encoder"})) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + self.text_encoder = _PlainEncoder() + + with pytest.raises(ValueError, match="must implement load_to_device"): + resolve_offload_plan(Pipeline(), _config({"text_encoder"})) + + def test_single_block_dit_is_rejected_for_layer_offload(self): + pipeline = nn.Module() + pipeline.transformer = _SingleBlockModel(num_blocks=1) + + with pytest.raises(ValueError, match="requires at least two streamable"): + resolve_offload_plan(pipeline, _config({"dit"})) + + def test_encoder_allgather_requires_declared_replication(self): + config = _config( + {"text_encoder"}, + strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE, + dp_size=2, + dlo_transfers={"dit": "rank-local", "text_encoder": "allgather"}, + ) + + with pytest.raises(ValueError, match="not declared replicated"): + resolve_offload_plan(_StagedComponentPipeline(), config) + + def test_two_dits_cannot_own_the_same_blocks(self): + class Pipeline(nn.Module): + _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = [] + _resident_modules: ClassVar[list[str]] = [] + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + self.transformer_2 = _SingleBlockModel(num_blocks=2) + self.transformer_2.blocks = self.transformer.blocks + + with pytest.raises(ValueError, match="claimed by both"): + resolve_offload_plan(Pipeline(), _config({"dit"})) + + def test_legacy_selection_keeps_unstreamable_components(self): + pipeline = nn.Module() + pipeline.transformer = _SingleBlockModel(num_blocks=1) + pipeline.text_encoder = _PlainEncoder() + + resolved = resolve_offload_plan(pipeline, _config()) + + assert resolved.dits[0].stacks == () + assert resolved.encoders[0].stacks == () + + +class TestPurity: + def _snapshot(self, pipeline: nn.Module): + return { + "devices": [(name, tensor.device) for name, tensor in pipeline.named_parameters()], + "attrs": {name: sorted(vars(module)) for name, module in pipeline.named_modules()}, + } + + def test_resolution_does_not_touch_the_model(self): + pipeline = _StagedComponentPipeline() + before = self._snapshot(pipeline) + + resolve_offload_plan(pipeline, _config({"dit", "text_encoder"})) + + assert self._snapshot(pipeline) == before + assert pipeline.text_encoder.offload_calls == 0 + assert pipeline.text_encoder.to_calls == 0 + assert pipeline.vae.to_calls == 0 + + def test_resolution_is_repeatable(self): + pipeline = _WanStylePipeline() + config = _config({"dit", "text_encoder"}) + + first = resolve_offload_plan(pipeline, config) + second = resolve_offload_plan(pipeline, config) + + assert [component.path for component in first.components] == [component.path for component in second.components] + for left, right in zip(first.components, second.components, strict=True): + assert left.selected == right.selected + assert left.on_demand == right.on_demand + assert [list(stack.blocks) for stack in left.stacks] == [list(stack.blocks) for stack in right.stacks] + + def test_backend_enable_rejects_before_touching_earlier_components(self, patched_offload_runtime): + class Pipeline(nn.Module): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder", "text_encoder_2"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] + _offload_plan = OffloadPlan( + encoder_component_types={ + "text_encoder": "text_encoder", + "text_encoder_2": "text_encoder", + }, + encoder_block_attrs={"text_encoder": ("encoder.block",)}, + ) + + def __init__(self): + super().__init__() + self.transformer = _SingleBlockModel(num_blocks=2) + self.text_encoder = _PlainEncoder() + self.text_encoder_2 = _PlainEncoder() + self.vae = _StagedVAE() + + pipeline = Pipeline() + backend = LayerWiseOffloadBackend(_config({"text_encoder"}), CPU) + + with pytest.raises(ValueError, match="Selected text encoder 'text_encoder_2' requires"): + backend.enable(pipeline) + + assert not backend.enabled + assert not getattr(pipeline.text_encoder, "_omni_layerwise_enabled", False) + for block in pipeline.text_encoder.encoder.block: + assert getattr(block, "_hook_registry", None) is None + assert pipeline.vae.to_calls == 0 diff --git a/vllm_omni/diffusion/offloader/__init__.py b/vllm_omni/diffusion/offloader/__init__.py index 0df4039ebb7..42ba45cc95b 100644 --- a/vllm_omni/diffusion/offloader/__init__.py +++ b/vllm_omni/diffusion/offloader/__init__.py @@ -20,6 +20,12 @@ from .layerwise_backend import LayerWiseOffloadBackend from .module_residency import BoundedAllocatorCache, PinnedModuleStager from .offload_plan import OffloadPlan, get_offload_plan +from .plan_resolver import ( + BlockStack, + ResolvedComponent, + ResolvedOffloadPlan, + resolve_offload_plan, +) from .sequential_backend import ( ModelLevelOffloadBackend, apply_sequential_offload, @@ -42,6 +48,10 @@ "OffloadConfig", "OffloadPlan", "OffloadStrategy", + "BlockStack", + "ResolvedComponent", + "ResolvedOffloadPlan", + "resolve_offload_plan", "SupportsModelCpuOffload", "LayerWiseOffloadBackend", "DistributedLayerwiseOffloadBackend", diff --git a/vllm_omni/diffusion/offloader/component_utils.py b/vllm_omni/diffusion/offloader/component_utils.py index 1000b29d651..cbc63159d7a 100644 --- a/vllm_omni/diffusion/offloader/component_utils.py +++ b/vllm_omni/diffusion/offloader/component_utils.py @@ -4,7 +4,7 @@ from __future__ import annotations -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Sequence from itertools import chain from operator import attrgetter from typing import TYPE_CHECKING, Any @@ -14,14 +14,12 @@ from torch.distributed.tensor import DTensor from vllm.logger import init_logger -from .block_discovery import get_blocks_from_dit -from .config import DIT_COMPONENT, TEXT_ENCODER_COMPONENT +from .config import TEXT_ENCODER_COMPONENT from .offload_plan import OffloadPlan from .tensor_utils import set_tensor_storage if TYPE_CHECKING: - from .base import OffloadConfig - from .module_collector import PipelineModules + from .plan_resolver import BlockStack, ResolvedComponent, ResolvedOffloadPlan logger = init_logger(__name__) @@ -76,30 +74,25 @@ def get_encoder_block_groups( def iter_streamable_dits( - modules: PipelineModules, - config: OffloadConfig, + resolved: ResolvedOffloadPlan, device: torch.device, - plan: OffloadPlan | None, -) -> Iterator[tuple[str, nn.Module, list[str], list[nn.Module]]]: - """Yield selected DiTs whose block metadata resolves successfully.""" - if not config.offloads(DIT_COMPONENT): - return - for name, module in zip(modules.dit_names, modules.dits): - logger.info("Applying hooks on %s (%s)", name, module.__class__.__name__) - planned_attrs = None if plan is None else plan.block_attrs.get(name) - block_attrs, blocks = get_blocks_from_dit(module, planned_attrs) - if blocks: - yield name, module, block_attrs, blocks +) -> Iterator[tuple[ResolvedComponent, BlockStack]]: + """Yield selected DiTs that resolved to a streamable block stack.""" + for component in resolved.dits: + if not component.selected: continue - if config.components is not None: - raise ValueError(f"Selected DiT {name!r} has no streamable layerwise-offload blocks") - logger.warning("Target layers (blocks) not found. Skipping offloading on %s (%s)", name, type(module).__name__) - module.to(device) + logger.info("Applying hooks on %s (%s)", component.path, component.module.__class__.__name__) + if not component.stacks: + # The resolver already warned; a legacy selection keeps the + # unstreamable DiT resident instead of failing the run. + component.module.to(device) + continue + yield component, component.stacks[0] def move_non_block_state_to_device( module: nn.Module, - block_groups: list[nn.ModuleList], + block_groups: Sequence[Sequence[nn.Module]], device: torch.device, ) -> None: """Keep component state outside streamed block lists resident on device.""" @@ -120,7 +113,7 @@ def move_non_block_state_to_device( def set_encoder_layerwise_state( module: nn.Module, hooks: list[Any], - block_groups: list[nn.ModuleList], + block_groups: Sequence[Sequence[nn.Module]], ) -> None: """Publish the backend-neutral state used by encoder stage lifecycles.""" module._omni_layerwise_hooks = hooks @@ -163,59 +156,38 @@ def prepare_component( def prepare_pipeline_components( - modules: PipelineModules, - config: OffloadConfig, - plan: OffloadPlan | None, + resolved: ResolvedOffloadPlan, *, device: torch.device, staged_components: list[nn.Module], - enable_encoder_blocks: Callable[[nn.Module, str, OffloadPlan | None, bool], bool], + enable_encoder_blocks: Callable[[ResolvedComponent], bool], ) -> None: """Apply the shared encoder/VAE/resident placement policy.""" - if config.components is not None and config.offloads(TEXT_ENCODER_COMPONENT): - selected_encoder_names = [name for name in modules.encoder_names if config.offloads_encoder(name, plan)] - if not selected_encoder_names: - raise ValueError("No text encoder modules found for selected text_encoder offload") - - if plan is not None: - for encoder, name in zip(modules.encoders, modules.encoder_names): - if config.should_offload_encoder(name, plan) and name in plan.on_demand_component_paths: - validate_on_demand_component(encoder, name) - - for encoder, name in zip(modules.encoders, modules.encoder_names): - selected = config.should_offload_encoder(name, plan) - stage_on_demand = bool(selected and plan is not None and name in plan.on_demand_component_paths) - blockwise = selected and enable_encoder_blocks(encoder, name, plan, stage_on_demand) - if stage_on_demand and not blockwise and config.uses_allgather(TEXT_ENCODER_COMPONENT): - raise ValueError( - f"Text encoder {name!r} cannot use AllGather without a model-declared streamable block plan" - ) - if selected and config.components is not None and not (blockwise or stage_on_demand): - raise ValueError(f"Selected text encoder {name!r} requires a model-declared streamable or on-demand plan") + for component in resolved.encoders: + blockwise = bool(component.stacks) and enable_encoder_blocks(component) prepare_component( - encoder, - name, + component.module, + component.path, device=device, - stage_on_demand=stage_on_demand, + stage_on_demand=component.on_demand, blockwise=blockwise, staged_components=staged_components, ) - for vae, name in zip(modules.vaes, modules.vae_names): - legacy_staged = config.components is None and plan is not None and name in plan.on_demand_component_paths + for component in resolved.vaes: prepare_component( - vae, - name, + component.module, + component.path, device=device, - stage_on_demand=legacy_staged, + stage_on_demand=component.on_demand, blockwise=False, staged_components=staged_components, ) - for name, module in zip(modules.resident_names, modules.resident_modules): - module.to(device) - logger.debug("Moved resident module %s to %s", name, device) + for component in resolved.residents: + component.module.to(device) + logger.debug("Moved resident module %s to %s", component.path, device) - if not config.offloads(DIT_COMPONENT): - for dit in modules.dits: - dit.to(device) + for component in resolved.dits: + if not component.selected: + component.module.to(device) diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index e626d87a543..299ec3d3a9f 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py @@ -21,6 +21,7 @@ import threading import time import weakref +from collections.abc import Sequence from itertools import chain from typing import Any @@ -39,7 +40,6 @@ from .base import OffloadBackend, OffloadConfig, run_cleanup_steps from .component_utils import ( clear_encoder_layerwise_state, - get_encoder_block_groups, iter_streamable_dits, move_non_block_state_to_device, prepare_component, @@ -53,11 +53,8 @@ HostRegistrationError, register_host_mappings, ) -from .module_collector import ModuleDiscovery -from .offload_plan import ( - OffloadPlan, - get_offload_plan, -) +from .offload_plan import OffloadPlan +from .plan_resolver import ResolvedComponent, resolve_offload_plan from .tensor_utils import ( clear_block_storage, clear_tensor_storage, @@ -1437,7 +1434,7 @@ def _has_multirank_allgather(self) -> bool: def _install_hook_group( self, - blocks: list[nn.Module] | nn.ModuleList, + blocks: Sequence[nn.Module] | nn.ModuleList, component: str, *, use_dit_mmap: bool = False, @@ -1483,48 +1480,27 @@ def _install_hook_group( hooks[1]._is_group_first = True return hooks - def _try_layerwise_offload_encoder( - self, - module: nn.Module, - name: str, - plan: OffloadPlan | None, - stage_on_demand: bool = False, - ) -> bool: - """Apply DLO hooks to a plan-declared text encoder. + def _try_layerwise_offload_encoder(self, component: ResolvedComponent) -> bool: + """Apply DLO hooks to a text encoder's resolved block stacks. Rank-local transfer is always safe because it preserves the tensors produced by the model loader. Multi-rank AllGather is enabled only when the model declares that those tensors are replicated across the - selected DLO group; encoder TP groups contain different shards and - are deliberately not used for this transport. + selected DLO group; the resolver rejects every other layout before + this method installs anything. """ - if plan is None: - return False - - group_size = self._component_transport(TEXT_ENCODER_COMPONENT)[1] - if group_size > 1 and name not in plan.encoder_dlo_weight_replication: - raise ValueError( - f"Text encoder {name!r} cannot use DLO AllGather across the DiT offload group: " - "its loader-produced weights are not declared replicated across that group. " - "Set layer_options.text_encoder.weight_transfer='rank-local' in diffusion_offload_config " - "for encoder-TP or rank-specific layouts." - ) - block_groups = get_encoder_block_groups( - module, - name, - plan, - strict=self.config.components is not None or group_size > 1, - ) + block_groups = [stack.blocks for stack in component.stacks] if not block_groups: return False + module = component.module encoder_hooks: list[DistributedLayerwiseOffloadHook] = [] for blocks in block_groups: encoder_hooks.extend(self._install_hook_group(blocks, TEXT_ENCODER_COMPONENT)) # Track the module before placement, which may fail, so outer rollback # also removes its partially-installed block hooks and marker state. self._encoder_modules.append(module) - if not stage_on_demand: + if not component.on_demand: move_non_block_state_to_device(module, block_groups, self.device) set_encoder_layerwise_state( module, @@ -1534,10 +1510,10 @@ def _try_layerwise_offload_encoder( logger.info( "Enabled %s DLO transfer for text encoder %s (%d blocks across %d stacks, group_size=%d)", self.config.transfer_for(TEXT_ENCODER_COMPONENT).value, - name, + component.path, sum(len(blocks) for blocks in block_groups), len(block_groups), - group_size, + self._component_transport(TEXT_ENCODER_COMPONENT)[1], ) return True @@ -1700,31 +1676,15 @@ def _enable(self, pipeline: nn.Module) -> None: if self.dp_group is None and self._has_multirank_allgather(): self._init_dp_group() - modules = ModuleDiscovery.discover(pipeline) - if not modules.dits and self.config.offloads(DIT_COMPONENT): + # One resolver call validates the whole topology; an explicit selection + # has already failed by the time any of the checks below run. + resolved = resolve_offload_plan(pipeline, self.config) + if not resolved.dits and self.config.offloads(DIT_COMPONENT): if self.host_weight_plan is not None: raise RuntimeError( "DLO received a loader-owned host-weight plan, but no DiT modules were discovered to consume it" ) - message = "No DiT/transformer modules found for selected DiT offload" - if self.config.components is not None: - raise ValueError(message) - logger.warning(message) - - # Retrieve optional declarative OffloadPlan from the pipeline. - # When present, replaces heuristic block discovery. - plan = get_offload_plan(pipeline) - - if self.config.dlo_resident_layers: - resident_paths = frozenset() if plan is None else plan.resident_dit_paths - if not resident_paths.intersection(modules.dit_names): - message = ( - f"resident_layers={self.config.dlo_resident_layers} was requested, but this model declares " - "no matching resident_dit_paths" - ) - if self.config.components is not None: - raise ValueError(message) - logger.warning("%s; all blocks will be streamed.", message) + logger.warning("No DiT/transformer modules found for selected DiT offload") # Storage selection belongs to the loader. DLO consumes the exact # prevalidated plan that caused the loader to skip materialization; @@ -1755,7 +1715,7 @@ def _enable(self, pipeline: nn.Module) -> None: elif host_weight_plan.backing_kind == "checkpoint_mmap": self._load_weights_via_mmap( pipeline, - modules, + resolved.modules, host_weight_plan, ) else: @@ -1768,10 +1728,10 @@ def _enable(self, pipeline: nn.Module) -> None: else: remaining_meta = [ name - for dit_name, dit_module in zip(modules.dit_names, modules.dits) + for component in resolved.dits for name, tensor in chain( - dit_module.named_parameters(), - dit_module.named_buffers(), + component.module.named_parameters(), + component.module.named_buffers(), ) if getattr(tensor, "is_meta", False) ] @@ -1784,66 +1744,46 @@ def _enable(self, pipeline: nn.Module) -> None: # Apply each selected encoder transfer while keeping explicit VAEs and # unselected components resident. prepare_pipeline_components( - modules, - self.config, - plan, + resolved, device=self.device, staged_components=self._staged_components, enable_encoder_blocks=self._try_layerwise_offload_encoder, ) if self.config.offloads(DIT_COMPONENT): - logger.info("Applying distributed layer-wise offloading on %s", modules.dit_names) + logger.info( + "Applying distributed layer-wise offloading on %s", + [component.path for component in resolved.dits], + ) # Collect all DiT module objects to detect submodules that are # already handled as a separate DiT module (avoids duplicate hooks). - all_dit_modules = set(id(m) for m in modules.dits) + all_dit_modules = set(id(component.module) for component in resolved.dits) # Apply hooks for each DiT module - for dit_name, dit_module, blocks_attr_names, blocks in iter_streamable_dits( - modules, self.config, self.device, plan - ): - resident_count = 0 - if plan is not None and dit_name in plan.resident_dit_paths: - resident_count = min(self.config.dlo_resident_layers, len(blocks)) - if resident_count: - resident_blocks = blocks[:resident_count] - self._resident_blocks.extend(resident_blocks) - blocks = blocks[resident_count:] + for component, stack in iter_streamable_dits(resolved, self.device): + streaming = list(stack.streaming) + if stack.resident_head: + self._resident_blocks.extend(stack.resident) logger.info( "Keeping %d leading blocks resident on %s; streaming %d tail blocks", - resident_count, - dit_name, - len(blocks), + stack.resident_head, + component.path, + len(streaming), ) self._prepare_dit_non_block_modules( - dit_module, - blocks_attr_names, + component.module, + list(stack.attrs), all_dit_modules, - plan, + resolved.declaration, ) - num_blocks = len(blocks) - if num_blocks == 0: - logger.info("All blocks for %s are resident; no streaming hooks required", dit_name) - continue - if num_blocks <= 1: - if self.config.components is not None: - raise ValueError( - f"Selected DiT {dit_name!r} leaves only one streaming block after " - f"resident_layers={resident_count}; choose a resident count that " - "leaves zero or at least two streaming blocks" - ) - logger.warning( - "#Streaming target layers (blocks) <= 1. Keeping the final block resident on %s (%s)", - dit_name, - dit_module.__class__.__name__, - ) - self._resident_blocks.extend(blocks) + if not streaming: + logger.info("All blocks for %s are resident; no streaming hooks required", component.path) continue - self._install_hook_group(blocks, DIT_COMPONENT, use_dit_mmap=True) + self._install_hook_group(streaming, DIT_COMPONENT, use_dit_mmap=True) if self._resident_blocks: self._resident_layer_group = PinnedResidentLayerGroup( self._resident_blocks, diff --git a/vllm_omni/diffusion/offloader/layerwise_backend.py b/vllm_omni/diffusion/offloader/layerwise_backend.py index 9b63f916263..a1e594f8b08 100644 --- a/vllm_omni/diffusion/offloader/layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/layerwise_backend.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations +from collections.abc import Sequence from itertools import chain from typing import Any @@ -20,15 +21,13 @@ ) from .component_utils import ( clear_encoder_layerwise_state, - get_encoder_block_groups, iter_streamable_dits, move_non_block_state_to_device, prepare_pipeline_components, set_encoder_layerwise_state, ) from .config import DIT_COMPONENT -from .module_collector import ModuleDiscovery -from .offload_plan import OffloadPlan, get_offload_plan +from .plan_resolver import ResolvedComponent, resolve_offload_plan from .tensor_utils import ( clear_block_storage, clear_tensor_storage, @@ -279,7 +278,7 @@ def remove_block_hook(module: nn.Module) -> None: def _install_layerwise_hook_group( - blocks: list[nn.Module] | nn.ModuleList, + blocks: Sequence[nn.Module] | nn.ModuleList, device: torch.device, stream: Any, pin_memory: bool, @@ -327,36 +326,29 @@ def _install_layerwise_hook_group( def enable_plan_encoder_layerwise_offload( - module: nn.Module, - name: str, - plan: OffloadPlan | None, + component: ResolvedComponent, *, device: torch.device, stream: current_omni_platform.Stream, pin_memory: bool, - stage_on_demand: bool = False, - strict: bool = False, ) -> bool: - """Apply rank-local layerwise hooks to plan-declared encoder stacks.""" + """Apply rank-local layerwise hooks to the encoder's resolved stacks.""" + module = component.module if getattr(module, "_omni_layerwise_enabled", False): return True - hooks: list[LayerwiseOffloadHook] = [] - hooked_blocks: list[nn.Module] = [] - block_groups = get_encoder_block_groups( - module, - name, - plan, - strict=strict, - ) + block_groups = [stack.blocks for stack in component.stacks] if not block_groups: return False + + hooks: list[LayerwiseOffloadHook] = [] + hooked_blocks: list[nn.Module] = [] try: for blocks in block_groups: group_hooks = _install_layerwise_hook_group(blocks, device, stream, pin_memory) hooks.extend(group_hooks) hooked_blocks.extend(blocks) - if not stage_on_demand: + if not component.on_demand: move_non_block_state_to_device(module, block_groups, device) except BaseException: run_cleanup_steps( @@ -376,7 +368,7 @@ def enable_plan_encoder_layerwise_offload( ) logger.info( "Enabled rank-local layerwise offload for encoder %s (%d blocks across %d stacks)", - name, + component.path, sum(len(blocks) for blocks in block_groups), len(block_groups), ) @@ -440,41 +432,27 @@ def _enable(self, pipeline: nn.Module) -> None: logger.warning("LayerWiseOffloadBackend already enabled") return - modules = ModuleDiscovery.discover(pipeline) - plan = get_offload_plan(pipeline) - if not modules.dits and self.config.offloads(DIT_COMPONENT): - message = "No DiT/transformer modules found for selected DiT layerwise offload" - if self.config.components is not None: - raise ValueError(message) - logger.warning(message) + resolved = resolve_offload_plan(pipeline, self.config) + if not resolved.dits and self.config.offloads(DIT_COMPONENT): + # An explicit selection already failed while resolving. + logger.warning("No DiT/transformer modules found for selected DiT layerwise offload") return - def enable_encoder_blocks( - module: nn.Module, - name: str, - component_plan: OffloadPlan | None, - stage_on_demand: bool, - ) -> bool: + def enable_encoder_blocks(component: ResolvedComponent) -> bool: enabled = enable_plan_encoder_layerwise_offload( - module, - name, - component_plan, + component, device=self.device, stream=self.copy_stream, pin_memory=self.config.pin_cpu_memory, - stage_on_demand=stage_on_demand, - strict=self.config.components is not None, ) if enabled: # Record each successful installation immediately so a later # component failure can remove these hooks transactionally. - self._encoder_modules.append(module) + self._encoder_modules.append(component.module) return enabled prepare_pipeline_components( - modules, - self.config, - plan, + resolved, device=self.device, staged_components=self._staged_components, enable_encoder_blocks=enable_encoder_blocks, @@ -489,30 +467,17 @@ def enable_encoder_blocks( ) return - logger.info("Applying layer-wise offloading on %s", modules.dit_names) + logger.info("Applying layer-wise offloading on %s", [component.path for component in resolved.dits]) # Apply block-wise offloading hook for each of the blocks in DiT model(s) # Note that there might exist multiple DiT models in specific pipelines - for dit_name, dit_module, blocks_attr_names, blocks in iter_streamable_dits( - modules, self.config, self.device, plan - ): - num_blocks = len(blocks) - if num_blocks <= 1: - if self.config.components is not None: - raise ValueError( - f"Selected DiT {dit_name!r} requires at least two streamable layerwise-offload blocks" - ) - logger.warning( - "#Target layers (blocks) <= 1. Skipping offloading on %s (%s)", - dit_name, - dit_module.__class__.__name__, - ) - dit_module.to(self.device) - continue + for component, stack in iter_streamable_dits(resolved, self.device): + dit_module = component.module + blocks = list(stack.blocks) # Move non-block modules to GPU (they stay resident) for name, m in dit_module.named_children(): - if name not in blocks_attr_names: + if name not in stack.attrs: m.to(self.device) logger.debug(f"Moved {name} to device {self.device}") else: @@ -540,7 +505,7 @@ def enable_encoder_blocks( # zero once; later denoising iterations prefetch it from the ring. block_hooks[0].prefetch_layer(non_blocking=False) - logger.info(f"Layer-wise offloading enabled on {num_blocks} layers (blocks)") + logger.info(f"Layer-wise offloading enabled on {len(blocks)} layers (blocks)") # Track hooked blocks for cleanup self._blocks.append(blocks) diff --git a/vllm_omni/diffusion/offloader/offload_plan.py b/vllm_omni/diffusion/offloader/offload_plan.py index 4aa37ae8ff4..5cc0c24c754 100644 --- a/vllm_omni/diffusion/offloader/offload_plan.py +++ b/vllm_omni/diffusion/offloader/offload_plan.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 """Declarative model capabilities for component layerwise offload. Models declare an ``_offload_plan`` class attribute on the pipeline. Both the @@ -22,9 +22,11 @@ class OffloadPlan: pipeline class. When present, both layerwise backends use it instead of model-specific backend branches, making new integrations data-driven. - If not declared, the offloader falls back to: - 1. ``_layerwise_offload_blocks_attrs`` on each DiT module class. - 2. Heuristic search for ``layers`` / ``blocks`` / ``h`` attributes. + If not declared, the offloader falls back to + ``_layerwise_offload_blocks_attrs`` on each DiT module class. + + :func:`~vllm_omni.diffusion.offloader.plan_resolver.resolve_offload_plan` + is the only consumer; backends read the resolved artifact it returns. Attributes: block_attrs: Maps DiT path → tuple of block-list attribute names. diff --git a/vllm_omni/diffusion/offloader/plan_resolver.py b/vllm_omni/diffusion/offloader/plan_resolver.py new file mode 100644 index 00000000000..b2e951b4e33 --- /dev/null +++ b/vllm_omni/diffusion/offloader/plan_resolver.py @@ -0,0 +1,330 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""One topology resolver shared by every diffusion offload backend. + +Model topology is declared through several mechanisms that grew independently: +:class:`~vllm_omni.diffusion.models.interface.SupportsComponentDiscovery` class +variables, the pipeline-level :class:`OffloadPlan`, per-DiT +``_layerwise_offload_blocks_attrs``, and a leaf-name rule for encoders. This +module runs them in one defined order and returns one backend-neutral artifact. + +The resolver is pure: it moves no tensor, installs no hook, writes no module +attribute, and touches no process group. Everything it can reject is rejected +here, so plan-dependent failures happen before a backend mutates the model. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import cache +from typing import TYPE_CHECKING + +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery + +from .block_discovery import get_blocks_from_dit +from .component_utils import get_encoder_block_groups, validate_on_demand_component +from .config import DIT_COMPONENT, TEXT_ENCODER_COMPONENT, OffloadStrategy +from .module_collector import ModuleDiscovery, PipelineModules +from .offload_plan import OffloadPlan, get_offload_plan + +if TYPE_CHECKING: + from .base import OffloadConfig + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class BlockStack: + """One ring of repeated blocks that a backend streams together. + + ``attrs`` names the owner attributes that contributed blocks; backends use + it to skip those children when they place the non-streamed remainder. It is + filled for DiT stacks only — encoder stacks are hooked as whole stacks and + their remainder is placed by tensor identity. + """ + + blocks: tuple[nn.Module, ...] + attrs: tuple[str, ...] = () + resident_head: int = 0 + + @property + def resident(self) -> tuple[nn.Module, ...]: + """Leading blocks that stay on the device instead of streaming.""" + return self.blocks[: self.resident_head] + + @property + def streaming(self) -> tuple[nn.Module, ...]: + """Blocks a hook ring transfers on demand.""" + return self.blocks[self.resident_head :] + + +@dataclass(frozen=True) +class ResolvedComponent: + """One pipeline component with its resolved, validated offload topology. + + ``selected`` means the active selector covers this component. ``on_demand`` + means the pipeline owns its residency through ``load_to_device`` / + ``offload_to_cpu``; a VAE staged by a legacy model plan is ``on_demand`` + without being selectable through the public component grammar. + """ + + path: str + module: nn.Module + selected: bool + stacks: tuple[BlockStack, ...] = () + on_demand: bool = False + + +@dataclass(frozen=True) +class ResolvedOffloadPlan: + """Backend-neutral topology for one pipeline under one offload config.""" + + dits: tuple[ResolvedComponent, ...] + encoders: tuple[ResolvedComponent, ...] + vaes: tuple[ResolvedComponent, ...] + residents: tuple[ResolvedComponent, ...] + # Raw discovery output and the model declaration remain reachable for the + # loader-owned mmap path and for DLO's nested-submodule staging, which + # still interpret topology themselves. + modules: PipelineModules + declaration: OffloadPlan | None = None + + @property + def components(self) -> tuple[ResolvedComponent, ...]: + return self.dits + self.encoders + self.vaes + self.residents + + +@cache +def _warn_legacy_discovery(pipeline_class: str) -> None: + """Warn once per pipeline class that component discovery guessed. + + The key is the fully qualified class name so two same-named pipeline + classes in different modules each get their own warning. + """ + logger.warning( + "%s does not declare SupportsComponentDiscovery; offload components were " + "discovered by scanning well-known attribute names. Declare _dit_modules, " + "_encoder_modules, and _vae_modules; the attribute scan is deprecated.", + pipeline_class, + ) + + +def _resolve_dit_stacks( + module: nn.Module, + path: str, + declaration: OffloadPlan | None, + config: OffloadConfig, + *, + explicit: bool, +) -> tuple[BlockStack, ...]: + """Resolve one DiT's streamed blocks and its device-resident head.""" + planned_attrs = None if declaration is None else declaration.block_attrs.get(path) + attrs, blocks = get_blocks_from_dit(module, planned_attrs) + if not blocks: + if explicit: + raise ValueError(f"Selected DiT {path!r} has no streamable layerwise-offload blocks") + logger.warning( + "Target layers (blocks) not found. Skipping offloading on %s (%s)", + path, + type(module).__name__, + ) + return () + + distributed = config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE + resident_head = 0 + if ( + distributed + and config.dlo_resident_layers + and declaration is not None + and path in declaration.resident_dit_paths + ): + resident_head = min(config.dlo_resident_layers, len(blocks)) + + if len(blocks) - resident_head == 1: + if explicit and distributed: + raise ValueError( + f"Selected DiT {path!r} leaves only one streaming block after " + f"resident_layers={resident_head}; choose a resident count that " + "leaves zero or at least two streaming blocks" + ) + if explicit: + raise ValueError(f"Selected DiT {path!r} requires at least two streamable layerwise-offload blocks") + if distributed: + # One streamed block cannot form a prefetch ring; keeping it + # resident preserves the model instead of skipping placement. + logger.warning( + "#Streaming target layers (blocks) <= 1. Keeping the final block resident on %s (%s)", + path, + type(module).__name__, + ) + resident_head = len(blocks) + else: + logger.warning( + "#Target layers (blocks) <= 1. Skipping offloading on %s (%s)", + path, + type(module).__name__, + ) + return () + + return (BlockStack(blocks=tuple(blocks), attrs=tuple(attrs), resident_head=resident_head),) + + +def _resolve_encoder_stacks( + module: nn.Module, + path: str, + declaration: OffloadPlan | None, + config: OffloadConfig, + *, + explicit: bool, +) -> tuple[BlockStack, ...]: + """Resolve one encoder's streamable stacks and check transfer safety.""" + if declaration is None: + return () + + group_size = 1 + if config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE and config.uses_allgather(TEXT_ENCODER_COMPONENT): + group_size = config.dp_size + if group_size > 1 and path not in declaration.encoder_dlo_weight_replication: + # Every rank must reject an unsafe group, including ranks whose local + # encoder is an unloaded stub, so the check precedes block discovery. + raise ValueError( + f"Text encoder {path!r} cannot use DLO AllGather across the DiT offload group: " + "its loader-produced weights are not declared replicated across that group. " + "Set layer_options.text_encoder.weight_transfer='rank-local' in diffusion_offload_config " + "for encoder-TP or rank-specific layouts." + ) + + groups = get_encoder_block_groups(module, path, declaration, strict=explicit or group_size > 1) + return tuple(BlockStack(blocks=tuple(blocks)) for blocks in groups) + + +def resolve_offload_plan(pipeline: nn.Module, config: OffloadConfig) -> ResolvedOffloadPlan: + """Resolve and validate one pipeline's offload topology. + + Raises ``ValueError`` for every topology that the requested configuration + cannot serve, before the caller places a module or installs a hook. + """ + modules = ModuleDiscovery.discover(pipeline) + declaration = get_offload_plan(pipeline) + if not isinstance(pipeline, SupportsComponentDiscovery) and (modules.dits or modules.encoders or modules.vaes): + pipeline_class = type(pipeline) + _warn_legacy_discovery(f"{pipeline_class.__module__}.{pipeline_class.__qualname__}") + + explicit = config.components is not None + layerwise = config.strategy in (OffloadStrategy.LAYER_WISE, OffloadStrategy.DISTRIBUTED_LAYER_WISE) + dit_selected = config.offloads(DIT_COMPONENT) + + if explicit: + if config.strategy is OffloadStrategy.MODEL_LEVEL: + # Model-level offload swaps the DiT against an encoder stage, so + # both sides must exist regardless of which side was selected. + if not modules.dits: + raise ValueError("Component-selective model offload requires a DiT/transformer module") + if not modules.encoders: + raise ValueError("Component-selective model offload requires an encoder execution stage") + elif dit_selected and not modules.dits: + raise ValueError("No DiT/transformer modules found for selected DiT offload") + if config.offloads(TEXT_ENCODER_COMPONENT) and not any( + config.offloads_encoder(path, declaration) for path in modules.encoder_names + ): + raise ValueError("No text encoder modules found for selected text_encoder offload") + + if config.strategy is OffloadStrategy.DISTRIBUTED_LAYER_WISE and config.dlo_resident_layers: + resident_paths = frozenset() if declaration is None else declaration.resident_dit_paths + if not resident_paths.intersection(modules.dit_names): + message = ( + f"resident_layers={config.dlo_resident_layers} was requested, but this model declares " + "no matching resident_dit_paths" + ) + if explicit: + raise ValueError(message) + logger.warning("%s; all blocks will be streamed.", message) + + dits: list[ResolvedComponent] = [] + encoders: list[ResolvedComponent] = [] + vaes: list[ResolvedComponent] = [] + residents: list[ResolvedComponent] = [] + + for path, module in zip(modules.dit_names, modules.dits): + stacks: tuple[BlockStack, ...] = () + if dit_selected and layerwise: + stacks = _resolve_dit_stacks(module, path, declaration, config, explicit=explicit) + dits.append(ResolvedComponent(path=path, module=module, selected=dit_selected, stacks=stacks)) + + for path, module in zip(modules.encoder_names, modules.encoders): + selected = config.should_offload_encoder(path, declaration) + # Staged residency belongs to the layer backends; model-level offload + # swaps whole modules through its own hooks. + on_demand = bool( + selected and layerwise and declaration is not None and path in declaration.on_demand_component_paths + ) + if on_demand: + validate_on_demand_component(module, path) + stacks = () + if selected and layerwise: + stacks = _resolve_encoder_stacks(module, path, declaration, config, explicit=explicit) + if on_demand and not stacks and config.uses_allgather(TEXT_ENCODER_COMPONENT): + raise ValueError( + f"Text encoder {path!r} cannot use AllGather without a model-declared streamable block plan" + ) + if explicit and not (stacks or on_demand): + raise ValueError( + f"Selected text encoder {path!r} requires a model-declared streamable or on-demand plan" + ) + encoders.append( + ResolvedComponent( + path=path, + module=module, + selected=selected, + stacks=stacks, + on_demand=on_demand, + ) + ) + + for path, module in zip(modules.vae_names, modules.vaes): + # VAEs are not part of the public selector. A model plan may still own + # their residency, which the compatibility topology preserves. + legacy_staged = ( + layerwise and not explicit and declaration is not None and path in declaration.on_demand_component_paths + ) + if legacy_staged: + validate_on_demand_component(module, path) + vaes.append(ResolvedComponent(path=path, module=module, selected=False, on_demand=legacy_staged)) + + for path, module in zip(modules.resident_names, modules.resident_modules): + residents.append(ResolvedComponent(path=path, module=module, selected=False)) + + resolved = ResolvedOffloadPlan( + dits=tuple(dits), + encoders=tuple(encoders), + vaes=tuple(vaes), + residents=tuple(residents), + modules=modules, + declaration=declaration, + ) + _validate_unique_ownership(resolved) + return resolved + + +def _validate_unique_ownership(resolved: ResolvedOffloadPlan) -> None: + """Reject topologies where two components would hook the same blocks.""" + owner_by_block: dict[int, str] = {} + for component in resolved.components: + for block in (block for stack in component.stacks for block in stack.blocks): + owner = owner_by_block.setdefault(id(block), component.path) + if owner != component.path: + raise ValueError( + f"Offload block is claimed by both {owner!r} and {component.path!r}; " + "declare one owner for each block list" + ) + + +__all__ = [ + "BlockStack", + "ResolvedComponent", + "ResolvedOffloadPlan", + "resolve_offload_plan", +] diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index b6b59607089..1187c8cee1c 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/vllm_omni/diffusion/offloader/sequential_backend.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 collections.abc import Collection, Iterator from contextlib import contextmanager @@ -13,9 +13,8 @@ from vllm_omni.platforms import current_omni_platform from .base import OffloadBackend, OffloadConfig, SupportsModelCpuOffload -from .config import DIT_COMPONENT, TEXT_ENCODER_COMPONENT -from .module_collector import ModuleDiscovery -from .offload_plan import get_offload_plan +from .config import DIT_COMPONENT +from .plan_resolver import resolve_offload_plan logger = init_logger(__name__) @@ -329,59 +328,44 @@ def enable(self, pipeline: nn.Module) -> None: ) return - modules = ModuleDiscovery.discover(pipeline) - plan = get_offload_plan(pipeline) - selected_encoders = [ - encoder - for encoder, name in zip(modules.encoders, modules.encoder_names) - if self.config.should_offload_encoder(name, plan) - ] - if self.config.components is not None: - if not modules.dits: - raise ValueError("Component-selective model offload requires a DiT/transformer module") - if not modules.encoders: - raise ValueError("Component-selective model offload requires an encoder execution stage") - if self.config.offloads(TEXT_ENCODER_COMPONENT) and not selected_encoders: - raise ValueError("No text encoder modules found for selected text_encoder module offload") - - all_modules = [ - *modules.dits, - *modules.encoders, - *modules.vaes, - *modules.resident_modules, - ] + resolved = resolve_offload_plan(pipeline, self.config) + dits = [component.module for component in resolved.dits] + encoders = [component.module for component in resolved.encoders] + vaes = [component.module for component in resolved.vaes] + residents = [component.module for component in resolved.residents] + selected_encoders = [component.module for component in resolved.encoders if component.selected] + + all_modules = [*dits, *encoders, *vaes, *residents] initial_devices = _capture_tensor_devices(all_modules) try: - for encoder in modules.encoders: + for encoder in encoders: encoder.to(self.device) - for vae in modules.vaes: + for vae in vaes: vae.to(self.device, non_blocking=True) - for resident in modules.resident_modules: + for resident in residents: resident.to(self.device) - if not modules.dits: + if not dits: logger.warning("No DiT/transformer modules found, skipping model-level offloading") return - if not modules.encoders: - for dit in modules.dits: + if not encoders: + for dit in dits: dit.to(self.device) logger.warning("No encoder modules found, skipping model-level offloading") return apply_sequential_offload( - dit_modules=modules.dits, - encoder_modules=modules.encoders, + dit_modules=dits, + encoder_modules=encoders, device=self.device, pin_memory=self.config.pin_cpu_memory, use_hsdp=self.config.use_hsdp, - offload_dit_modules=( - modules.dits if self.config.components is None or self.config.offloads(DIT_COMPONENT) else () - ), + offload_dit_modules=(dits if self.config.offloads(DIT_COMPONENT) else ()), offload_encoder_modules=selected_encoders, ) except BaseException: try: - remove_sequential_offload([*modules.dits, *modules.encoders]) + remove_sequential_offload([*dits, *encoders]) except BaseException: logger.exception("Failed to remove every model-level hook during rollback") try: @@ -391,15 +375,19 @@ def enable(self, pipeline: nn.Module) -> None: raise # Track modules for cleanup - self._offload_modules = [*modules.dits, *modules.encoders] + self._offload_modules = [*dits, *encoders] self.enabled = True logger.info( "Model-level offloading enabled: %s <-> %s (mutual exclusion)%s", - ", ".join(modules.dit_names), - ", ".join(modules.encoder_names), - f"; resident on GPU: {', '.join(modules.resident_names)}" if modules.resident_names else "", + ", ".join(component.path for component in resolved.dits), + ", ".join(component.path for component in resolved.encoders), + ( + f"; resident on GPU: {', '.join(component.path for component in resolved.residents)}" + if resolved.residents + else "" + ), ) def disable(self) -> None: From f67cc52f7449d5f5dbb9d7790c7a049fa6aa2ba1 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Wed, 9 Sep 2026 09:17:24 +0800 Subject: [PATCH 09/93] [Doc] Add AI usage policy for contributions (#7305) Signed-off-by: hsliuustc0106 --- CONTRIBUTING.md | 2 ++ docs/contributing/README.md | 22 ++++++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 36cefcf0da5..542db3fe9b9 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -2,6 +2,8 @@ You may find information about contributing to vLLM-Omni on [Contributing](https://vllm-omni.readthedocs.io/en/latest/contributing/). +AI-assisted contributions must follow the [AI Usage Policy](docs/contributing/README.md#ai-usage-policy), including disclosure, human review, and validation requirements. + Local `pre-commit` gates (SPDX, forbidden imports, `torch.cuda`, shellcheck, mypy, test marks, TTS adapter ratchet, markdownlint, Buildkite schema, and how to extend allowlists) are documented in the diff --git a/docs/contributing/README.md b/docs/contributing/README.md index dbfc917bfc8..a45854b1429 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -262,6 +262,28 @@ or serving, explain the required L2/L3 environment and provide the exact command Repository skills guide the agent, but they do not replace contributor judgment, the accepted issue or RFC, required test evidence, or maintainer review. Review all generated changes before submitting them. +## AI Usage Policy + +AI tools are welcome in vLLM-Omni contributions. Contributors, including maintainers, remain responsible for everything they submit. These requirements apply to pull requests, issues, discussions, and review comments: + +- **Disclose AI assistance.** Name the tools used and describe their role in the submission, including assistance with code, tests, documentation, or writing. +- **Understand your changes.** Review every generated change and be able to explain its behavior, design choices, and interaction with vLLM-Omni without relying on AI to answer for you. +- **Validate before submitting.** Run relevant checks and report actual results. For model, kernel, or performance changes, include the hardware, model, configuration, and commands needed to reproduce accuracy or benchmark evidence. Clearly identify checks you could not run; generated claims are not test results. +- **Keep communication useful.** Verify technical claims, search for existing reports, and edit AI-assisted text for accuracy and brevity. Submit only content you have personally reviewed, and take responsibility for responding to review feedback. +- **Label generated media.** Model-generated images, video, and audio are valid examples or test artifacts for this multimodal project. Identify them as generated output and provide the model, inputs, and reproduction settings. Do not present them as ground-truth references or independent evidence of correctness. + +For example, a PR description could include: + +```text +AI assistance: Used Codex to draft the regression test and PR description. +I reviewed the changes and verified the test reproduces the reported failure. +Validation: . +``` + +Maintainer review time is limited. Contributions that are unverified, misleading, or beyond the author's understanding may be closed. Repeated violations may lead to contribution restrictions. New contributors are welcome to ask for help and submit small, well-understood changes. + +This policy is adapted from [Ghostty's AI Usage Policy](https://github.com/ghostty-org/ghostty/blob/main/AI_POLICY.md) for vLLM-Omni's contribution and multimodal validation workflows. The [coding-agent workflow](#using-repository-skills-with-coding-agents) provides practical guidance for applying it. + ## Issues If you encounter a bug or have a feature request, please search existing issues first to see if it has already been reported. If not, please file a new issue, providing as much relevant information as possible. From 5078c7ad66c393f16bf9092c2a66d1c179354008 Mon Sep 17 00:00:00 2001 From: Zhichao Zhang <60429419+smartDream-chao@users.noreply.github.com> Date: Wed, 9 Sep 2026 09:38:11 +0800 Subject: [PATCH 10/93] [Bugfix][MiMo-Audio] Align code2wav decode with tokenizer device (#6539) Signed-off-by: chaosansui Signed-off-by: Zhichao Zhang <60429419+smartDream-chao@users.noreply.github.com> --- .../test_mimo_audio_code2wav_batch_decode.py | 6 +- .../test_mimo_audio_tokenizer_worker_cache.py | 115 ++++++++++++++++++ .../models/mimo_audio/mimo_audio_code2wav.py | 61 ++++++---- 3 files changed, 160 insertions(+), 22 deletions(-) create mode 100644 tests/model_executor/models/mimo_audio/test_mimo_audio_tokenizer_worker_cache.py diff --git a/tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py b/tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py index 194a024dbc4..c20ee7f4896 100644 --- a/tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py +++ b/tests/model_executor/models/mimo_audio/test_mimo_audio_code2wav_batch_decode.py @@ -91,7 +91,11 @@ def _minimal_model(mocker: MockerFixture): decoder=decoder, config=SimpleNamespace(avg_pooler=2, stride_size=2, hop_length=240), ) - model._tokenizer_service = SimpleNamespace(audio_tokenizer=audio_tok, cuda_graph_wrapper=None) + model._tokenizer_service = SimpleNamespace( + audio_tokenizer=audio_tok, + cuda_graph_wrapper=None, + device="cpu", + ) return model, audio_tok diff --git a/tests/model_executor/models/mimo_audio/test_mimo_audio_tokenizer_worker_cache.py b/tests/model_executor/models/mimo_audio/test_mimo_audio_tokenizer_worker_cache.py new file mode 100644 index 00000000000..06b16f85353 --- /dev/null +++ b/tests/model_executor/models/mimo_audio/test_mimo_audio_tokenizer_worker_cache.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for the MiMo-Audio tokenizer-worker cache key. + +The worker cache key used to keep only the device *type*, so an explicit +``MIMO_AUDIO_TOKENIZER_DEVICE=cuda:1`` silently loaded the tokenizer on the +process's current CUDA device and shared one cache entry with ``cuda:0``. +The index must survive both the cache key and the worker constructor. + +CPU-only and weight-free: the worker class is replaced by a stub so no +tokenizer weights are loaded. +""" + +from __future__ import annotations + +import functools + +import pytest + +torch = pytest.importorskip("torch") + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_TOK_PATH = "/models/mimo-audio-tokenizer" + + +@functools.lru_cache(maxsize=1) +def _code2wav_module(): + """Defer the import (pulls vLLM model_executor) until first use.""" + from vllm_omni.model_executor.models.mimo_audio import mimo_audio_code2wav + + return mimo_audio_code2wav + + +class _StubWorker: + def __init__(self, device_str: str, config_path: str | None, audio_tokenizer_path: str): + self.device = device_str + self.config_path = config_path + self.audio_tokenizer_path = audio_tokenizer_path + + +@pytest.fixture +def isolated_worker_cache(monkeypatch): + """Empty worker cache with the real worker replaced by a stub.""" + mod = _code2wav_module() + monkeypatch.setattr(mod, "MiMoAudioTokenizerWorker", _StubWorker) + monkeypatch.setattr(mod, "_TOKENIZER_WORKER_CACHE", {}) + return mod + + +@pytest.mark.parametrize("device", [torch.device("cuda:1"), "cuda:1"]) +def test_cache_key_keeps_cuda_index(device): + """Both torch.device and string spellings retain the index.""" + mod = _code2wav_module() + key = mod._normalize_tokenizer_worker_cache_key(device, None, _TOK_PATH) + assert key[0] == "cuda:1" + + +def test_cache_key_cpu_unchanged(): + mod = _code2wav_module() + key = mod._normalize_tokenizer_worker_cache_key(torch.device("cpu"), None, _TOK_PATH) + assert key[0] == "cpu" + + +def test_cache_key_bare_cuda_resolves_to_current_device(monkeypatch): + """A bare "cuda" shares the entry with its equivalent indexed spelling.""" + mod = _code2wav_module() + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 3) + + bare = mod._normalize_tokenizer_worker_cache_key(torch.device("cuda"), None, _TOK_PATH) + indexed = mod._normalize_tokenizer_worker_cache_key(torch.device("cuda:3"), None, _TOK_PATH) + assert bare[0] == "cuda:3" + assert bare == indexed + + +def test_worker_constructed_on_requested_index(isolated_worker_cache): + mod = isolated_worker_cache + worker = mod.get_tokenizer_worker( + device=torch.device("cuda:1"), + config_path=None, + audio_tokenizer_path=_TOK_PATH, + ) + assert worker.device == "cuda:1" + + +def test_distinct_cuda_indices_get_distinct_workers(isolated_worker_cache): + mod = isolated_worker_cache + first = mod.get_tokenizer_worker( + device=torch.device("cuda:0"), + config_path=None, + audio_tokenizer_path=_TOK_PATH, + ) + second = mod.get_tokenizer_worker( + device=torch.device("cuda:1"), + config_path=None, + audio_tokenizer_path=_TOK_PATH, + ) + assert first is not second + assert (first.device, second.device) == ("cuda:0", "cuda:1") + + +def test_same_device_reuses_cached_worker(isolated_worker_cache): + mod = isolated_worker_cache + first = mod.get_tokenizer_worker( + device=torch.device("cuda:1"), + config_path=None, + audio_tokenizer_path=_TOK_PATH, + ) + second = mod.get_tokenizer_worker( + device="cuda:1", + config_path=None, + audio_tokenizer_path=_TOK_PATH, + ) + assert first is second diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py index 019670cf9d9..5423788a0af 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py @@ -99,8 +99,10 @@ def __init__( self.audio_tokenizer.config.nfft, ) mel_start = time.monotonic() - self.mel_transform = ( - MelSpectrogram( + # Build on CPU first: torchaudio MelSpectrogram init can mix CPU/CUDA + # tensors on newer PyTorch when default device is CUDA. + with torch.device("cpu"): + self.mel_transform = MelSpectrogram( sample_rate=self.audio_tokenizer.config.sampling_rate, n_fft=self.audio_tokenizer.config.nfft, hop_length=self.audio_tokenizer.config.hop_length, @@ -110,10 +112,9 @@ def __init__( n_mels=self.audio_tokenizer.config.n_mels, power=1.0, center=True, - ) - .to(self.device) - .to(torch.float32) - ) + ).to(torch.float32) + if self.device != "cpu": + self.mel_transform = self.mel_transform.to(self.device) logger.info( "[tokenizer worker] MelSpectrogram ready in %.2fs", time.monotonic() - mel_start, @@ -382,12 +383,19 @@ def extract_audio_code_tensor( def _normalize_tokenizer_worker_cache_key( - device: torch.device, + device: torch.device | str, config_path: str | None, audio_tokenizer_path: str, ) -> tuple[str, str, str]: """Normalize cache key so that same tokenizer always hits the same cache entry.""" - device_type = device.type if isinstance(device, torch.device) else str(device).split(":")[0] + dev = device if isinstance(device, torch.device) else torch.device(str(device)) + # Keep the index: an explicit "cuda:1" must not collapse onto the process's + # current CUDA device. Resolve a bare "cuda" so it shares the cache entry + # with the equivalent indexed spelling. + if dev.type == "cuda" and dev.index is None and torch.cuda.is_available(): + # Repo bans torch.cuda.current_device; match mimo_audio.py. + dev = torch.device("cuda", torch.accelerator.current_device_index()) + device_key = str(dev) # Use realpath so symlinks / trailing slash don't create duplicate entries ap = audio_tokenizer_path or "" if ap and os.path.exists(ap): @@ -398,22 +406,21 @@ def _normalize_tokenizer_worker_cache_key( if not cp and ap: cp = os.path.dirname(ap) - return (device_type, cp, ap) + return (device_key, cp, ap) _TOKENIZER_WORKER_CACHE: dict[tuple[str, str, str], MiMoAudioTokenizerWorker] = {} def get_tokenizer_worker( - device: torch.device, - config_path: str, + device: torch.device | str, + config_path: str | None, audio_tokenizer_path: str, ) -> MiMoAudioTokenizerWorker: key = _normalize_tokenizer_worker_cache_key(device, config_path, audio_tokenizer_path) if key not in _TOKENIZER_WORKER_CACHE: - device_type = key[0] _TOKENIZER_WORKER_CACHE[key] = MiMoAudioTokenizerWorker( - device_str=device_type, + device_str=key[0], config_path=config_path, audio_tokenizer_path=audio_tokenizer_path, ) @@ -469,8 +476,14 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = ""): or self.config.name_or_path ) + _tok_dev = os.environ.get("MIMO_AUDIO_TOKENIZER_DEVICE") + if _tok_dev: + tokenizer_device = torch.device("cpu" if _tok_dev.lower() == "cpu" else _tok_dev) + else: + tokenizer_device = self.device + self._tokenizer_service: MiMoAudioTokenizerWorker | None = get_tokenizer_worker( - device=self.device, + device=tokenizer_device, config_path=self.tokenizer_config_path, audio_tokenizer_path=self.audio_tokenizer_path, ) @@ -723,7 +736,7 @@ def _batch_decode_waveforms( Instead of calling _decode_waveform_from_codes per request (which incurs 4 GPU↔CPU round-trips each), this method: 1. Extracts audio_codes for all valid requests on CPU (cheap). - 2. Runs quantizer.decode_vq (embedding lookup) for each on GPU. + 2. Runs quantizer.decode_vq (embedding lookup) on the tokenizer device. 3. Packs all hidden-states into one tensor and calls decoder(packed_hs, input_lengths) once. 4. Splits the output waveforms back to per-request tensors. @@ -734,6 +747,9 @@ def _batch_decode_waveforms( return [empty] tokenizer = self._tokenizer_service.audio_tokenizer + # Codes/weights must share the tokenizer device (may be CPU via + # MIMO_AUDIO_TOKENIZER_DEVICE), not the stage CUDA device. + tok_device = torch.device(self._tokenizer_service.device) group_size = self.streamer_config.group_size audio_channels = self.streamer_config.audio_channels @@ -768,7 +784,7 @@ def _batch_decode_waveforms( use_cuda_graph = cg_ready and len(extracted) == 1 if use_cuda_graph: - wav_out = cg_wrapper.decode(extracted[0][1].to(self.device)) + wav_out = cg_wrapper.decode(extracted[0][1].to(tok_device)) wav = wav_out.squeeze(0).squeeze(0) cfg = tokenizer.config frames_per_token = cfg.avg_pooler * cfg.stride_size * cfg.hop_length @@ -782,7 +798,7 @@ def _batch_decode_waveforms( hidden_list: list[torch.Tensor] = [] lengths: list[int] = [] for _, audio_codes in extracted: - hs = tokenizer.encoder.decode_vq(audio_codes.to(self.device)) + hs = tokenizer.encoder.decode_vq(audio_codes.to(tok_device)) hidden_list.append(hs) lengths.append(hs.size(0)) @@ -790,7 +806,7 @@ def _batch_decode_waveforms( packed_hs = hidden_list[0] else: packed_hs = torch.cat(hidden_list, dim=0) - input_lengths = torch.tensor(lengths, device=self.device) + input_lengths = torch.tensor(lengths, device=tok_device) recon_wav = tokenizer.decoder(packed_hs, input_lengths) @@ -831,6 +847,9 @@ def _batch_chunked_decode_streaming( return [empty] tokenizer = self._tokenizer_service.audio_tokenizer + # Codes/weights must share the tokenizer device (may be CPU via + # MIMO_AUDIO_TOKENIZER_DEVICE), not the stage CUDA device. + tok_device = torch.device(self._tokenizer_service.device) group_size = self.streamer_config.group_size audio_channels = self.streamer_config.audio_channels @@ -886,7 +905,7 @@ def _batch_chunked_decode_streaming( frames_per_token = cfg.avg_pooler * cfg.stride_size * cfg.hop_length if use_cuda_graph: - wav_out = cg_wrapper.decode(extracted[0][1].to(self.device)) + wav_out = cg_wrapper.decode(extracted[0][1].to(tok_device)) wav = wav_out.squeeze(0).squeeze(0) valid_len = extracted[0][1].shape[-1] * frames_per_token if wav.numel() > valid_len: @@ -903,7 +922,7 @@ def _batch_chunked_decode_streaming( hidden_list: list[torch.Tensor] = [] lengths: list[int] = [] for _, audio_codes in extracted: - hs = tokenizer.encoder.decode_vq(audio_codes.to(self.device)) + hs = tokenizer.encoder.decode_vq(audio_codes.to(tok_device)) hidden_list.append(hs) lengths.append(hs.size(0)) @@ -911,7 +930,7 @@ def _batch_chunked_decode_streaming( packed_hs = hidden_list[0] else: packed_hs = torch.cat(hidden_list, dim=0) - input_lengths = torch.tensor(lengths, device=self.device) + input_lengths = torch.tensor(lengths, device=tok_device) recon_wav = tokenizer.decoder(packed_hs, input_lengths) From d7678e2bb30df56aad929e290f194b930f2b73b1 Mon Sep 17 00:00:00 2001 From: eval <74645252+eval-dev@users.noreply.github.com> Date: Tue, 8 Sep 2026 18:41:42 -0700 Subject: [PATCH 11/93] [Bugfix][MiniCPM-o] Fix the audio_embeds input path (#5730) Signed-off-by: eval-dev <0xe5bca0@gmail.com> Signed-off-by: eval <74645252+eval-dev@users.noreply.github.com> --- .../minicpmo_4_5/test_audio_embeds_input.py | 138 ++++++++++++++++++ .../minicpmo_4_5/minicpmo_4_5_omni_llm.py | 5 +- 2 files changed, 141 insertions(+), 2 deletions(-) create mode 100644 tests/model_executor/models/minicpmo_4_5/test_audio_embeds_input.py diff --git a/tests/model_executor/models/minicpmo_4_5/test_audio_embeds_input.py b/tests/model_executor/models/minicpmo_4_5/test_audio_embeds_input.py new file mode 100644 index 00000000000..d25b58e4858 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_audio_embeds_input.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for the MiniCPM-o 4.5 ``audio_embeds`` input path. + +Two defects made the OpenAI ``audio_embeds`` content part unusable with this model: + +1. ``MiniCPMOAudioEmbeddingItems`` registered itself under ``modality="image"`` + (copy-pasted from ``MiniCPMVImageEmbeddingItems``), so every request failed with + ``Modality 'image' not found. Available modalities: {'audio'}``. +2. ``_get_prompt_updates`` sized the placeholder from + ``sum(map(len, single_audio_embeds))``. ``MiniCPMOAudioEmbeddingInputs`` declares + ``TensorShape("bn", "s", "h")``, so one item is ``(s, h)`` and that expression + iterates rows and sums the hidden size -- ``s * h`` instead of ``s``. A one-second + slice of shape ``(10, 4096)`` asked for 40960 placeholder tokens instead of 10, + pushing every real request past ``max_model_len``. +""" + +from __future__ import annotations + +import math +from types import SimpleNamespace + +import pytest +import torch +from vllm.multimodal.parse import MultiModalDataItems + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + MiniCPMO45OmniLLMMultiModalProcessor, + MiniCPMO45OmniLLMProcessingInfo, + MiniCPMOAudioEmbeddingItems, + _minicpmo_field_config, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +POOL_STEP = 5 +HIDDEN_SIZE = 4096 +#: Audio embedding vectors produced per second: 100 mel -> 50 cnn -> (50-5)//5+1. +EMBEDS_PER_SECOND = 10 + + +class _FakeProcessor: + """Minimal processor surface used by ``get_audio_placeholder``.""" + + def __init__(self) -> None: + self.pool_step = POOL_STEP + self.audio_processor = SimpleNamespace(hop_length=160) + self.image_processor = SimpleNamespace(mean=[0.0], std=[1.0]) + + def get_audio_placeholder( + self, + audio_lens: int, + chunk_input: bool = True, + chunk_length: int = 1, + ) -> str: + del chunk_input, chunk_length + feature_lens = math.ceil(audio_lens / self.audio_processor.hop_length) + cnn_feature_lens = (feature_lens - 1) // 2 + 1 + output_lens = (cnn_feature_lens - self.pool_step) // self.pool_step + 1 + return "" * output_lens + + +class _IdentityTokenizer: + """``_get_prompt_updates`` round-trips the image/video patterns through the + tokenizer to detect encode/decode drift; an identity pair keeps that a no-op.""" + + def encode(self, text: str, add_special_tokens: bool = True) -> str: + del add_special_tokens + return text + + def decode(self, tokens: str) -> str: + return tokens + + +def _processing_info() -> MiniCPMO45OmniLLMProcessingInfo: + processor = _FakeProcessor() + tokenizer = _IdentityTokenizer() + info = object.__new__(MiniCPMO45OmniLLMProcessingInfo) + info.ctx = SimpleNamespace( + tokenizer=tokenizer, + get_tokenizer=lambda: tokenizer, + get_hf_config=lambda: SimpleNamespace(audio_pool_step=POOL_STEP), + get_hf_processor=lambda **kwargs: processor, + ) + return info + + +def _audio_placeholder_for(single_audio_embeds: torch.Tensor) -> str: + """Run the real prompt-update path and return the audio placeholder text. + + Goes through ``_get_prompt_updates`` rather than recomputing the length, so the + ``sum(map(len, ...))`` regression is actually exercised. + """ + info = _processing_info() + proc = object.__new__(MiniCPMO45OmniLLMMultiModalProcessor) + proc.info = info + + items = MiniCPMOAudioEmbeddingItems( + {"audio_embeds": single_audio_embeds.unsqueeze(0)}, + fields_factory=_minicpmo_field_config, + ) + mm_items = MultiModalDataItems({"audio": items}) + + updates = proc._get_prompt_updates(mm_items, {}, {}) + audio_update = next(u for u in updates if u.modality == "audio") + return audio_update.content(0).full + + +def test_audio_embedding_items_use_the_audio_modality() -> None: + """Registering under ``image`` makes every ``audio_embeds`` request fail.""" + items = MiniCPMOAudioEmbeddingItems( + {"audio_embeds": torch.zeros(1, EMBEDS_PER_SECOND, HIDDEN_SIZE)}, + fields_factory=_minicpmo_field_config, + ) + + assert items.modality == "audio" + + +@pytest.mark.parametrize("num_seconds", [1, 3, 30]) +def test_placeholder_length_matches_embedding_count(num_seconds: int) -> None: + """An ``(s, h)`` item must yield exactly ``s`` placeholder tokens.""" + num_embeds = num_seconds * EMBEDS_PER_SECOND + + placeholder = _audio_placeholder_for(torch.zeros(num_embeds, HIDDEN_SIZE)) + + assert placeholder.count("") == num_embeds + + +def test_placeholder_does_not_scale_with_hidden_size() -> None: + """The count must depend on ``s`` alone, never on ``h``. + + ``sum(map(len, ...))`` returned ``s * h``, so the placeholder grew with the model's + hidden size. Two items with the same ``s`` but different ``h`` must agree. + """ + narrow = _audio_placeholder_for(torch.zeros(EMBEDS_PER_SECOND, 8)) + wide = _audio_placeholder_for(torch.zeros(EMBEDS_PER_SECOND, HIDDEN_SIZE)) + + assert narrow.count("") == wide.count("") == EMBEDS_PER_SECOND diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index f6c9b1fa948..9dfff5589c7 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -3310,7 +3310,7 @@ def __init__( ) -> None: super().__init__( data, - modality="image", + modality="audio", required_fields={"audio_embeds"}, fields_factory=fields_factory, ) @@ -3733,7 +3733,8 @@ def get_audio_replacement(item_idx: int): if isinstance(audios, MiniCPMOAudioEmbeddingItems): single_audio_embeds = audios.get(item_idx)["audio_embeds"] - audio_len = self.info.get_audio_len_by_num_chunks(sum(map(len, single_audio_embeds))) + # One item is ``(s, h)``, so its leading dim is the audio embedding count. + audio_len = self.info.get_audio_len_by_num_chunks(len(single_audio_embeds)) else: audio_len = audios.get_audio_length(item_idx) From cdf027874c9fe070888b33127987c9a19e8c888c Mon Sep 17 00:00:00 2001 From: boatman <1930807094@qq.com> Date: Wed, 9 Sep 2026 09:58:38 +0800 Subject: [PATCH 12/93] [Feat][OmniVoice]Support Varlen Attn, Request-Batch and Step-Execution (#6408) Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com> --- docs/serving/speech_api.md | 4 + .../test_omnivoice_expansion.py | 96 +++- .../online_serving/test_omnivoice_parity.py | 96 ++++ .../openai_api/test_serving_speech.py | 112 +++- tests/helpers/client.py | 1 + .../omnivoice/test_cuda_graph_generator.py | 291 +++++----- .../omnivoice/test_fused_projection_load.py | 27 +- .../models/omnivoice/test_mask_dtype.py | 196 ++----- .../omnivoice/test_pipeline_batching.py | 228 ++++++++ vllm_omni/deploy/omnivoice.yaml | 2 + .../models/omnivoice/pipeline_omnivoice.py | 384 ++++++++++--- vllm_omni/diffusion/worker/utils.py | 2 +- .../entrypoints/openai/serving_speech.py | 19 + .../models/omnivoice/omnivoice_generator.py | 542 +++++++++++------- 14 files changed, 1393 insertions(+), 607 deletions(-) create mode 100644 tests/e2e/online_serving/test_omnivoice_parity.py create mode 100644 tests/model_executor/models/omnivoice/test_pipeline_batching.py diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index d4871ce15fa..a515e655768 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -752,6 +752,10 @@ Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` ne | ------- | ------------- | | `k2-fsa/OmniVoice` | Pure-diffusion TTS. Supports voice cloning via `ref_audio` (with optional `ref_text`); no built-in voice presets. | +OmniVoice uses packed variable-length attention for batched generator execution. The attention operator accepts FP16 and BF16 inputs, so its query, +key, and value tensors are evaluated in BF16 even when the stage is configured with `dtype: float32`; the attention output is converted back to the model's +hidden-state dtype before the output projection. Consequently, a float32 stage configuration does not imply FP32 attention arithmetic. + ### VoxCPM2 | Model | Description | diff --git a/tests/e2e/online_serving/test_omnivoice_expansion.py b/tests/e2e/online_serving/test_omnivoice_expansion.py index 3a042082e63..d77405a735a 100644 --- a/tests/e2e/online_serving/test_omnivoice_expansion.py +++ b/tests/e2e/online_serving/test_omnivoice_expansion.py @@ -7,7 +7,10 @@ accessed through the standard OpenAI-compatible speech API. """ +import io import os +import wave +from collections.abc import Sequence os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" @@ -15,7 +18,7 @@ from tests.helpers.mark import hardware_test from tests.helpers.media import get_asset_path -from tests.helpers.runtime import OmniServerParams +from tests.helpers.runtime import OmniResponse, OmniServerParams from tests.helpers.stage_config import get_deploy_config_path from vllm_omni.entrypoints.openai.serving_speech import _DEFAULT_VOICE_NAME @@ -34,6 +37,10 @@ EXTRA_ARGS = [ "--trust-remote-code", "--disable-log-stats", + "--max-num-seqs", + "8", + "--request-batch-max-wait-ms", + "50", ] TEST_PARAMS = [ OmniServerParams( @@ -42,7 +49,14 @@ server_args=EXTRA_ARGS, ) ] - +STEP_EXECUTION_ARGS = EXTRA_ARGS + ["--step-execution"] +STEP_EXECUTION_PARAMS = [ + OmniServerParams( + model=MODEL, + stage_config_path=STAGE_CONFIG, + server_args=STEP_EXECUTION_ARGS, + ) +] # Lower this in ``request_config`` via ``min_audio_bytes`` if a run produces legitimately short WAVs. _DEFAULT_MIN_AUDIO_BYTES = 5000 @@ -58,6 +72,33 @@ def get_prompt(prompt_type="text"): return prompts.get(prompt_type, prompts["text"]) +def _assert_valid_omnivoice_wav_responses(responses: Sequence[OmniResponse]) -> None: + """Validate that every concurrent response contains plausible OmniVoice audio.""" + audio_shapes: list[tuple[int, int, int, int]] = [] + for response in responses: + audio_bytes = response.audio_bytes + assert audio_bytes is not None + assert audio_bytes.startswith(b"RIFF") + with wave.open(io.BytesIO(audio_bytes), "rb") as wav_file: + num_channels = wav_file.getnchannels() + sample_width = wav_file.getsampwidth() + sample_rate = wav_file.getframerate() + num_frames = wav_file.getnframes() + assert num_channels == 1 + assert sample_width == 2 + assert sample_rate == 24000 + assert num_frames > 0 + audio_shapes.append((num_channels, sample_width, sample_rate, num_frames)) + + duration_s = num_frames / sample_rate + assert 0.1 <= duration_s <= 30.0 + + # Identical prompts use the same estimated target length. Their samples may + # differ numerically across batch layouts, but the decoded audio shape must + # remain consistent across concurrent requests. + assert len(set(audio_shapes)) == 1 + + @pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) class TestOmniVoiceTTS: """E2E tests for OmniVoice TTS model.""" @@ -74,6 +115,57 @@ def test_speech_auto_voice(self, omni_server, online_client) -> None: } online_client.send_audio_speech_request(request_config) + @hardware_test(res={"cuda": "L4"}, num_cards=1) + def test_speech_auto_voice_batch(self, omni_server, openai_client) -> None: + """Test concurrent request-batch TTS generation.""" + batch_request_config = { + "model": omni_server.model, + "input": get_prompt("text"), + "response_format": "wav", + "seed": 42, + "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES, + } + batch_r = openai_client.send_audio_speech_request(batch_request_config, request_num=2) + assert len(batch_r) == 2 + _assert_valid_omnivoice_wav_responses(batch_r) + + +@pytest.mark.parametrize("omni_server", STEP_EXECUTION_PARAMS, indirect=True) +class TestOmniVoiceStepExecution: + """E2E tests for OmniVoice TTS model.""" + + @hardware_test(res={"cuda": "L4"}, num_cards=1) + def test_speech_auto_voice_step_execution(self, omni_server, openai_client) -> None: + """Test auto voice TTS generation (text only, no reference audio).""" + request_config = { + "model": omni_server.model, + "input": get_prompt("text"), + "response_format": "wav", + "timeout": 180.0, + "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES, + "extra_params": { + "num_inference_steps": 32, + }, + } + openai_client.send_audio_speech_request(request_config) + + @hardware_test(res={"cuda": "L4"}, num_cards=1) + def test_speech_auto_voice_batch_step_execution(self, omni_server, openai_client) -> None: + """Test concurrent step-execution TTS generation.""" + batch_request_config = { + "model": omni_server.model, + "input": get_prompt("text"), + "response_format": "wav", + "seed": 42, + "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES, + "extra_params": { + "num_inference_steps": 32, + }, + } + batch_r = openai_client.send_audio_speech_request(batch_request_config, request_num=2) + assert len(batch_r) == 2 + _assert_valid_omnivoice_wav_responses(batch_r) + @pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) class TestOmniVoiceSeed: diff --git a/tests/e2e/online_serving/test_omnivoice_parity.py b/tests/e2e/online_serving/test_omnivoice_parity.py new file mode 100644 index 00000000000..e9b267f749d --- /dev/null +++ b/tests/e2e/online_serving/test_omnivoice_parity.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cross-mode full-model parity tests for OmniVoice.""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import pytest +import requests + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServer +from tests.helpers.stage_config import get_deploy_config_path + +pytestmark = [pytest.mark.slow, pytest.mark.tts] + +MODEL = "k2-fsa/OmniVoice" +STAGE_CONFIG = get_deploy_config_path("omnivoice.yaml") +PROMPT = "The weather is nice today, perfect for a walk in the park." + +payload = { + "model": MODEL, + "input": PROMPT, + "language": "English", + "seed": 42, + "response_format": "wav", + "extra_params": {"num_inference_steps": 32}, +} + + +def _generate_without_graph(server_args: list[str]) -> bytes: + with OmniServer( + MODEL, + server_args, + use_omni=True, + env_dict={"OMNIVOICE_CUDA_GRAPH": "0"}, + ) as server: + response = requests.post( + f"http://{server.host}:{server.port}/v1/audio/speech", + json=payload, + timeout=600, + ) + response.raise_for_status() + assert response.content.startswith(b"RIFF") + return response.content + + +def _generate_with_graph(server_args: list[str]) -> bytes: + with OmniServer( + MODEL, + server_args, + use_omni=True, + ) as server: + response = requests.post( + f"http://{server.host}:{server.port}/v1/audio/speech", + json=payload, + timeout=600, + ) + response.raise_for_status() + assert response.content.startswith(b"RIFF") + return response.content + + +def _common_args() -> list[str]: + """Return the shared B=1 server arguments for parity tests.""" + return [ + "--trust-remote-code", + "--disable-log-stats", + "--deploy-config", + STAGE_CONFIG, + "--max-num-seqs", + "1", + ] + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_request_mode_and_step_execution_b1_parity_without_graph() -> None: + """B=1 eager request and step modes must produce identical seeded WAV bytes.""" + common_args = _common_args() + + request_audio = _generate_without_graph(common_args) + step_audio = _generate_without_graph([*common_args, "--step-execution", "--enforce-eager"]) + + assert request_audio == step_audio + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_request_mode_and_step_execution_b1_parity_with_graph() -> None: + """B=1 Graph request and step modes must produce identical seeded WAV bytes.""" + common_args = _common_args() + + request_audio = _generate_with_graph(common_args) + step_audio = _generate_with_graph([*common_args, "--step-execution", "--enforce-eager"]) + + assert request_audio == step_audio diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 7e1a9ccec42..6a5ebcf7d3b 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -27,6 +27,9 @@ from pytest_mock import MockerFixture from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler +from vllm_omni.diffusion.sched.step_scheduler import StepScheduler from vllm_omni.entrypoints.omni_base import OmniEngineDeadError from vllm_omni.entrypoints.openai import api_server as api_server_module from vllm_omni.entrypoints.openai import serving_speech as serving_speech_module @@ -55,6 +58,7 @@ from vllm_omni.entrypoints.openai.tts_adapters.ming_tts import MingTTSAdapter from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSAdapter, Qwen3TTSCodecLimitError from vllm_omni.entrypoints.openai.tts_adapters.voxtral import VoxtralTTSAdapter +from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.model_executor.models.fish_speech.prompt_utils import ( FISH_TEXT_ONLY_SYSTEM_PROMPT, build_fish_voice_clone_prompt_ids, @@ -753,13 +757,12 @@ async def test_diffusion_create_speech_with_unknown_voice(self, mocker: MockerFi @pytest.mark.asyncio async def test_create_diffusion_speech_extra_params(self, mocker: MockerFixture): - """Test public diffusion speech success and extra_params propagation.""" + """Test diffusion parameters reach StepScheduler as standard fields.""" # Mock the engine client mock_engine = mocker.MagicMock() # Mock default sampling params - mock_sampling_param = mocker.MagicMock() - mock_sampling_param.extra_args = {"existing_arg": "value"} + mock_sampling_param = OmniDiffusionSamplingParams(extra_args={"existing_arg": "value"}) mock_engine.default_sampling_params_list = [mock_sampling_param] # Mock generate to yield a valid OmniRequestOutput @@ -775,7 +778,15 @@ async def mock_generate(*args, **kwargs): server, "create_audio", return_value=mocker.MagicMock(audio_data=b"dummy", media_type="audio/wav") ) - req = OpenAICreateSpeechRequest(input="Hello", extra_params={"new_arg": 123, "existing_arg": "new_value"}) + req = OpenAICreateSpeechRequest( + input="Hello", + extra_params={ + "new_arg": 123, + "existing_arg": "new_value", + "num_inference_steps": 12, + "guidance_scale": 7.0, + }, + ) response = await server.create_speech(req) @@ -792,7 +803,98 @@ async def mock_generate(*args, **kwargs): # Verify it was deepcopied and updated assert passed_params is not mock_engine.default_sampling_params_list - assert passed_params[0].extra_args == {"existing_arg": "new_value", "new_arg": 123} + assert passed_params[0].extra_args == { + "existing_arg": "new_value", + "new_arg": 123, + "num_inference_steps": 12, + "guidance_scale": 7.0, + } + assert passed_params[0].num_inference_steps == 12 + assert passed_params[0].guidance_scale == 7.0 + + # Regression: StepScheduler.add_request() used to receive + # num_inference_steps=None and fail while converting it to int. + scheduler = StepScheduler() + scheduler.add_request( + OmniDiffusionRequest( + prompt="Hello", + sampling_params=passed_params[0], + request_id="speech-test", + ) + ) + assert scheduler._request_progress["speech-test"].total_steps == 12 + + @pytest.mark.asyncio + async def test_diffusion_speech_guidance_promotion_controls_request_batch_admission( + self, + mocker: MockerFixture, + ) -> None: + """Different request guidance values must not enter one request batch.""" + mock_engine = mocker.MagicMock() + mock_engine.default_sampling_params_list = [OmniDiffusionSamplingParams(num_inference_steps=12)] + passed_sampling_params = [] + + async def mock_generate(*args, **kwargs): + passed_sampling_params.append(kwargs["sampling_params_list"][0]) + yield create_mock_audio_output_for_test() + + mock_engine.generate = mocker.MagicMock(side_effect=mock_generate) + server = OmniOpenAIServingSpeech.for_diffusion(diffusion_engine=mock_engine, model_name="test-model") + mocker.patch.object( + server, + "create_audio", + return_value=mocker.MagicMock(audio_data=b"dummy", media_type="audio/wav"), + ) + + for guidance_scale in (2.0, 7.0): + response = await server.create_speech( + OpenAICreateSpeechRequest( + input="Hello", + extra_params={"guidance_scale": guidance_scale}, + ) + ) + assert response.status_code == 200 + + scheduler = RequestScheduler() + scheduler.initialize(SimpleNamespace(max_num_seqs=2)) + for index, sampling_params in enumerate(passed_sampling_params): + scheduler.add_request( + OmniDiffusionRequest( + prompt="Hello", + sampling_params=sampling_params, + request_id=f"speech-{index}", + ) + ) + + first = scheduler.schedule() + + assert [request.request_id for request in first.scheduled_new_reqs] == ["speech-0"] + assert first.num_running_reqs == 1 + assert first.num_waiting_reqs == 1 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("extra_params", "expected_message"), + [ + ({"num_inference_steps": "invalid"}, "num_inference_steps must be an integer"), + ({"guidance_scale": "invalid"}, "guidance_scale must be a number"), + ], + ) + async def test_create_diffusion_speech_rejects_invalid_scheduler_params( + self, + mocker: MockerFixture, + extra_params: dict[str, str], + expected_message: str, + ) -> None: + mock_engine = mocker.MagicMock() + mock_engine.default_sampling_params_list = [OmniDiffusionSamplingParams()] + server = OmniOpenAIServingSpeech.for_diffusion(diffusion_engine=mock_engine, model_name="test-model") + + response = await server.create_speech(OpenAICreateSpeechRequest(input="Hello", extra_params=extra_params)) + + assert response.status_code == 400 + assert expected_message in response.body.decode() + mock_engine.generate.assert_not_called() class TestTTSMethods: diff --git a/tests/helpers/client.py b/tests/helpers/client.py index 5339f7195b9..6a6a9b7ed31 100644 --- a/tests/helpers/client.py +++ b/tests/helpers/client.py @@ -1468,6 +1468,7 @@ def send_audio_speech_request(self, request_config: dict[str, Any], request_num: "instructions", "speed", "sample_rate", + "extra_params", "stream_format", "x_vector_only_mode", ): diff --git a/tests/model_executor/models/omnivoice/test_cuda_graph_generator.py b/tests/model_executor/models/omnivoice/test_cuda_graph_generator.py index b55f8279c44..ef05bf7383c 100644 --- a/tests/model_executor/models/omnivoice/test_cuda_graph_generator.py +++ b/tests/model_executor/models/omnivoice/test_cuda_graph_generator.py @@ -6,19 +6,27 @@ Verifies that _OmniVoiceCUDAGraphForward produces results equivalent to eager mode across three scenarios: - Exact-size inputs (no padding) → bit-identical - - Padded inputs (padded to nearest bucket) → correct slicing, exact match - for position-independent models - - Oversized inputs (fallback to eager) → bit-identical + - Padded inputs (padded to nearest bucket) → correct slicing and exact match + - Oversized inputs (128-aligned lazy capture) → bit-identical -Uses a lightweight synthetic generator to keep warmup fast and avoid -loading actual model weights. +Uses a small randomly initialized OmniVoiceGenerator so every test exercises +the production embedding, transformer, varlen-attention, and logits paths +without loading checkpoint weights. """ from __future__ import annotations +from types import SimpleNamespace + import pytest import torch -import torch.nn as nn +from vllm.utils.math_utils import round_up + +from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import ( + OmniVoiceGenerator, + _OmniVoiceCUDAGraphForward, +) +from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig pytestmark = [ pytest.mark.core_model, @@ -29,72 +37,45 @@ DEVICE = torch.device("cuda:0") NUM_CB = 8 VOCAB = 1025 -HIDDEN = 64 CAPTURE_SIZES = [32, 64, 128] -HEAD_DIM = 64 # --------------------------------------------------------------------------- -# Synthetic generator: matches _OmniVoiceCUDAGraphForward's interface +# Helpers # --------------------------------------------------------------------------- -class _SyntheticGenerator: - """Position-independent Embedding+Linear replacing OmniVoiceGenerator's 28-layer transformer. - - Each token position is processed independently (no attention), so padded - runs produce bit-identical results to non-padded runs on the same positions. - This makes the padding/slicing logic easy to verify exactly. - """ - - class _Cfg: - num_audio_codebook = NUM_CB - - def __init__(self, device: torch.device): - self.config = self._Cfg() - self.text_embedding = nn.Embedding(1000, HIDDEN).to(device).eval() - self._linear = nn.Linear(HIDDEN, NUM_CB * VOCAB, bias=False).to(device).eval() - self._rope_table: torch.Tensor | None = None - - @property - def model_dtype(self) -> torch.dtype: - """Part of the interface _OmniVoiceCUDAGraphForward expects of a generator.""" - return self.text_embedding.weight.dtype +def _eager(gen: OmniVoiceGenerator, ids: torch.Tensor, mask: torch.Tensor, cu_seqs) -> torch.Tensor: + """Run the production varlen step outside CUDA Graph.""" + seq_len = ids.shape[0] + rope_table = gen._rope_table_for(seq_len, ids.device, gen.model_dtype) + return gen._step_forward(ids, mask, cu_seqs, rope_table) - def _rope_table_for(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: - """Part of the interface _OmniVoiceCUDAGraphForward expects of a generator.""" - return torch.zeros(seq_len, HEAD_DIM, device=device, dtype=dtype) - - def _step_forward( - self, - input_ids: torch.Tensor, - audio_mask: torch.Tensor, - attention_mask: torch.Tensor | None, - rope_table: torch.Tensor, - ) -> torch.Tensor: - two_b, _, S = input_ids.shape - x = self.text_embedding(input_ids[:, 0, :].clamp(0, 999)) # [two_b, S, H] - logits = self._linear(x) # [two_b, S, 8*1025] - return logits.view(two_b, S, NUM_CB, VOCAB).permute(0, 2, 1, 3) # [two_b, 8, S, 1025] - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - -def _eager(gen: _SyntheticGenerator, ids: torch.Tensor, mask: torch.Tensor, attn: torch.Tensor) -> torch.Tensor: - """Run _step_forward with a zero RoPE table (synthetic gen ignores it).""" - two_b, _, S = ids.shape - return gen._step_forward(ids, mask, attn, gen._rope_table_for(S, ids.device, torch.float32)) +def _bucketed_eager(gen, wrapper, ids, mask, cu_seqs, batch_size) -> torch.Tensor: + """Run eager with the exact padded shape and metadata used by Graph.""" + seq_len = ids.shape[0] + bucket = wrapper._find_bucket(batch_size, seq_len) + if bucket is None: + bucket = round_up(seq_len, wrapper._LAZY_CAPTURE_ALIGNMENT) + padded_ids, padded_mask = wrapper._pad_inputs(ids, mask, bucket) + padded_cu_seqs = cu_seqs.clone() + padded_cu_seqs[-1] = bucket + return _eager(gen, padded_ids, padded_mask, padded_cu_seqs)[:, :seq_len, :] def _make_inputs(seq_len: int, device: torch.device = DEVICE): - two_b = 2 - ids = torch.randint(0, 100, (two_b, NUM_CB, seq_len), dtype=torch.long, device=device) - mask = torch.ones(two_b, seq_len, dtype=torch.bool, device=device) - attn = torch.ones(two_b, 1, seq_len, seq_len, dtype=torch.bool, device=device) - return ids, mask, attn + ids = torch.randint(0, 100, (seq_len, NUM_CB), dtype=torch.long, device=device) + mask = torch.ones(seq_len, dtype=torch.bool, device=device) + cond_len = (seq_len + 1) // 2 + uncond_len = seq_len - cond_len + + cu_seqs = torch.tensor( + [0, cond_len, cond_len + uncond_len, cond_len + uncond_len], + dtype=torch.int32, + device=device, + ) + return ids, mask, cu_seqs # --------------------------------------------------------------------------- @@ -105,15 +86,26 @@ def _make_inputs(seq_len: int, device: torch.device = DEVICE): @pytest.fixture(scope="module") def gen(): torch.manual_seed(42) - return _SyntheticGenerator(DEVICE) + config = OmniVoiceConfig( + audio_vocab_size=VOCAB, + num_audio_codebook=NUM_CB, + llm_config={ + "hidden_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "intermediate_size": 128, + "vocab_size": 128, + "max_position_embeddings": 4096, + "head_dim": 16, + }, + enable_cuda_graph=False, + ) + return OmniVoiceGenerator(config, SimpleNamespace(max_num_seqs=2)).to(DEVICE).eval() @pytest.fixture(scope="module") def wrapper(gen): - from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import ( - _OmniVoiceCUDAGraphForward, - ) - w = _OmniVoiceCUDAGraphForward(gen, capture_sizes=CAPTURE_SIZES) w.warmup(DEVICE) return w @@ -127,10 +119,10 @@ def wrapper(gen): @pytest.mark.parametrize("seq_len", CAPTURE_SIZES) def test_exact_size_bit_identical(gen, wrapper, seq_len): """When input exactly matches a captured bucket, output must be bit-identical to eager.""" - ids, mask, attn = _make_inputs(seq_len) + ids, mask, cu_seqs = _make_inputs(seq_len) with torch.no_grad(): - eager_out = _eager(gen, ids, mask, attn) - graph_out = wrapper(ids, mask, attn) + eager_out = _eager(gen, ids, mask, cu_seqs) + graph_out = wrapper(ids, mask, cu_seqs, 1) torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0) @@ -140,40 +132,36 @@ def test_exact_size_bit_identical(gen, wrapper, seq_len): @pytest.mark.parametrize("seq_len", [1, 15, 33, 60, 100]) -def test_padded_output_shape(gen, wrapper, seq_len): +def test_padded_output_shape(wrapper, seq_len): """Graph output must be sliced back to actual seq_len, not the bucket size.""" - ids, mask, attn = _make_inputs(seq_len) + ids, mask, cu_seqs = _make_inputs(seq_len) with torch.no_grad(): - graph_out = wrapper(ids, mask, attn) - assert graph_out.shape == (2, NUM_CB, seq_len, VOCAB) + graph_out = wrapper(ids, mask, cu_seqs, 1) + assert graph_out.shape == (NUM_CB, seq_len, VOCAB) @pytest.mark.parametrize("seq_len", [15, 33, 60, 100]) def test_padded_output_matches_eager(gen, wrapper, seq_len): - """Padded graph output must equal eager output at actual positions. - - The synthetic model has no attention across positions, so zero-padding - does not affect non-padded positions — exact match is expected. - """ - ids, mask, attn = _make_inputs(seq_len) + """Padded graph output must equal eager output at actual positions.""" + ids, mask, cu_seqs = _make_inputs(seq_len) with torch.no_grad(): - eager_out = _eager(gen, ids, mask, attn) - graph_out = wrapper(ids, mask, attn) + eager_out = _bucketed_eager(gen, wrapper, ids, mask, cu_seqs, 1) + graph_out = wrapper(ids, mask, cu_seqs, 1) torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0) # --------------------------------------------------------------------------- -# 3. Oversized inputs → fallback to eager (lazy capture), bit-identical +# 3. Oversized inputs → aligned lazy capture, bit-identical # --------------------------------------------------------------------------- @pytest.mark.parametrize("seq_len", [129, 200, 256]) -def test_fallback_eager_bit_identical(gen, wrapper, seq_len): - """Sequences exceeding the largest bucket fall back to lazy capture → bit-identical.""" - ids, mask, attn = _make_inputs(seq_len) +def test_aligned_lazy_capture_bit_identical(gen, wrapper, seq_len): + """Sequences beyond the static plan use aligned lazy capture and remain bit-identical.""" + ids, mask, cu_seqs = _make_inputs(seq_len) with torch.no_grad(): - eager_out = _eager(gen, ids, mask, attn) - graph_out = wrapper(ids, mask, attn) + eager_out = _bucketed_eager(gen, wrapper, ids, mask, cu_seqs, 1) + graph_out = wrapper(ids, mask, cu_seqs, 1) torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0) @@ -184,77 +172,96 @@ def test_fallback_eager_bit_identical(gen, wrapper, seq_len): def test_deterministic_across_calls(wrapper): """Same input must produce identical output on repeated CUDA graph replays.""" - ids, mask, attn = _make_inputs(32) + ids, mask, cu_seqs = _make_inputs(32) with torch.no_grad(): - out1 = wrapper(ids, mask, attn) - out2 = wrapper(ids, mask, attn) + out1 = wrapper(ids, mask, cu_seqs, 1).clone() + out2 = wrapper(ids, mask, cu_seqs, 1).clone() torch.testing.assert_close(out1, out2, atol=0, rtol=0) -# --------------------------------------------------------------------------- -# 5. _find_bucket logic (CPU, no CUDA graph) -# --------------------------------------------------------------------------- - - -def test_find_bucket_returns_nearest_bucket(): - """_find_bucket must return the smallest bucket >= seq_len, or None if all are smaller.""" - from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import ( - _OmniVoiceCUDAGraphForward, - ) +def test_replay_uses_updated_cu_seqs(gen, wrapper): + """The same graph key must honor new sequence boundaries on replay.""" + ids, mask, _ = _make_inputs(48) + cu_seqs_a = torch.tensor([0, 32, 48, 48], dtype=torch.int32, device=DEVICE) + cu_seqs_b = torch.tensor([0, 17, 48, 48], dtype=torch.int32, device=DEVICE) - w = _OmniVoiceCUDAGraphForward.__new__(_OmniVoiceCUDAGraphForward) - w._capture_sizes = [32, 64, 128] - w._graphs = {} - - assert w._find_bucket(1) == 32 - assert w._find_bucket(32) == 32 - assert w._find_bucket(33) == 64 - assert w._find_bucket(64) == 64 - assert w._find_bucket(128) == 128 - assert w._find_bucket(129) is None - assert w._find_bucket(1000) is None + with torch.no_grad(): + eager_a = _bucketed_eager(gen, wrapper, ids, mask, cu_seqs_a, 1) + graph_a = wrapper(ids, mask, cu_seqs_a, 1).clone() + eager_b = _bucketed_eager(gen, wrapper, ids, mask, cu_seqs_b, 1) + graph_b = wrapper(ids, mask, cu_seqs_b, 1).clone() + torch.testing.assert_close(graph_a, eager_a, atol=0, rtol=0) + torch.testing.assert_close(graph_b, eager_b, atol=0, rtol=0) + assert not torch.equal(graph_a, graph_b) -# --------------------------------------------------------------------------- -# 6. enable_cuda_graph=False produces same tokens as enable_cuda_graph=True -# --------------------------------------------------------------------------- +def test_batch_two_cu_seqs_matches_eager(gen, wrapper): + """B=2 uses four real sequence segments plus the fixed tail slot.""" + ids, mask, _ = _make_inputs(48) + cu_seqs = torch.tensor([0, 10, 18, 32, 48, 48], dtype=torch.int32, device=DEVICE) -def test_cuda_graph_disabled_matches_eager_generator(): - """OmniVoiceGenerator with enable_cuda_graph=False must produce the same - _step_forward output as one with enable_cuda_graph=True (after warmup). + assert cu_seqs.numel() == 2 * 2 + 2 + with torch.no_grad(): + eager_out = _bucketed_eager(gen, wrapper, ids, mask, cu_seqs, 2) + graph_out = wrapper(ids, mask, cu_seqs, 2).clone() - Uses a 2-layer config so warmup completes quickly without loading weights. - """ - from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import OmniVoiceGenerator - from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig + torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0) - cfg = OmniVoiceConfig() - cfg.llm_num_hidden_layers = 2 - cfg.num_hidden_layers = 2 - torch.manual_seed(0) - gen_eager = OmniVoiceGenerator(cfg).to(DEVICE).eval() - gen_eager.config.enable_cuda_graph = False - gen_eager._cuda_graph_fwd = None +def test_lazy_graph_cache_uses_lru_eviction(wrapper): + """A lazy-cache hit must protect that graph from the next eviction.""" + original_limit = wrapper._MAX_LAZY_GRAPHS + wrapper._lazy_graphs.clear() + wrapper._MAX_LAZY_GRAPHS = 2 + try: + # Static B=1 coverage ends at 128. These lengths round to three + # distinct 128-aligned lazy keys: 256, 384, and 512. + for seq_len in (129, 257): + ids, mask, cu_seqs = _make_inputs(seq_len) + with torch.no_grad(): + wrapper(ids, mask, cu_seqs, 1) + assert list(wrapper._lazy_graphs) == [(1, 256), (1, 384)] + + # Refresh key 256, making key 384 the least recently used entry. + ids, mask, cu_seqs = _make_inputs(129) + with torch.no_grad(): + wrapper(ids, mask, cu_seqs, 1) + assert list(wrapper._lazy_graphs) == [(1, 384), (1, 256)] + + # Inserting key 512 must evict key 384, not the recently hit key 256. + ids, mask, cu_seqs = _make_inputs(385) + with torch.no_grad(): + wrapper(ids, mask, cu_seqs, 1) + assert list(wrapper._lazy_graphs) == [(1, 256), (1, 512)] + finally: + wrapper._MAX_LAZY_GRAPHS = original_limit + wrapper._lazy_graphs.clear() - torch.manual_seed(0) - gen_graph = OmniVoiceGenerator(cfg).to(DEVICE).eval() - gen_graph.config.cuda_graph_capture_sizes = [64] - from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import _OmniVoiceCUDAGraphForward - gen_graph._cuda_graph_fwd = _OmniVoiceCUDAGraphForward(gen_graph, capture_sizes=[64]) - gen_graph._cuda_graph_fwd.warmup(DEVICE) +# --------------------------------------------------------------------------- +# 5. _find_bucket logic (CPU, no CUDA graph) +# --------------------------------------------------------------------------- - seq_len = 64 - ids = torch.zeros(2, cfg.num_audio_codebook, seq_len, dtype=torch.long, device=DEVICE) - mask = torch.ones(2, seq_len, dtype=torch.bool, device=DEVICE) - attn = torch.ones(2, 1, seq_len, seq_len, dtype=torch.bool, device=DEVICE) - rope_table = gen_eager._rope_table_for(seq_len, DEVICE, gen_eager.model_dtype) +def test_find_bucket_returns_nearest_bucket(): + """_find_bucket must return the smallest bucket >= seq_len, or None if all are smaller.""" + from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import ( + _OmniVoiceCUDAGraphForward, + ) - with torch.no_grad(): - eager_logits = gen_eager._step_forward(ids, mask, attn, rope_table) - graph_logits = gen_graph._cuda_graph_fwd(ids, mask, attn) + w = _OmniVoiceCUDAGraphForward.__new__(_OmniVoiceCUDAGraphForward) + w.capture_bucket_sizes_by_batch = { + 1: [32, 64, 128], + 2: [64, 128, 256], + } + w._graphs = {} - torch.testing.assert_close(graph_logits, eager_logits, atol=0, rtol=0) + assert w._find_bucket(1, 1) == 32 + assert w._find_bucket(1, 32) == 32 + assert w._find_bucket(1, 33) == 64 + assert w._find_bucket(1, 128) == 128 + assert w._find_bucket(1, 129) is None + assert w._find_bucket(2, 33) == 64 + assert w._find_bucket(2, 129) == 256 + assert w._find_bucket(3, 64) is None diff --git a/tests/model_executor/models/omnivoice/test_fused_projection_load.py b/tests/model_executor/models/omnivoice/test_fused_projection_load.py index fc5db01c7ab..c80f0e65bd2 100644 --- a/tests/model_executor/models/omnivoice/test_fused_projection_load.py +++ b/tests/model_executor/models/omnivoice/test_fused_projection_load.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 """The OmniVoice generator packs q/k/v and gate/up into fused projections. The HF checkpoint stores those five tensors separately, so packing them is a @@ -14,6 +14,8 @@ from __future__ import annotations +from types import SimpleNamespace + import pytest import torch @@ -49,6 +51,13 @@ def _config() -> OmniVoiceConfig: ) +def _generator() -> OmniVoiceGenerator: + return OmniVoiceGenerator( + _config(), + SimpleNamespace(max_num_seqs=1), + ) + + def _checkpoint_shards() -> dict[str, torch.Tensor]: """The five per-layer tensors an OmniVoice checkpoint actually stores.""" torch.manual_seed(0) @@ -64,7 +73,7 @@ def _checkpoint_shards() -> dict[str, torch.Tensor]: def test_every_fused_parameter_is_written() -> None: - generator = OmniVoiceGenerator(_config()) + generator = _generator() shards = _checkpoint_shards() loaded = generator._load_fused_projections(shards) @@ -77,7 +86,7 @@ def test_every_fused_parameter_is_written() -> None: def test_packed_weight_matches_the_shards_it_came_from() -> None: """Packing order is q,k,v and gate,up -- the split in forward() assumes it.""" - generator = OmniVoiceGenerator(_config()) + generator = _generator() shards = _checkpoint_shards() generator._load_fused_projections(shards) @@ -101,7 +110,7 @@ def test_packed_weight_matches_the_shards_it_came_from() -> None: def test_no_fused_parameter_keeps_its_random_init() -> None: """The corruption this guards against: a fused param never written at all.""" - generator = OmniVoiceGenerator(_config()) + generator = _generator() before = { name: param.detach().clone() for name, param in generator.named_parameters() @@ -121,7 +130,7 @@ def test_no_fused_parameter_keeps_its_random_init() -> None: ["self_attn.k_proj", "self_attn.v_proj", "self_attn.q_proj", "mlp.gate_proj", "mlp.up_proj"], ) def test_a_missing_shard_raises_instead_of_loading_partially(dropped: str) -> None: - generator = OmniVoiceGenerator(_config()) + generator = _generator() shards = _checkpoint_shards() del shards[f"llm.layers.0.{dropped}.weight"] @@ -130,7 +139,7 @@ def test_a_missing_shard_raises_instead_of_loading_partially(dropped: str) -> No def test_a_wrong_shaped_shard_raises() -> None: - generator = OmniVoiceGenerator(_config()) + generator = _generator() shards = _checkpoint_shards() shards["llm.layers.0.self_attn.q_proj.weight"] = torch.randn(NUM_HEADS * HEAD_DIM + 1, HIDDEN) @@ -140,7 +149,7 @@ def test_a_wrong_shaped_shard_raises() -> None: def test_a_checkpoint_missing_a_whole_layer_raises() -> None: """Half-loaded is the dangerous state: it neither errors nor works.""" - generator = OmniVoiceGenerator(_config()) + generator = _generator() shards = {k: v for k, v in _checkpoint_shards().items() if not k.startswith("llm.layers.1.")} with pytest.raises(ValueError, match="fused projections"): @@ -149,7 +158,7 @@ def test_a_checkpoint_missing_a_whole_layer_raises() -> None: def test_a_checkpoint_with_no_fused_shards_is_left_alone() -> None: """An unrelated state_dict must not trip the completeness check.""" - generator = OmniVoiceGenerator(_config()) + generator = _generator() assert generator._load_fused_projections({"audio_heads.weight": torch.randn(4, 4)}) == set() @@ -162,7 +171,7 @@ def test_load_weights_wires_the_packing_in(tmp_path, monkeypatch: pytest.MonkeyP """ import safetensors.torch - generator = OmniVoiceGenerator(_config()) + generator = _generator() shards = _checkpoint_shards() before = generator.layers[0].self_attn.qkv_proj.weight.detach().clone() diff --git a/tests/model_executor/models/omnivoice/test_mask_dtype.py b/tests/model_executor/models/omnivoice/test_mask_dtype.py index ff899e2a2f1..ff8dda49d7e 100644 --- a/tests/model_executor/models/omnivoice/test_mask_dtype.py +++ b/tests/model_executor/models/omnivoice/test_mask_dtype.py @@ -1,183 +1,79 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""OmniVoice attention masks must follow the model dtype, not a hardcoded float32. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""OmniVoice SDPA fallback masks must not depend on the model dtype. -SDPA requires the additive attention mask to carry the same dtype as the query. -The generator used to materialize that mask as float32 unconditionally, in three -places (the shared per-forward mask, the CUDA-graph capture buffer, and the -replay-time normalization), so every half-precision deployment died on the first -attention call with:: +The previous padded attention path used an additive floating-point mask. SDPA +requires such a mask to have the same dtype as the query, and constructing it +as float32 caused half-precision OmniVoice inference to fail on CUDA. - RuntimeError: invalid dtype for bias - should match query's dtype - -(the exact wording depends on which SDPA backend is selected for the shape; the -math backend words the same rejection as ``attn_mask.dtype``.) - -That left OmniVoice servable only in float32, even though the upstream k2-fsa -implementation runs it in float16. - -Note the split between the CPU and CUDA tests below: the CPU SDPA backend -silently accepts the mismatched mask, so only the CUDA tests can pin the actual -crash. The CPU tests cover the mask contract itself, which is what the fix -changes and what a future regression would break first. +The packed-varlen path no longer needs an additive mask. When SDPA fallback is +required, ``_attention_metadata_from_cu_seqs`` builds a boolean block mask. +Boolean SDPA masks have no dtype coupling with the query, so no mask coercion +is required for float16 or bfloat16 inference. """ from __future__ import annotations import pytest import torch +import torch.nn.functional as F from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import ( - OmniVoiceAttention, - OmniVoiceGenerator, - _additive_float_mask, + _attention_metadata_from_cu_seqs, ) -from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig HALF_DTYPES = [torch.float16, torch.bfloat16] -ALL_DTYPES = HALF_DTYPES + [torch.float32] cpu_test = pytest.mark.core_model cuda_test = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def _tiny_config() -> OmniVoiceConfig: - """A real OmniVoiceConfig sized so the module fits in a unit test.""" - return OmniVoiceConfig( - llm_config={ - "hidden_size": 32, - "num_hidden_layers": 2, - "num_attention_heads": 4, - "num_key_value_heads": 2, - "head_dim": 8, - "intermediate_size": 64, - "vocab_size": 64, - "max_position_embeddings": 128, - }, - enable_cuda_graph=False, - ) - - -def _inputs(dtype: torch.dtype, device: torch.device) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - hidden = torch.randn(2, 6, 32, device=device, dtype=dtype) - bool_mask = torch.ones(2, 1, 6, 6, dtype=torch.bool, device=device) - bool_mask[:, :, :, 4:] = False - rope_table = torch.zeros(6, 8, device=device, dtype=dtype) - rope_table[:, :4] = 1.0 # cos = 1, sin = 0: no rotation, so the test is about the mask - return hidden, bool_mask, rope_table - - -# -------------------------------------------------------------------------- -# The mask contract (CPU) -# -------------------------------------------------------------------------- - - -@cpu_test -@pytest.mark.cpu -@pytest.mark.parametrize("dtype", ALL_DTYPES) -def test_additive_float_mask_uses_requested_dtype(dtype: torch.dtype) -> None: - out = _additive_float_mask(torch.tensor([[True, False]]), dtype) - assert out.dtype == dtype - - -@cpu_test -@pytest.mark.cpu -def test_additive_float_mask_keeps_mask_semantics() -> None: - """Attend positions stay at 0.0 and masked ones at -inf, in half precision too.""" - out = _additive_float_mask(torch.tensor([[True, False]]), torch.float16) - assert out[0, 0].item() == 0.0 - assert out[0, 1].item() == float("-inf") - - -@cpu_test -@pytest.mark.cpu -def test_additive_float_mask_requires_an_explicit_dtype() -> None: - """The float32 default is the bug; keeping the parameter required is the fix.""" - with pytest.raises(TypeError): - _additive_float_mask(torch.tensor([[True]])) # type: ignore[call-arg] - - -@cpu_test -@pytest.mark.cpu -@pytest.mark.parametrize("dtype", ALL_DTYPES) -def test_generator_model_dtype_tracks_its_weights(dtype: torch.dtype) -> None: - """The single source of truth the three former float32 literals now defer to.""" - assert OmniVoiceGenerator(_tiny_config()).to(dtype).model_dtype == dtype - - @cpu_test @pytest.mark.cpu -@pytest.mark.parametrize("dtype", HALF_DTYPES) -def test_attention_accepts_a_model_dtype_mask(dtype: torch.dtype) -> None: - attn = OmniVoiceAttention(_tiny_config()).to(dtype).eval() - hidden, bool_mask, rope_table = _inputs(dtype, torch.device("cpu")) - - with torch.inference_mode(): - out = attn(hidden, rope_table, attention_mask=_additive_float_mask(bool_mask, dtype)) - - assert out.shape == hidden.shape - assert out.dtype == dtype - assert torch.isfinite(out).all() - - -# -------------------------------------------------------------------------- -# The crash itself (CUDA only — the CPU backend does not reject the mismatch) -# -------------------------------------------------------------------------- - - -@cuda_test -@pytest.mark.parametrize("dtype", HALF_DTYPES) -def test_float32_mask_is_what_broke_half_precision(dtype: torch.dtype) -> None: - """Pin the original failure so a float32 default cannot come back unnoticed.""" - attn = OmniVoiceAttention(_tiny_config()).to(device="cuda:0", dtype=dtype).eval() - hidden, bool_mask, rope_table = _inputs(dtype, torch.device("cuda:0")) +def test_sdpa_fallback_mask_is_dtype_independent() -> None: + """The fallback mask is boolean rather than an additive model-dtype mask.""" + cu_seqs = torch.tensor([0, 2, 3, 6, 6], dtype=torch.int32) + + metadata = _attention_metadata_from_cu_seqs( + cu_seqs, + 6, + needs_sdpa_mask=True, + ) - with pytest.raises(RuntimeError, match=r"(invalid dtype for bias|attn_mask)"), torch.inference_mode(): - attn(hidden, rope_table, attention_mask=_additive_float_mask(bool_mask, torch.float32)) + assert metadata.attn_mask is not None + assert metadata.attn_mask.dtype == torch.bool @cuda_test @pytest.mark.parametrize("dtype", HALF_DTYPES) -def test_real_path_forward_runs_in_half_precision(dtype: torch.dtype) -> None: - """The regression, on the production path with the Triton norm kernels live.""" - attn = OmniVoiceAttention(_tiny_config()).to(device="cuda:0", dtype=dtype).eval() - hidden, bool_mask, rope_table = _inputs(dtype, torch.device("cuda:0")) - - with torch.inference_mode(): - out = attn(hidden, rope_table, attention_mask=_additive_float_mask(bool_mask, dtype)) - - assert out.dtype == dtype - assert torch.isfinite(out).all() - +def test_sdpa_fallback_mask_needs_no_half_precision_coercion( + dtype: torch.dtype, +) -> None: + """Half-precision SDPA accepts the production bool mask without casting.""" + device = torch.device("cuda:0") + seq_len = 6 + cu_seqs = torch.tensor([0, 2, 3, 6, 6], dtype=torch.int32, device=device) -@cuda_test -@pytest.mark.parametrize("dtype", HALF_DTYPES) -def test_generator_forward_runs_in_half_precision(dtype: torch.dtype) -> None: - """End-to-end over the real iterative loop, which is where the float32 mask was built. + metadata = _attention_metadata_from_cu_seqs( + cu_seqs, + seq_len, + needs_sdpa_mask=True, + ) - This is the test that fails on an unfixed tree: ``forward`` materialized the - shared SDPA mask as float32 regardless of the weights' dtype. - """ - device = torch.device("cuda:0") - config = _tiny_config() - generator = OmniVoiceGenerator(config).to(device=device, dtype=dtype).eval() + assert metadata.attn_mask is not None + assert metadata.attn_mask.dtype == torch.bool - seq_len, target_len = 12, 4 - input_ids = torch.zeros(2, config.num_audio_codebook, seq_len, dtype=torch.long, device=device) - input_ids[:, 1:, :] = config.audio_mask_id - audio_mask = torch.zeros(2, seq_len, dtype=torch.bool, device=device) - audio_mask[:, seq_len - target_len :] = True - attention_mask = torch.ones(2, 1, seq_len, seq_len, dtype=torch.bool, device=device) + query = torch.randn(1, 2, seq_len, 8, device=device, dtype=dtype) + key = torch.randn_like(query) + value = torch.randn_like(query) with torch.inference_mode(): - tokens = generator( - input_ids, - audio_mask, - attention_mask, - target_lens=[target_len], - seed=0, - num_step=2, + output = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=metadata.attn_mask, ) - assert tokens.shape == (1, config.num_audio_codebook, target_len) - assert tokens.dtype == torch.long + assert output.dtype == dtype + assert torch.isfinite(output).all() diff --git a/tests/model_executor/models/omnivoice/test_pipeline_batching.py b/tests/model_executor/models/omnivoice/test_pipeline_batching.py new file mode 100644 index 00000000000..6e611f2ff09 --- /dev/null +++ b/tests/model_executor/models/omnivoice/test_pipeline_batching.py @@ -0,0 +1,228 @@ +# 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.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.models.omnivoice.pipeline_omnivoice import ( + OmniVoicePipeline, + _PreparedOmniVoiceRequest, +) +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import ( + _attention_metadata_from_cu_seqs, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_sdpa_fallback_mask_preserves_packed_sequence_boundaries() -> None: + cu_seqs = torch.tensor([0, 2, 3, 6, 6], dtype=torch.int32) + + metadata = _attention_metadata_from_cu_seqs(cu_seqs, 6, needs_sdpa_mask=True) + + assert metadata.attn_mask is not None + expected = torch.tensor( + [ + [1, 1, 0, 0, 0, 0], + [1, 1, 0, 0, 0, 0], + [0, 0, 1, 0, 0, 0], + [0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 1, 1], + [0, 0, 0, 1, 1, 1], + ], + dtype=torch.bool, + ) + torch.testing.assert_close(metadata.attn_mask[0, 0], expected) + + +def test_eager_attention_metadata_honors_exact_max_seqlen() -> None: + cu_seqs = torch.tensor([0, 2, 3, 6, 6], dtype=torch.int32) + + metadata = _attention_metadata_from_cu_seqs( + cu_seqs, + 6, + needs_sdpa_mask=False, + max_seqlen=3, + ) + + assert metadata.extra["max_seqlen_q"] == 3 + assert metadata.extra["max_seqlen_k"] == 3 + + +def test_request_batch_prepare_error_does_not_skip_valid_request() -> None: + """A malformed request must not prevent another request from completing.""" + prepared = _PreparedOmniVoiceRequest( + input_ids=torch.zeros(5, 8, dtype=torch.long), + audio_mask=torch.ones(5, dtype=torch.bool), + cond_len=3, + target_len=2, + seed=42, + ) + + def prepare_request(prompt, extra): + del extra + if prompt == "bad": + return DiffusionOutput(error="invalid prompt") + return prepared + + def collate_requests(requests): + assert requests == [prepared] + return prepared.input_ids, prepared.audio_mask, [prepared.cond_len] + + generator_calls = [] + + def generate_tokens(**kwargs): + generator_calls.append(kwargs) + assert kwargs["target_lens"] == [prepared.target_len] + return torch.ones(1, 8, prepared.target_len, dtype=torch.long) + + pipeline = SimpleNamespace( + _prepare_request_input=prepare_request, + _collate_request_inputs=collate_requests, + generator=generate_tokens, + decoder=lambda tokens: tokens.float(), + num_step=32, + guidance_scale=7.0, + t_shift=1.0, + layer_penalty_factor=0.0, + position_temperature=0.0, + class_temperature=0.0, + ) + sampling = OmniDiffusionSamplingParams(num_inference_steps=32, guidance_scale=7.0) + batch = DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest(prompt="bad", sampling_params=sampling, request_id="bad"), + OmniDiffusionRequest( + prompt="good", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=32, guidance_scale=7.0), + request_id="good", + ), + ] + ) + + outputs = OmniVoicePipeline.forward(pipeline, batch) + + assert len(outputs) == 2 + assert outputs[0].error == "invalid prompt" + assert outputs[1].error is None + torch.testing.assert_close(outputs[1].output, torch.ones(1, 8, 2)) + assert generator_calls[0]["guidance_scale"] == 7.0 + + +def test_request_batch_honors_explicit_sampling_overrides() -> None: + """Request steps and guidance must override the OmniVoice defaults.""" + prepared = _PreparedOmniVoiceRequest( + input_ids=torch.zeros(5, 8, dtype=torch.long), + audio_mask=torch.ones(5, dtype=torch.bool), + cond_len=3, + target_len=2, + seed=42, + ) + captured_sampling = [] + + def generate_tokens(**kwargs): + captured_sampling.append((kwargs["num_step"], kwargs["guidance_scale"])) + return torch.ones(1, 8, prepared.target_len, dtype=torch.long) + + pipeline = SimpleNamespace( + _prepare_request_input=lambda prompt, extra: prepared, + _collate_request_inputs=lambda requests: (prepared.input_ids, prepared.audio_mask, [prepared.cond_len]), + generator=generate_tokens, + decoder=lambda tokens: tokens.float(), + num_step=32, + guidance_scale=2.0, + t_shift=1.0, + layer_penalty_factor=0.0, + position_temperature=0.0, + class_temperature=0.0, + ) + batch = DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest( + prompt="hello", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=12, guidance_scale=6.5), + request_id="request", + ) + ] + ) + + outputs = OmniVoicePipeline.forward(pipeline, batch) + + assert outputs[0].error is None + assert captured_sampling == [(12, 6.5)] + + +def test_request_batch_prepare_error_preserves_surrounding_output_indices() -> None: + """A middle prepare error must not shift valid outputs into the wrong slots.""" + prepared_before = _PreparedOmniVoiceRequest( + input_ids=torch.zeros(3, 8, dtype=torch.long), + audio_mask=torch.ones(3, dtype=torch.bool), + cond_len=2, + target_len=1, + seed=1, + ) + prepared_after = _PreparedOmniVoiceRequest( + input_ids=torch.zeros(5, 8, dtype=torch.long), + audio_mask=torch.ones(5, dtype=torch.bool), + cond_len=3, + target_len=2, + seed=2, + ) + + def prepare_request(prompt, extra): + del extra + if prompt == "before": + return prepared_before + if prompt == "after": + return prepared_after + return DiffusionOutput(error="invalid middle request") + + def collate_requests(requests): + assert requests == [prepared_before, prepared_after] + return ( + torch.cat([prepared_before.input_ids, prepared_after.input_ids]), + torch.cat([prepared_before.audio_mask, prepared_after.audio_mask]), + [prepared_before.cond_len, prepared_after.cond_len], + ) + + def generate_tokens(**kwargs): + assert kwargs["target_lens"] == [1, 2] + before = torch.full((1, 8, 1), 11, dtype=torch.long) + after = torch.full((1, 8, 2), 22, dtype=torch.long) + return torch.cat([before, after], dim=-1) + + pipeline = SimpleNamespace( + _prepare_request_input=prepare_request, + _collate_request_inputs=collate_requests, + generator=generate_tokens, + decoder=lambda tokens: tokens.float(), + num_step=32, + guidance_scale=2.0, + t_shift=1.0, + layer_penalty_factor=0.0, + position_temperature=0.0, + class_temperature=0.0, + ) + batch = DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest( + prompt=prompt, + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=32), + request_id=prompt, + ) + for prompt in ("before", "bad", "after") + ] + ) + + outputs = OmniVoicePipeline.forward(pipeline, batch) + + assert len(outputs) == 3 + torch.testing.assert_close(outputs[0].output, torch.full((1, 8, 1), 11.0)) + assert outputs[1].error == "invalid middle request" + torch.testing.assert_close(outputs[2].output, torch.full((1, 8, 2), 22.0)) diff --git a/vllm_omni/deploy/omnivoice.yaml b/vllm_omni/deploy/omnivoice.yaml index 6769dd5e995..df9e6a5daf0 100644 --- a/vllm_omni/deploy/omnivoice.yaml +++ b/vllm_omni/deploy/omnivoice.yaml @@ -11,3 +11,5 @@ stages: trust_remote_code: true distributed_executor_backend: "mp" dtype: "float32" + max_num_seqs: 8 + request_batch_max_wait_ms: 10 diff --git a/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py b/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py index cfc40fcb6bc..bad1c6154e7 100644 --- a/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py +++ b/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py @@ -12,10 +12,13 @@ from __future__ import annotations import json +import math import os +import random import re -from collections.abc import Iterable -from typing import ClassVar +from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from typing import Any, ClassVar import numpy as np import torch @@ -26,10 +29,17 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.models.interface import SupportAudioOutput +from vllm_omni.diffusion.worker.input_batch import InputBatch from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.diffusion.worker.utils import StepRequestState +from vllm_omni.errors import OmniClientError from vllm_omni.model_executor.models.omnivoice.duration import RuleDurationEstimator from vllm_omni.model_executor.models.omnivoice.omnivoice_decoder import OmniVoiceDecoder -from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import OmniVoiceGenerator +from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import ( + OmniVoiceGenerator, + _build_cu_seqs, + _get_time_steps, +) from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig from vllm_omni.utils.speaker_cache import get_speaker_cache @@ -43,6 +53,15 @@ logger = init_logger(__name__) +@dataclass +class _PreparedOmniVoiceRequest: + input_ids: torch.Tensor + audio_mask: torch.Tensor + cond_len: int + target_len: int + seed: int | None + + def get_omnivoice_post_process_func(od_config: OmniDiffusionConfig): """Post-processing: convert audio tensor to numpy for WAV encoding.""" @@ -132,6 +151,8 @@ class OmniVoicePipeline(nn.Module, SupportAudioOutput): """ support_audio_output: ClassVar[bool] = True + supports_request_batch: ClassVar[bool] = True + supports_step_execution: ClassVar[bool] = True def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): super().__init__() @@ -152,7 +173,7 @@ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): self.config = OmniVoiceConfig(**hf_config) # Build generator and decoder - self.generator = OmniVoiceGenerator(self.config) + self.generator = OmniVoiceGenerator(self.config, od_config) self.decoder = OmniVoiceDecoder(self.config) # Tokenizer (low-level, avoids HF tokenizer extra_special_tokens issue) @@ -205,50 +226,36 @@ def _encode_ref_audio(self, audio_signal: torch.Tensor, sr: int) -> torch.Tensor tokens = tokens.squeeze(0) # [8, T_ref] return tokens - @torch.inference_mode() - def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: - """Generate speech audio from text, optionally with voice cloning. - - Accepts either a plain text prompt or a structured dict: - {"text": "...", "ref_audio": (samples, sr), "ref_text": "...", - "lang": "...", "instruct": "..."} - """ - prompt = req.prompts[0] if req.prompts else "" + def _prepare_request_input( + self, + prompt: Any, + extra: dict[str, Any], + ) -> _PreparedOmniVoiceRequest | DiffusionOutput: + """Build one request's conditional/unconditional model inputs.""" ref_audio = None ref_text = None lang = "None" instruct = "None" - extra = req.sampling_params.extra_args or {} + voice_name = None seed = extra.get("seed", None) - voice_name = None if isinstance(prompt, dict): - # Top-level keys (used by serving_speech.py /v1/audio/speech path) text = prompt.get("input") or prompt.get("text") or prompt.get("prompt") ref_audio = prompt.get("ref_audio") ref_text = prompt.get("ref_text") voice_name = prompt.get("voice_name") lang = prompt.get("lang") instruct = prompt.get("instruct") - # OmniTextPrompt format (used by offline Omni.generate path): - # ref_audio comes via multi_modal_data["audio"] and the rest via - # mm_processor_kwargs. Fall back to those when top-level keys are - # absent so both invocation styles work. mm_data = prompt.get("multi_modal_data") or {} mm_kwargs = prompt.get("mm_processor_kwargs") or {} if ref_audio is None: audio_field = mm_data.get("audio") - # Standard multimodal shape allows a list of audios; OmniVoice - # voice cloning conditions on a single reference clip, so - # unwrap a length-1 list and reject multi-reference prompts up - # front (otherwise a list would later crash inside - # ``_encode_ref_audio`` when it calls ``audio.dim()``). if isinstance(audio_field, list): if len(audio_field) == 1: audio_field = audio_field[0] elif len(audio_field) > 1: return DiffusionOutput( - error=f"OmniVoice voice cloning supports a single reference audio; got {len(audio_field)}" # noqa: E501 + error=f"OmniVoice voice cloning supports a single reference audio; got {len(audio_field)}" ) else: audio_field = None @@ -264,7 +271,6 @@ def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: lang = mm_kwargs.get("lang") if instruct is None: instruct = mm_kwargs.get("instruct") - if not text: return DiffusionOutput(error="Empty text prompt") lang = lang or "None" @@ -274,105 +280,325 @@ def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: if not text: return DiffusionOutput(error="Empty text prompt") - device = self.device - num_cb = self.config.num_audio_codebook - mask_id = self.config.audio_mask_id - - # Estimate target duration target_len = self.duration_estimator.estimate_duration(text, "Nice to meet you.", 25) target_len = max(1, int(target_len)) - # Build text prompt with control tokens style_text = f"<|denoise|><|lang_start|>{lang}<|lang_end|><|instruct_start|>{instruct}<|instruct_end|>" full_text = _combine_text(ref_text=ref_text, text=text) wrapped_text = f"<|text_start|>{full_text}<|text_end|>" style_tokens = self.tokenizer.encode(style_text).ids text_tokens = _tokenize_with_nonverbal_tags(wrapped_text, self.tokenizer) encoding_ids = style_tokens + text_tokens - text_tokens = torch.tensor(encoding_ids, dtype=torch.long, device=device) - text_len = text_tokens.shape[0] + text_tokens_tensor = torch.tensor(encoding_ids, dtype=torch.long, device=self.device) + text_len = text_tokens_tensor.shape[0] - # Encode reference audio tokens if provided (with voice caching) ref_audio_tokens = None if ref_audio is not None: if self.audio_tokenizer is None: raise RuntimeError( "Voice cloning requires transformers>=5.3.0. Try: uv pip install 'transformers>=5.3.0'" ) - # Check speaker cache first - _cache_key = None + cache_key = None if voice_name: - _cache_key = self._speaker_cache.make_cache_key( + cache_key = self._speaker_cache.make_cache_key( voice_name, model_type="omnivoice", created_at=int(prompt.get("voice_created_at") or 0), ) - cached = self._speaker_cache.get(_cache_key) + cached = self._speaker_cache.get(cache_key) if cached is not None: - ref_audio_tokens = cached["ref_audio_tokens"].to(device) - _cache_key = None # hit → don't store again + ref_audio_tokens = cached["ref_audio_tokens"].to(self.device) + cache_key = None logger.debug("Speaker cache HIT for OmniVoice speaker '%s'", voice_name) if ref_audio_tokens is None: audio_signal, sr = ref_audio if isinstance(audio_signal, np.ndarray): audio_signal = torch.from_numpy(audio_signal).float() - ref_audio_tokens = self._encode_ref_audio(audio_signal, int(sr)).to(device) - - # Store in cache for next request - if _cache_key is not None: - self._speaker_cache.put(_cache_key, {"ref_audio_tokens": ref_audio_tokens.cpu()}) + ref_audio_tokens = self._encode_ref_audio(audio_signal, int(sr)).to(self.device) + if cache_key is not None: + self._speaker_cache.put(cache_key, {"ref_audio_tokens": ref_audio_tokens.cpu()}) logger.debug("Speaker cache STORE for OmniVoice speaker '%s'", voice_name) - # Build conditional + unconditional batches [2, 8, max_len] - text_ids = text_tokens.unsqueeze(0).repeat(num_cb, 1) - target_ids = torch.full((num_cb, target_len), mask_id, dtype=torch.long, device=device) - - if ref_audio_tokens is not None: - cond_ids = torch.cat([text_ids, ref_audio_tokens, target_ids], dim=1) - else: - cond_ids = torch.cat([text_ids, target_ids], dim=1) + num_cb = self.config.num_audio_codebook + mask_id = self.config.audio_mask_id + text_ids = text_tokens_tensor.unsqueeze(0).repeat(num_cb, 1) + target_ids = torch.full((num_cb, target_len), mask_id, dtype=torch.long, device=self.device) + cond_ids = ( + torch.cat([text_ids, ref_audio_tokens, target_ids], dim=1) + if ref_audio_tokens is not None + else torch.cat([text_ids, target_ids], dim=1) + ) cond_len = cond_ids.shape[1] uncond_ids = target_ids.clone() - uncond_len = target_len - max_len = max(cond_len, uncond_len) - if uncond_len < max_len: - pad = torch.full( - (num_cb, max_len - uncond_len), - mask_id, - dtype=torch.long, - device=device, + input_ids = torch.cat([cond_ids, uncond_ids], dim=1).transpose(0, 1).contiguous() + + max_len = input_ids.shape[0] + audio_mask = torch.zeros(max_len, dtype=torch.bool, device=self.device) + audio_mask[text_len:] = True + + return _PreparedOmniVoiceRequest( + input_ids=input_ids, + audio_mask=audio_mask, + cond_len=cond_len, + target_len=target_len, + seed=seed, + ) + + def _collate_request_inputs( + self, + prepared_requests: Sequence[_PreparedOmniVoiceRequest], + ) -> tuple[torch.Tensor, torch.Tensor, list[int]]: + """Pack request-major [cond, uncond] token sequences.""" + input_ids: list[torch.Tensor] = [] + audio_masks: list[torch.Tensor] = [] + cond_lens: list[int] = [] + + for request in prepared_requests: + input_id = request.input_ids + audio_mask = request.audio_mask + cond_len = request.cond_len + input_ids.append(input_id) + audio_masks.append(audio_mask) + cond_lens.append(cond_len) + + input_ids = torch.cat(input_ids, dim=0) + audio_masks = torch.cat(audio_masks, dim=0) + + return input_ids, audio_masks, cond_lens + + def prepare_encode(self, state: StepRequestState) -> StepRequestState: + prompt = state.prompt if state.prompt else "" + extra = state.sampling.extra_args or {} + prepared = self._prepare_request_input(prompt, extra) + if isinstance(prepared, DiffusionOutput): + raise OmniClientError(prepared.error or "OmniVoice request preparation failed") + + prepared_request = prepared + cond_len = prepared_request.cond_len + target_len = prepared_request.target_len + input_ids = prepared_request.input_ids + audio_mask = prepared_request.audio_mask + seed = prepared_request.seed + device = self.device + mask_id = self.config.audio_mask_id + num_codebooks = self.config.num_audio_codebook + if seed is None: + seed = random.randint(0, 2**63 - 1) + num_step = ( + state.sampling.num_inference_steps if state.sampling.num_inference_steps is not None else self.num_step + ) + + t_shift = self.t_shift + + # Initialize all target tokens as [MASK] + tokens = torch.full((1, num_codebooks, target_len), mask_id, dtype=torch.long, device=device) + + timesteps = _get_time_steps(0.0, 1.0, num_step + 1, t_shift) + + # Compute unmasking schedule + schedules = [] + total_mask = target_len * num_codebooks + rem = total_mask + sched = [] + for step in range(num_step): + num = ( + rem + if step == num_step - 1 + else min( + math.ceil(total_mask * (timesteps[step + 1] - timesteps[step])), + rem, + ) ) - uncond_ids = torch.cat([uncond_ids, pad], dim=1) + sched.append(int(num)) + rem -= int(num) + schedules = torch.tensor(sched, dtype=torch.long, device=device) - batch_input_ids = torch.stack([cond_ids, uncond_ids]) + layer_ids = torch.arange(num_codebooks, device=device).view(1, -1, 1) + generator = torch.Generator(device=device).manual_seed(seed) - batch_audio_mask = torch.zeros(2, max_len, dtype=torch.bool, device=device) - batch_audio_mask[0, text_len:cond_len] = True - batch_audio_mask[1, :uncond_len] = True + guidance_scale = ( + state.sampling.guidance_scale if state.sampling.guidance_scale is not None else self.guidance_scale + ) + state.latents = input_ids + state.timesteps = schedules + state.guidance = guidance_scale + state.extra["schedules"] = schedules + state.extra["layer_ids"] = layer_ids + state.extra["generator"] = generator + state.extra["t_shift"] = t_shift + state.extra["cond_len"] = cond_len + state.extra["target_len"] = target_len + state.extra["audio_mask"] = audio_mask + state.extra["tokens"] = tokens + return state + + def denoise_step(self, input_batch: InputBatch, *, states: Sequence[StepRequestState] | None = None, **kwargs: Any): + input_ids = input_batch.latents + use_cuda_graph = self.generator._cuda_graph_fwd is not None and input_ids.is_cuda + layer_ids = states[0].extra["layer_ids"] + + audio_masks: list[torch.Tensor] = [] + target_lens: list[int] = [] + batch_tokens: list[torch.Tensor] = [] + cond_lens: list[int] = [] + + steps: list[int] = [] + schedules: list[torch.Tensor] = [] + generators: list[torch.Generator] = [] + guidance_scales: list[float] = [] + + for state in states: + audio_masks.append(state.extra.get("audio_mask", None)) + cond_lens.append(state.extra["cond_len"]) + target_lens.append(state.extra["target_len"]) + batch_tokens.append(state.extra["tokens"]) + guidance_scales.append(state.guidance) + generators.append(state.extra.get("generator", None)) + schedules.append(state.extra["schedules"]) + steps.append(state.step_index) + + audio_masks = torch.cat(audio_masks, dim=0) + + B = len(target_lens) + cu_seqs = _build_cu_seqs(cond_lens, target_lens, input_ids.device) + + position_temperature = self.position_temperature + class_temperature = self.class_temperature + layer_penalty_factor = self.layer_penalty_factor + if use_cuda_graph: + # Replay a fixed packed-token bucket with dynamic varlen metadata. + batch_logits = self.generator._cuda_graph_fwd(input_ids, audio_masks, cu_seqs, B) + else: + # Run packed eager attention for the current active requests. + inputs_embeds = self.generator._prepare_embeddings(input_ids, audio_masks) + hidden_states = self.generator._transformer_forward( + inputs_embeds, + cu_seqs, + max_seqlen=max(cond_lens), + ) + # fp32 cast deferred to the per-item slices below. + batch_logits = self.generator._get_logits(hidden_states) + # batch_logits: [8, total_seq_len, 1025] + + target_offsets: list[int] = [] + target_offset = 0 + for target_len in target_lens: + target_offsets.append(target_offset) + target_offset += target_len + + sequence_offsets: list[int] = [] + sequence_offset = 0 + for cond_len, target_len in zip(cond_lens, target_lens): + sequence_offsets.append(sequence_offset) + sequence_offset += cond_len + target_len + + for i in range(B): + k = schedules[i][steps[i]] + if k <= 0: + continue + + c_len = cond_lens[i] + t_len = target_lens[i] + + # Extract logits for target region; upcast only the slices we actually consume. + request_start = sequence_offsets[i] + cond_end = request_start + c_len + uncond_start = cond_end + + # Extract logits for target region; upcast only the slices we actually consume. + c_logits = batch_logits[:, cond_end - t_len : cond_end, :].unsqueeze(0).to(torch.float32) + u_logits = batch_logits[:, uncond_start : uncond_start + t_len, :].unsqueeze(0).to(torch.float32) + sample = batch_tokens[i] + sample_tokens = sample[..., :t_len] + self.generator._unmask_one_request( + c_logits, + u_logits, + sample_tokens, + num_to_unmask=k, + guidance_scale=guidance_scales[i], + generator=generators[i], + class_temperature=class_temperature, + position_temperature=position_temperature, + layer_penalty_factor=layer_penalty_factor, + layer_ids=layer_ids, + ) + + # Mirror update into both cond and uncond input_ids halves for the next step. + packed_sample_tokens = sample_tokens.squeeze(0).transpose(0, 1) + input_ids[cond_end - t_len : cond_end] = packed_sample_tokens + input_ids[uncond_start : uncond_start + t_len] = packed_sample_tokens + states[i].extra["tokens"] = sample_tokens + + # InputBatch reuses its latents buffer across steps. Returning that + # same storage would make the Runner persist per-request views into the + # cached destination; the next make_batch() would then copy overlapping + # source/destination slices. Break the alias at the lifecycle boundary. + return input_ids.clone() + + def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor, **kwargs: Any): + state.latents = noise_pred + state.step_index += 1 + + def post_decode(self, state: StepRequestState, **kwargs: Any): + tokens = state.extra["tokens"] + if tokens.dim() == 2: + tokens = tokens.unsqueeze(0) + audio = self.decoder(tokens) + return DiffusionOutput(output=audio) - batch_attn_mask = torch.zeros(2, 1, max_len, max_len, dtype=torch.bool, device=device) - batch_attn_mask[0, :, :cond_len, :cond_len] = True - batch_attn_mask[1, :, :uncond_len, :uncond_len] = True + @torch.inference_mode() + def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: + """Generate speech audio from text, optionally with voice cloning. + Accepts either a plain text prompt or a structured dict: + {"text": "...", "ref_audio": (samples, sr), "ref_text": "...", + "lang": "...", "instruct": "..."} + """ + prepared_requests: list[_PreparedOmniVoiceRequest] = [] + outputs = [None] * len(req.requests) + prepared_indices: list[int] = [] + for i, request in enumerate(req.requests): + prompt = request.prompt if request.prompt else "" + extra = request.sampling_params.extra_args or {} + prepared = self._prepare_request_input(prompt, extra) + if isinstance(prepared, DiffusionOutput): + outputs[i] = prepared + continue + prepared_indices.append(i) + prepared_requests.append(prepared) + + if not prepared_requests: + return outputs + + batch_target_len = [request.target_len for request in prepared_requests] + batch_seeds = [request.seed for request in prepared_requests] + batch_input_ids, batch_audio_mask, batch_cond_lens = self._collate_request_inputs(prepared_requests) # Run 32-step iterative unmasking + sampling = req.requests[0].sampling_params + num_step = sampling.num_inference_steps if sampling.num_inference_steps is not None else self.num_step + guidance_scale = sampling.guidance_scale if sampling.guidance_scale is not None else self.guidance_scale tokens = self.generator( input_ids=batch_input_ids, audio_mask=batch_audio_mask, - attention_mask=batch_attn_mask, - target_lens=[target_len], - num_step=self.num_step, - guidance_scale=self.guidance_scale, + cond_lens=batch_cond_lens, + target_lens=batch_target_len, + num_step=num_step, + guidance_scale=guidance_scale, t_shift=self.t_shift, layer_penalty_factor=self.layer_penalty_factor, position_temperature=self.position_temperature, class_temperature=self.class_temperature, - seed=seed, + seed=batch_seeds, ) - # Decode tokens to audio - audio = self.decoder(tokens) # [1, 1, samples] - return DiffusionOutput(output=audio) + target_offset = 0 + for i, target_len in enumerate(batch_target_len): + request_tokens = tokens[:, :, target_offset : target_offset + target_len] + audio = self.decoder(request_tokens) + outputs[prepared_indices[i]] = DiffusionOutput(output=audio) + target_offset += target_len + return outputs def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights from model directory (not from the iterator). diff --git a/vllm_omni/diffusion/worker/utils.py b/vllm_omni/diffusion/worker/utils.py index 5158dbfbf37..efc935a5dbd 100644 --- a/vllm_omni/diffusion/worker/utils.py +++ b/vllm_omni/diffusion/worker/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 """Per-request mutable state for step-wise diffusion execution.""" from __future__ import annotations diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 50f459a5863..75a8a435e94 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -2113,6 +2113,25 @@ async def _create_diffusion_speech( if sampling_params_list[0].extra_args is None: sampling_params_list[0].extra_args = {} sampling_params_list[0].extra_args.update(extra) + + sampling = sampling_params_list[0] + + # This change allows StepScheduler read total_steps from upper + # sampling.num_inference_steps, check diffusion/sched/step_scheduler:_get_total_steps + if "num_inference_steps" in extra: + value = extra["num_inference_steps"] + try: + sampling.num_inference_steps = int(value) + except (TypeError, ValueError) as exc: + raise ValueError("num_inference_steps must be an integer") from exc + + if "guidance_scale" in extra: + value = extra["guidance_scale"] + try: + sampling.guidance_scale = float(value) + except (TypeError, ValueError) as exc: + raise ValueError("guidance_scale must be a number") from exc + logger.info("Applied extra_params to diffusion: %s", extra) generator = self._diffusion_engine.generate( diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py b/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py index a75c4e6f7fc..965d021b82f 100644 --- a/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py +++ b/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py @@ -6,9 +6,8 @@ Generates 8-codebook audio tokens from text via 32-step non-autoregressive iterative masked prediction with classifier-free guidance. -Uses full bidirectional attention computed directly with PyTorch SDPA -(torch.nn.functional.scaled_dot_product_attention); no auto-selected -FlashAttention/SageAttention/DiffusionAttention backend is used. +Uses backend-dispatched variable-length full bidirectional attention over +packed conditional and unconditional sequences, with an SDPA mask fallback. """ from __future__ import annotations @@ -22,7 +21,11 @@ import torch.nn as nn import torch.nn.functional as F from vllm.logger import init_logger +from vllm.utils.math_utils import round_up +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.model_executor.models.omnivoice.fused_qkv_rope import fused_qkv_norm_rope from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig @@ -241,7 +244,7 @@ def _gumbel_sample(logits: torch.Tensor, temperature: float, generator: torch.Ge # --------------------------------------------------------------------------- -# Qwen3-style transformer blocks using PyTorch SDPA +# Qwen3-style transformer blocks using PyTorch variable-length attention # --------------------------------------------------------------------------- @@ -261,9 +264,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class OmniVoiceAttention(nn.Module): - """Qwen3-style GQA attention using PyTorch SDPA (full bidirectional).""" + """Qwen3-style GQA using packed, full-bidirectional varlen attention.""" - def __init__(self, config: OmniVoiceConfig): + def __init__(self, config: OmniVoiceConfig, layer_idx: int): super().__init__() self.hidden_size = config.llm_hidden_size self.num_heads = config.llm_num_attention_heads @@ -282,16 +285,27 @@ def __init__(self, config: OmniVoiceConfig): self.k_norm = OmniVoiceRMSNorm(self.head_dim) self.scale = 1.0 / math.sqrt(self.head_dim) + self.attention_op = Attention( + num_heads=self.num_heads, + num_kv_heads=self.num_heads, + head_size=self.head_dim, + causal=False, + softmax_scale=self.scale, + prefix=f"layers.{layer_idx}.self_attn.attention_op", + qkv_layout="BSND", + skip_sequence_parallel=True, + ) + self.use_packed_varlen = self.attention_op.attn_backend.supports_multi_doc_packed_varlen() def forward( self, hidden_states: torch.Tensor, rope_table: torch.Tensor, - attention_mask: torch.Tensor | None = None, + attn_metadata: AttentionMetadata, ) -> torch.Tensor: - batch_size, seq_len, _ = hidden_states.shape + seq_len, _ = hidden_states.shape - qkv = self.qkv_proj(hidden_states).view(batch_size, seq_len, self.num_qkv_heads, self.head_dim) + qkv = self.qkv_proj(hidden_states).view(1, seq_len, self.num_qkv_heads, self.head_dim) # One kernel for the whole prologue: split the packed projection, RMSNorm # Q and K per head, rotate both, broadcast K and V across their query @@ -306,22 +320,16 @@ def forward( self.num_kv_heads, ) - # Caller passes a float mask; materialize float form if a bool slips through. - sdpa_mask = attention_mask - if sdpa_mask is not None and sdpa_mask.dtype == torch.bool: - sdpa_mask = torch.zeros_like(attention_mask, dtype=q.dtype).masked_fill_(~attention_mask, float("-inf")) - - out = F.scaled_dot_product_attention( - q, - k, - v, - attn_mask=sdpa_mask, - scale=self.scale, - ) + # The fused prologue emits [B, N, S, D]; Omni attention backends use + # [B, S, N, D]. Keep the conversion outside the backend. + q, k, v = (tensor.permute(0, 2, 1, 3).contiguous().to(torch.bfloat16) for tensor in (q, k, v)) + if self.use_packed_varlen: + out = self.attention_op(q, k, v, attn_metadata) + else: + out = self.attention_op.sdpa_fallback.forward(q, k, v, attn_metadata) - # Back to (batch, seq, heads * head_dim) - out = out.permute(0, 2, 1, 3).contiguous() - out = out.view(batch_size, seq_len, self.num_heads * self.head_dim) + out = out.squeeze(0).to(hidden_states.dtype) + out = out.reshape(seq_len, self.num_heads * self.head_dim) return self.o_proj(out) @@ -351,19 +359,19 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class OmniVoiceTransformerBlock(nn.Module): - """Single Qwen3 transformer block with PyTorch SDPA attention.""" + """Single Qwen3 transformer block with variable-length attention.""" - def __init__(self, config: OmniVoiceConfig): + def __init__(self, config: OmniVoiceConfig, layer_idx: int): super().__init__() self.input_layernorm = OmniVoiceRMSNorm(config.llm_hidden_size, eps=config.llm_rms_norm_eps) - self.self_attn = OmniVoiceAttention(config) + self.self_attn = OmniVoiceAttention(config, layer_idx) self.post_attention_layernorm = OmniVoiceRMSNorm(config.llm_hidden_size, eps=config.llm_rms_norm_eps) self.mlp = OmniVoiceMLP(config) def forward( self, hidden_states: torch.Tensor, - attention_mask: torch.Tensor | None = None, + attn_metadata: AttentionMetadata, rope_table: torch.Tensor | None = None, residual: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: @@ -391,7 +399,7 @@ def forward( residual = hidden_states hidden_states = self.input_layernorm(hidden_states) - hidden_states = self.self_attn(hidden_states, rope_table, attention_mask=attention_mask) + hidden_states = self.self_attn(hidden_states, rope_table, attn_metadata) if _TRITON_AVAILABLE: # Fused: (attn_out + residual) + RMSNorm in one kernel @@ -432,6 +440,63 @@ def _precompute_rope_table( return torch.cat([freqs.cos(), freqs.sin()], dim=-1) +def _position_ids_from_cu_seqs(cu_seqs: torch.Tensor, seq_len: int) -> torch.Tensor: + """Return each packed token's zero-based position within its sequence.""" + token_ids = torch.arange(seq_len, device=cu_seqs.device, dtype=cu_seqs.dtype) + sequence_ids = torch.searchsorted(cu_seqs[1:], token_ids, right=True) + sequence_starts = cu_seqs.index_select(0, sequence_ids.to(torch.long)) + return (token_ids - sequence_starts).to(torch.long) + + +def _attention_metadata_from_cu_seqs( + cu_seqs: torch.Tensor, + seq_len: int, + *, + needs_sdpa_mask: bool, + max_seqlen: int | None = None, +) -> AttentionMetadata: + """Build packed-varlen metadata and, when needed, an SDPA block mask.""" + # Local benchmarks found no measurable performance difference between the + # exact longest segment and the packed total. Eager uses the exact bound; + # CUDA Graph uses the fixed token-bucket maximum. + kernel_max_seqlen = seq_len if max_seqlen is None else max_seqlen + extra = { + "cu_seqlens_q": cu_seqs, + "cu_seqlens_k": cu_seqs, + "max_seqlen_q": kernel_max_seqlen, + "max_seqlen_k": kernel_max_seqlen, + } + if not needs_sdpa_mask: + return AttentionMetadata(extra=extra) + + token_ids = torch.arange(seq_len, device=cu_seqs.device, dtype=cu_seqs.dtype) + sequence_ids = torch.searchsorted(cu_seqs[1:], token_ids, right=True) + attn_mask = sequence_ids.view(1, 1, seq_len, 1) == sequence_ids.view(1, 1, 1, seq_len) + return AttentionMetadata(attn_mask=attn_mask, extra=extra) + + +def _build_cu_seqs( + cond_lens: list[int], + uncond_lens: list[int], + device: torch.device, + *, + tail_end: int | None = None, +) -> torch.Tensor: + """Build request-major [cond0, uncond0, ...] cumulative offsets.""" + if len(cond_lens) != len(uncond_lens): + raise ValueError(f"Mismatched cond/uncond lengths: {len(cond_lens)} != {len(uncond_lens)}.") + offsets = [0] + for cond_len, uncond_len in zip(cond_lens, uncond_lens): + offsets.append(offsets[-1] + cond_len) + offsets.append(offsets[-1] + uncond_len) + if tail_end is None: + tail_end = offsets[-1] + if offsets[-1] > tail_end: + raise ValueError(f"Packed length {offsets[-1]} exceeds tail end {tail_end}.") + offsets.append(tail_end) + return torch.tensor(offsets, device=device, dtype=torch.int32) + + # --------------------------------------------------------------------------- # TF32 opt-in (process-wide; default off) # --------------------------------------------------------------------------- @@ -461,23 +526,6 @@ def _maybe_enable_tf32() -> None: # --------------------------------------------------------------------------- -def _additive_float_mask(mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: - """Convert a boolean attention mask to its additive float form. - - ``True`` (attend) maps to ``0.0`` and ``False`` (masked) to ``-inf``. A bool - mask must never be copied straight into a float buffer: the implicit cast - maps True/False to 1.0/0.0, which leaves masked positions at 0.0 and so - silently *unmasks* them. - - ``dtype`` is required rather than defaulting to float32: SDPA rejects an - additive mask whose dtype differs from the query, so the mask has to follow - the model dtype and a default would just hide that coupling. - """ - if mask.dtype != torch.bool: - return mask - return torch.zeros_like(mask, dtype=dtype).masked_fill_(~mask, float("-inf")) - - class _OmniVoiceCUDAGraphForward: """Pre-captures CUDA graphs for predefined sequence-length buckets. @@ -486,24 +534,47 @@ class _OmniVoiceCUDAGraphForward: (one step at a time) means pool sharing is safe. """ - # Default bucket count is 10; 16 gives modest headroom for edge cases - # (seq_len > max bucket or non-CFG batch) without unbounded GPU growth. - _MAX_LAZY_GRAPHS: int = 16 + _MAX_LAZY_GRAPHS = 16 + _STATIC_CAPTURE_TOKEN_LIMIT = 1024 + _LAZY_CAPTURE_ALIGNMENT = 128 def __init__(self, generator: OmniVoiceGenerator, capture_sizes: list[int]) -> None: self._gen = generator - self._capture_sizes = sorted(capture_sizes) - # Pre-warmed graphs keyed by (two_b, bucket); fixed set, never evicted. + self.capture_batch_sizes = self._derive_capture_batch_size() + self.capture_bucket_sizes_by_batch = self._derive_capture_bucket_sizes(capture_sizes) self._graphs: dict[tuple[int, int], dict] = {} - # Lazy-captured graphs for oversized / non-CFG shapes; capped via LRU. self._lazy_graphs: OrderedDict[tuple[int, int], dict] = OrderedDict() self._lock = threading.Lock() # Per-instance pool handle: isolates OmniVoice CUDA memory from other # vllm modules while still allowing safe re-use across sequential replays. self._pool_handle: int | None = None - def _find_bucket(self, seq_len: int) -> int | None: - for bucket in self._capture_sizes: + def _derive_capture_batch_size(self) -> list[int]: + return list(range(1, self._gen.od_config.max_num_seqs + 1)) + + def _derive_capture_bucket_sizes(self, capture_sizes: list[int]) -> dict[int, list[int]]: + """Build a triangular, batch-aware token-bucket capture plan.""" + base_sizes = sorted(set(capture_sizes)) + if not base_sizes: + return {batch_size: [] for batch_size in self.capture_batch_sizes} + alignment = base_sizes[0] + single_request_cap = min(base_sizes[-1], 512) + max_graph_tokens = min(base_sizes[-1], self._STATIC_CAPTURE_TOKEN_LIMIT) + growth_per_batch = max(alignment, single_request_cap // 2) + candidates = sorted(set(base_sizes) | set(range(alignment, max_graph_tokens + alignment, alignment))) + return { + batch_size: [ + bucket + for bucket in candidates + if alignment * batch_size + <= bucket + <= min(single_request_cap + (batch_size - 1) * growth_per_batch, max_graph_tokens) + ] + for batch_size in self.capture_batch_sizes + } + + def _find_bucket(self, batch_size: int, seq_len: int) -> int | None: + for bucket in self.capture_bucket_sizes_by_batch.get(batch_size, ()): if bucket >= seq_len: return bucket return None @@ -512,57 +583,36 @@ def _pad_inputs( self, input_ids: torch.Tensor, audio_mask: torch.Tensor, - attention_mask: torch.Tensor | None, bucket: int, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - S = input_ids.shape[-1] - if S == bucket: - return input_ids, audio_mask, attention_mask - - two_b = input_ids.shape[0] - num_cb = input_ids.shape[1] - - ids_padded = torch.zeros(two_b, num_cb, bucket, dtype=input_ids.dtype, device=input_ids.device) - ids_padded[:, :, :S] = input_ids - - mask_padded = torch.zeros(two_b, bucket, dtype=torch.bool, device=audio_mask.device) - mask_padded[:, :S] = audio_mask - - if attention_mask is not None: - # Callers normalize to the additive float form first, so pad with -inf. - attn_padded = torch.full( - (two_b, 1, bucket, bucket), - float("-inf"), - dtype=attention_mask.dtype, - device=attention_mask.device, - ) - attn_padded[:, :, :S, :S] = attention_mask - else: - attn_padded = None - - return ids_padded, mask_padded, attn_padded + ) -> tuple[torch.Tensor, torch.Tensor]: + seq_len = input_ids.shape[0] + if seq_len == bucket: + return input_ids, audio_mask + return ( + F.pad(input_ids, (0, 0, 0, bucket - seq_len), value=0), + F.pad(audio_mask, (0, bucket - seq_len), value=False), + ) def _capture_for_key( self, key: tuple[int, int], input_ids: torch.Tensor, audio_mask: torch.Tensor, - attention_mask: torch.Tensor | None, + cu_seqs: torch.Tensor, ) -> dict: _, bucket = key device = input_ids.device - static_rope_table = self._gen._rope_table_for(bucket, device, self._gen.model_dtype) - static_input_ids = input_ids.clone() static_audio_mask = audio_mask.clone() - static_attn_mask = attention_mask.clone() if attention_mask is not None else None + static_cu_seqs = cu_seqs.clone() + static_rope_table = self._gen._rope_table_for(bucket, device, self._gen.model_dtype) with torch.no_grad(): _ = self._gen._step_forward( static_input_ids, static_audio_mask, - static_attn_mask, + static_cu_seqs, static_rope_table, ) torch.accelerator.synchronize(device) @@ -580,7 +630,7 @@ def _capture_for_key( static_output = self._gen._step_forward( static_input_ids, static_audio_mask, - static_attn_mask, + static_cu_seqs, static_rope_table, ) @@ -588,87 +638,83 @@ def _capture_for_key( "graph": graph, "static_input_ids": static_input_ids, "static_audio_mask": static_audio_mask, - "static_attn_mask": static_attn_mask, + "static_cu_seqs": static_cu_seqs, "static_rope_table": static_rope_table, "static_output": static_output, } logger.info("OmniVoice CUDA Graph captured for key %s", key) return entry + def make_capture_cu_seq(self, batch_size: int, bucket_size: int, device: torch.device) -> torch.Tensor: + num_real_sequences = 2 * batch_size + base, rem = divmod(bucket_size, num_real_sequences) + lengths = torch.full((num_real_sequences,), base, dtype=torch.int32, device=device) + lengths[:rem] += 1 + cu_seq = torch.empty(num_real_sequences + 2, dtype=torch.int32, device=device) + cu_seq[0] = 0 + cu_seq[1:-1] = lengths.cumsum(0) + cu_seq[-1] = bucket_size + return cu_seq + def warmup(self, device: torch.device) -> None: - """Pre-capture graphs for all bucket sizes with B=1 (two_b=2 for CFG).""" + """Pre-capture common request batch sizes for every token bucket.""" if not torch.cuda.is_available(): return logger.info( - "OmniVoice CUDA Graph warmup: capturing %d bucket sizes %s", - len(self._capture_sizes), - self._capture_sizes, + "OmniVoice CUDA Graph warmup: batch-aware capture plan %s", + self.capture_bucket_sizes_by_batch, ) - two_b = 2 num_cb = self._gen.config.num_audio_codebook - for bucket in self._capture_sizes: - key = (two_b, bucket) - dummy_ids = torch.zeros(two_b, num_cb, bucket, dtype=torch.long, device=device) - dummy_mask = torch.zeros(two_b, bucket, dtype=torch.bool, device=device) - # Capture with a float mask to match what forward() feeds at replay time, - # in the model dtype so replay can copy_ into it without a cast. - dummy_attn = torch.zeros(two_b, 1, bucket, bucket, dtype=self._gen.model_dtype, device=device) - self._graphs[key] = self._capture_for_key(key, dummy_ids, dummy_mask, dummy_attn) + for batch_size in self.capture_batch_sizes: + for bucket in self.capture_bucket_sizes_by_batch[batch_size]: + key = (batch_size, bucket) + dummy_ids = torch.zeros(bucket, num_cb, dtype=torch.long, device=device) + dummy_mask = torch.zeros(bucket, dtype=torch.bool, device=device) + dummy_cu_seqs = self.make_capture_cu_seq(batch_size, bucket, device) + self._graphs[key] = self._capture_for_key(key, dummy_ids, dummy_mask, dummy_cu_seqs) logger.info("OmniVoice CUDA Graph warmup complete (%d graphs)", len(self._graphs)) def __call__( self, input_ids: torch.Tensor, audio_mask: torch.Tensor, - attention_mask: torch.Tensor | None, + cu_seqs: torch.Tensor, + batch_size: int, ) -> torch.Tensor: if torch.cuda.is_current_stream_capturing(): - rope_table = self._gen._rope_table_for(input_ids.shape[-1], input_ids.device, self._gen.model_dtype) - return self._gen._step_forward(input_ids, audio_mask, attention_mask, rope_table) - - seq_len = input_ids.shape[-1] - two_b = input_ids.shape[0] - bucket = self._find_bucket(seq_len) if two_b == 2 else None - - # Graphs are captured with (and their static buffers hold) the additive - # float mask, so normalize here, before padding or any copy_ into them. - if attention_mask is not None: - attention_mask = _additive_float_mask(attention_mask, self._gen.model_dtype) + rope_table = self._gen._rope_table_for(input_ids.shape[0], input_ids.device, self._gen.model_dtype) + return self._gen._step_forward(input_ids, audio_mask, cu_seqs, rope_table) + seq_len = input_ids.shape[0] + bucket = self._find_bucket(batch_size, seq_len) + is_lazy = bucket is None if bucket is None: - # Lazy capture: oversized sequence or non-unit batch (no pre-warmed bucket). - # Lock prevents concurrent threads from double-capturing the same key. - # _lazy_graphs is capped at _MAX_LAZY_GRAPHS with LRU eviction to - # prevent unbounded GPU memory growth when seq_len varies widely. - key = (two_b, seq_len) - ids_in, mask_in, attn_in = input_ids, audio_mask, attention_mask - with self._lock: - entry = self._lazy_graphs.get(key) - if entry is None: - entry = self._capture_for_key(key, ids_in, mask_in, attn_in) - if len(self._lazy_graphs) >= self._MAX_LAZY_GRAPHS: - evicted_key, _ = self._lazy_graphs.popitem(last=False) - logger.warning("OmniVoice CUDA Graph lazy cache full; evicted key %s", evicted_key) - self._lazy_graphs[key] = entry - else: - key = (two_b, bucket) - ids_in, mask_in, attn_in = self._pad_inputs(input_ids, audio_mask, attention_mask, bucket) - with self._lock: - entry = self._graphs.get(key) - if entry is None: - entry = self._capture_for_key(key, ids_in, mask_in, attn_in) - self._graphs[key] = entry + bucket = round_up(seq_len, self._LAZY_CAPTURE_ALIGNMENT) + ids_in, mask_in = self._pad_inputs(input_ids, audio_mask, bucket) + runtime_cu_seqs = cu_seqs.clone() + runtime_cu_seqs[-1] = bucket + key = (batch_size, bucket) + cache = self._lazy_graphs if is_lazy else self._graphs + with self._lock: + entry = cache.get(key) + if entry is None: + if is_lazy and len(self._lazy_graphs) >= self._MAX_LAZY_GRAPHS: + evicted_key, _ = self._lazy_graphs.popitem(last=False) + logger.info("Evicted OmniVoice lazy CUDA Graph key %s", evicted_key) + entry = self._capture_for_key(key, ids_in, mask_in, runtime_cu_seqs) + cache[key] = entry + elif is_lazy: + self._lazy_graphs.move_to_end(key) entry["static_input_ids"].copy_(ids_in) entry["static_audio_mask"].copy_(mask_in) - if attn_in is not None and entry["static_attn_mask"] is not None: - entry["static_attn_mask"].copy_(attn_in) + entry["static_cu_seqs"].copy_(runtime_cu_seqs) entry["graph"].replay() output = entry["static_output"] - if bucket is not None and bucket != seq_len: - output = output[:, :, :seq_len, :] + if bucket != seq_len: + output = output[:, :seq_len, :] return output def clear(self) -> None: @@ -692,17 +738,17 @@ class OmniVoiceGenerator(nn.Module): - 32-step iterative unmasking with classifier-free guidance Optimizations: - - Full bidirectional attention via PyTorch SDPA (no auto-selected - FlashAttn/SageAttn/DiffusionAttention backend) + - Packed full-bidirectional varlen attention with SDPA fallback - regionally_compile() compatible for torch.compile on repeated blocks """ # For regionally_compile() support _repeated_blocks = ["layers"] - def __init__(self, config: OmniVoiceConfig): + def __init__(self, config: OmniVoiceConfig, od_config: OmniDiffusionConfig): super().__init__() self.config = config + self.od_config = od_config # Opt-in TF32; must run before any CUDA-graph capture so captured kernels honour it. if getattr(config, "enable_tf32", False): @@ -722,7 +768,10 @@ def __init__(self, config: OmniVoiceConfig): ) # Transformer layers - self.layers = nn.ModuleList([OmniVoiceTransformerBlock(config) for _ in range(config.llm_num_hidden_layers)]) + self.layers = nn.ModuleList( + [OmniVoiceTransformerBlock(config, layer_idx) for layer_idx in range(config.llm_num_hidden_layers)] + ) + self._needs_sdpa_mask = any(not layer.self_attn.use_packed_varlen for layer in self.layers) self.norm = OmniVoiceRMSNorm(config.llm_hidden_size, eps=config.llm_rms_norm_eps) # Prediction head: hidden → 8 * 1025 @@ -775,23 +824,23 @@ def _prepare_embeddings( """Prepare mixed text+audio embeddings. Args: - input_ids: [B, 8, S] - text tokens replicated across codebooks, + input_ids: [T, 8] - text tokens replicated across codebooks, audio positions have per-codebook token IDs - audio_mask: [B, S] - True for audio positions, False for text - text_embeds: optional cached [B, S, H] text-position embeddings - audio_mask_3d: optional cached [B, S, 1] audio_mask.unsqueeze(-1) + audio_mask: [T] - True for audio positions, False for text + text_embeds: optional cached [T, H] text-position embeddings + audio_mask_3d: optional cached [T, 1] audio_mask.unsqueeze(-1) Returns: - embeddings: [B, S, hidden_size] + embeddings: [T, hidden_size] """ # Cached across the denoising loop since text ids don't change. if text_embeds is None: - text_embeds = self.text_embedding(input_ids[:, 0, :]) + text_embeds = self.text_embedding(input_ids[:, 0]) if audio_mask_3d is None: audio_mask_3d = audio_mask.unsqueeze(-1) # Audio embeddings: offset per codebook, then sum across codebooks - shifted_ids = (input_ids * audio_mask.unsqueeze(1)) + self.codebook_layer_offsets.view(1, -1, 1) + shifted_ids = (input_ids * audio_mask.unsqueeze(1)) + self.codebook_layer_offsets.view(1, -1) audio_embeds = self.audio_embeddings(shifted_ids).sum(dim=1) # Merge: audio where audio_mask=True, text elsewhere @@ -800,35 +849,39 @@ def _prepare_embeddings( def _transformer_forward( self, inputs_embeds: torch.Tensor, - attention_mask: torch.Tensor | None = None, + cu_seqs: torch.Tensor, + max_seqlen: int | None = None, rope_table: torch.Tensor | None = None, ) -> torch.Tensor: """Run through transformer layers. Args: - inputs_embeds: [B, S, hidden_size] - attention_mask: [B, 1, S, S] or None - rope_table: optional precomputed [B * S, head_dim] RoPE table + inputs_embeds: [T, hidden_size] + cu_seqs: cumulative boundaries for packed sequences + rope_table: optional base [T, head_dim] RoPE table Returns: - hidden_states: [B, S, hidden_size] + hidden_states: [T, hidden_size] """ hidden_states = inputs_embeds if rope_table is None: - rope_table = self._rope_table_for(inputs_embeds.shape[1], inputs_embeds.device, hidden_states.dtype) - - # Safety: convert bool mask if caller hasn't (e.g. external paths beyond forward()). - if attention_mask is not None and attention_mask.dtype == torch.bool: - attention_mask = torch.zeros_like(attention_mask, dtype=hidden_states.dtype).masked_fill_( - ~attention_mask, float("-inf") - ) + rope_table = self._rope_table_for(inputs_embeds.shape[0], inputs_embeds.device, hidden_states.dtype) + seq_len = inputs_embeds.shape[0] + position_ids = _position_ids_from_cu_seqs(cu_seqs, seq_len) + packed_rope_table = rope_table.index_select(0, position_ids).contiguous() + attn_metadata = _attention_metadata_from_cu_seqs( + cu_seqs, + seq_len, + needs_sdpa_mask=self._needs_sdpa_mask, + max_seqlen=max_seqlen, + ) residual = None for layer in self.layers: hidden_states, residual = layer( hidden_states, - attention_mask=attention_mask, - rope_table=rope_table, + attn_metadata=attn_metadata, + rope_table=packed_rope_table, residual=residual, ) @@ -838,44 +891,77 @@ def _get_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: """Project hidden states to per-codebook logits. Args: - hidden_states: [B, S, hidden_size] + hidden_states: [T, hidden_size] Returns: - logits: [B, 8, S, 1025] + logits: [8, T, 1025] """ - batch_size, seq_len, _ = hidden_states.shape - logits_flat = self.audio_heads(hidden_states) # [B, S, 8*1025] + seq_len, _ = hidden_states.shape + logits_flat = self.audio_heads(hidden_states) # [T, 8*1025] return logits_flat.view( - batch_size, seq_len, self.config.num_audio_codebook, self.config.audio_vocab_size, - ).permute(0, 2, 1, 3) # [B, 8, S, 1025] + ).permute(1, 0, 2) # [8, T, 1025] def _step_forward( self, input_ids: torch.Tensor, audio_mask: torch.Tensor, - attention_mask: torch.Tensor | None, + cu_seqs: torch.Tensor, rope_table: torch.Tensor, ) -> torch.Tensor: """Single unmasking-step forward using a pre-cast RoPE table (CUDA graph safe).""" hidden_states = self._prepare_embeddings(input_ids, audio_mask) - residual = None - for layer in self.layers: - hidden_states, residual = layer( - hidden_states, attention_mask=attention_mask, rope_table=rope_table, residual=residual + hidden_states = self._transformer_forward(hidden_states, cu_seqs, rope_table=rope_table) + return self._get_logits(hidden_states) + + def _unmask_one_request( + self, + c_logits: torch.Tensor, + u_logits: torch.Tensor, + sample_tokens: torch.Tensor, + *, + num_to_unmask: int | torch.Tensor, + guidance_scale: float, + generator: torch.Generator, + class_temperature: float, + position_temperature: float, + layer_penalty_factor: float, + layer_ids: torch.Tensor, + ) -> None: + """Sample and unmask one request in place from FP32 target logits.""" + mask_id = self.config.audio_mask_id + if guidance_scale != 0: + log_probs = F.log_softmax( + (1.0 + guidance_scale) * c_logits - guidance_scale * u_logits, + dim=-1, ) - return self._get_logits(self.norm(hidden_states + residual)) + else: + log_probs = F.log_softmax(c_logits, dim=-1) + log_probs[..., mask_id] = -float("inf") + if class_temperature > 0.0: + pred_tokens = _gumbel_sample(log_probs, class_temperature, generator).argmax(dim=-1) + else: + pred_tokens = log_probs.argmax(dim=-1) + scores = log_probs.max(dim=-1)[0] + scores = scores - (layer_ids * layer_penalty_factor) + if position_temperature > 0.0: + scores = _gumbel_sample(scores, position_temperature, generator) + scores.masked_fill_(sample_tokens != mask_id, -float("inf")) + _, topk_idx = torch.topk(scores.flatten(), num_to_unmask) + flat_tokens = sample_tokens.flatten() + flat_tokens[topk_idx] = pred_tokens.flatten()[topk_idx] + sample_tokens.copy_(flat_tokens.view_as(sample_tokens)) @torch.inference_mode() def forward( self, input_ids: torch.Tensor, audio_mask: torch.Tensor, - attention_mask: torch.Tensor, + cond_lens: list[int], target_lens: list[int], - seed: int | None = None, + seed: int | list[int | None] | None = None, num_step: int = 32, guidance_scale: float = 2.0, t_shift: float = 0.1, @@ -886,36 +972,51 @@ def forward( """Run the full 32-step iterative unmasking generation. Args: - input_ids: [2*B, 8, S] - conditional (0:B) + unconditional (B:2B) - audio_mask: [2*B, S] - True for audio positions - attention_mask: [2*B, 1, S, S] - attention mask - target_lens: List of target audio lengths per batch item - num_step: Number of unmasking steps - guidance_scale: CFG scale - t_shift: Time shift for schedule - layer_penalty_factor: Penalty for later codebooks - position_temperature: Gumbel temperature for position selection - class_temperature: Temperature for token prediction (0=greedy) + input_ids: Packed token IDs with shape ``[total_seq_len, 8]`` in + request-major ``[cond0, uncond0, ...]`` order. + audio_mask: Boolean audio-position mask with shape + ``[total_seq_len]``. + cond_lens: Conditional sequence length for each request. + target_lens: Target length for each request; also the corresponding + unconditional sequence length. + seed: One seed per request, a shared scalar seed, or ``None``. + num_step: Number of iterative unmasking steps. + guidance_scale: Classifier-free guidance scale. + t_shift: Time shift used to construct the unmasking schedule. + layer_penalty_factor: Penalty applied to later codebooks. + position_temperature: Gumbel temperature for position selection. + class_temperature: Token sampling temperature; zero selects greedy + decoding. Returns: - tokens: [B, 8, max_target_len] - generated audio tokens + Packed generated audio tokens with shape + ``[1, 8, sum(target_lens)]``. """ B = len(target_lens) device = input_ids.device - max_target_len = max(target_lens) + total_target_lens = sum(target_lens) mask_id = self.config.audio_mask_id num_codebooks = self.config.num_audio_codebook - if seed is None: - seed = random.randint(0, 2**63 - 1) - generator = torch.Generator(device=device).manual_seed(seed) + seeds = seed if isinstance(seed, list) else [seed] * B + generators = [ + torch.Generator(device=device).manual_seed( + request_seed if request_seed is not None else random.randint(0, 2**63 - 1) + ) + for request_seed in seeds + ] # Initialize all target tokens as [MASK] - tokens = torch.full( - (B, num_codebooks, max_target_len), - mask_id, - dtype=torch.long, - device=device, - ) + tokens = torch.full((1, num_codebooks, total_target_lens), mask_id, dtype=torch.long, device=device) + target_offsets: list[int] = [] + target_offset = 0 + sequence_offsets: list[int] = [] + sequence_offset = 0 + for cond_len, target_len in zip(cond_lens, target_lens): + target_offsets.append(target_offset) + target_offset += target_len + sequence_offsets.append(sequence_offset) + sequence_offset += cond_len + target_len + cu_seqs = _build_cu_seqs(cond_lens, target_lens, device) # Compute unmasking schedule timesteps = _get_time_steps(0.0, 1.0, num_step + 1, t_shift).tolist() @@ -939,45 +1040,46 @@ def forward( layer_ids = torch.arange(num_codebooks, device=device).view(1, -1, 1) - # Single D2H pull for all conditional lengths instead of B per-item .item() syncs. - c_lens = attention_mask[:B, 0, 0].sum(dim=-1).tolist() - - # Materialize the SDPA float mask once so the captured graph (and eager path) skip per-layer conversion. - sdpa_attn_mask = _additive_float_mask(attention_mask, self.model_dtype) - use_cuda_graph = self._cuda_graph_fwd is not None and input_ids.is_cuda if not use_cuda_graph: # Eager-path-only constants (the cuda-graph captures its own). - text_embeds_cached = self.text_embedding(input_ids[:, 0, :]) + text_embeds_cached = self.text_embedding(input_ids[:, 0]) audio_mask_3d = audio_mask.unsqueeze(-1) - rope_table = self._rope_table_for(input_ids.shape[-1], device, text_embeds_cached.dtype) + rope_table = self._rope_table_for(input_ids.shape[0], device, text_embeds_cached.dtype) # Main iterative loop for step in range(num_step): if use_cuda_graph: # Float mask skips per-layer conversion; fp32 cast deferred to the per-item slices below. - batch_logits = self._cuda_graph_fwd(input_ids, audio_mask, sdpa_attn_mask) + batch_logits = self._cuda_graph_fwd(input_ids, audio_mask, cu_seqs, B) else: # Eager fallback reuses hoisted constants (text embeds, sdpa mask, rope table). inputs_embeds = self._prepare_embeddings( input_ids, audio_mask, text_embeds=text_embeds_cached, audio_mask_3d=audio_mask_3d ) - hidden_states = self._transformer_forward(inputs_embeds, sdpa_attn_mask, rope_table=rope_table) + hidden_states = self._transformer_forward( + inputs_embeds, + cu_seqs, + max_seqlen=max(cond_lens), + rope_table=rope_table, + ) # fp32 cast deferred to the per-item slices below. batch_logits = self._get_logits(hidden_states) - # batch_logits: [2*B, 8, S, 1025] + # batch_logits: [8, T, 1025] for i in range(B): k = schedules[i][step] if k <= 0: continue - c_len = c_lens[i] + c_len = cond_lens[i] t_len = target_lens[i] + request_start = sequence_offsets[i] + cond_end = request_start + c_len # Extract logits for target region; upcast only the slices we actually consume. - c_logits = batch_logits[i : i + 1, :, c_len - t_len : c_len, :].to(torch.float32) - u_logits = batch_logits[B + i : B + i + 1, :, :t_len, :].to(torch.float32) + c_logits = batch_logits[:, cond_end - t_len : cond_end, :].unsqueeze(0).to(torch.float32) + u_logits = batch_logits[:, cond_end : cond_end + t_len, :].unsqueeze(0).to(torch.float32) # Classifier-free guidance. Fuse the chain: the two inner # log_softmax normalizers are per-position scalars that the final @@ -996,7 +1098,7 @@ def forward( # Token prediction if class_temperature > 0.0: - pred_tokens = _gumbel_sample(log_probs, class_temperature, generator).argmax(dim=-1) + pred_tokens = _gumbel_sample(log_probs, class_temperature, generators[i]).argmax(dim=-1) else: pred_tokens = log_probs.argmax(dim=-1) # [1, 8, T] @@ -1008,10 +1110,11 @@ def forward( # Gumbel noise for position selection if position_temperature > 0.0: - scores = _gumbel_sample(scores, position_temperature, generator) + scores = _gumbel_sample(scores, position_temperature, generators[i]) # Mask out already unmasked positions - sample_tokens = tokens[i : i + 1, :, :t_len] + target_start = target_offsets[i] + sample_tokens = tokens[:, :, target_start : target_start + t_len] scores.masked_fill_(sample_tokens != mask_id, -float("inf")) # Select top-k positions to unmask. .flatten() on this non-contiguous view already copies. @@ -1021,8 +1124,9 @@ def forward( sample_tokens.copy_(flat_tokens.view_as(sample_tokens)) # Mirror update into both cond and uncond input_ids halves for the next step. - input_ids[i, :, c_len - t_len : c_len] = sample_tokens.squeeze(0) - input_ids[B + i, :, :t_len] = sample_tokens.squeeze(0) + packed_sample_tokens = sample_tokens.squeeze(0).transpose(0, 1) + input_ids[cond_end - t_len : cond_end] = packed_sample_tokens + input_ids[cond_end : cond_end + t_len] = packed_sample_tokens return tokens From e74fe3c6e86d664cf165d22252163f20df587065 Mon Sep 17 00:00:00 2001 From: NancyFyong <88076188+NancyFyong@users.noreply.github.com> Date: Wed, 9 Sep 2026 10:00:25 +0800 Subject: [PATCH 13/93] [Model] Add Audio8 TTS Preview 0.6B (DualAR, 44.1 kHz codec) (#6157) Signed-off-by: NancyFyong Co-authored-by: NancyFyong --- benchmarks/tts/README.md | 19 + benchmarks/tts/model_configs.yaml | 11 + docs/models/supported_models.md | 1 + .../text_to_speech/README.md | 52 ++ .../text_to_speech/audio8_tts/end2end.py | 237 ++++++ .../online_serving/text_to_speech/README.md | 133 +++ .../text_to_speech/audio8_tts/gradio_demo.py | 258 ++++++ .../audio8_tts/run_gradio_demo.sh | 44 + .../text_to_speech/audio8_tts/run_server.sh | 24 + .../audio8_tts/speech_client.py | 131 +++ .../e2e/offline_inference/test_audio8_tts.py | 112 +++ tests/e2e/online_serving/test_audio8_tts.py | 166 ++++ .../openai/test_audio8_tts_adapter.py | 85 ++ .../openai_api/test_tts_detection.py | 2 + tests/helpers/assertions.py | 24 +- tests/helpers/client.py | 118 +++ .../models/audio8_tts/__init__.py | 0 .../test_audio8_tts_codec_decoder.py | 141 ++++ .../test_audio8_tts_prompt_utils.py | 112 +++ .../audio8_tts/test_audio8_tts_sampling.py | 156 ++++ .../test_audio8_tts_async_chunk.py | 163 ++++ tests/worker/test_gpu_ar_model_runner.py | 1 + vllm_omni/config/pipeline_registry.py | 2 + vllm_omni/deploy/audio8_tts.yaml | 80 ++ .../entrypoints/openai/serving_speech.py | 2 + .../openai/tts_adapters/__init__.py | 1 + .../openai/tts_adapters/audio8_tts.py | 190 +++++ .../models/audio8_tts/__init__.py | 3 + .../audio8_tts/audio8_tts_codec_decoder.py | 284 +++++++ .../models/audio8_tts/audio8_tts_fast_ar.py | 346 ++++++++ .../models/audio8_tts/audio8_tts_slow_ar.py | 755 ++++++++++++++++++ .../model_executor/models/audio8_tts/codec.py | 600 ++++++++++++++ .../models/audio8_tts/codec_utils.py | 224 ++++++ .../audio8_tts/configuration_audio8_tts.py | 267 +++++++ .../models/audio8_tts/pipeline.py | 50 ++ .../models/audio8_tts/prompt_utils.py | 159 ++++ .../models/audio8_tts/sampling.py | 227 ++++++ vllm_omni/model_executor/models/registry.py | 11 + .../stage_input_processors/audio8_tts.py | 173 ++++ .../transformers_utils/configs/__init__.py | 9 +- .../transformers_utils/configs/audio8_tts.py | 30 + 41 files changed, 5398 insertions(+), 5 deletions(-) create mode 100644 examples/offline_inference/text_to_speech/audio8_tts/end2end.py create mode 100644 examples/online_serving/text_to_speech/audio8_tts/gradio_demo.py create mode 100644 examples/online_serving/text_to_speech/audio8_tts/run_gradio_demo.sh create mode 100644 examples/online_serving/text_to_speech/audio8_tts/run_server.sh create mode 100644 examples/online_serving/text_to_speech/audio8_tts/speech_client.py create mode 100644 tests/e2e/offline_inference/test_audio8_tts.py create mode 100644 tests/e2e/online_serving/test_audio8_tts.py create mode 100644 tests/entrypoints/openai/test_audio8_tts_adapter.py create mode 100644 tests/model_executor/models/audio8_tts/__init__.py create mode 100644 tests/model_executor/models/audio8_tts/test_audio8_tts_codec_decoder.py create mode 100644 tests/model_executor/models/audio8_tts/test_audio8_tts_prompt_utils.py create mode 100644 tests/model_executor/models/audio8_tts/test_audio8_tts_sampling.py create mode 100644 tests/model_executor/stage_input_processors/test_audio8_tts_async_chunk.py create mode 100644 vllm_omni/deploy/audio8_tts.yaml create mode 100644 vllm_omni/entrypoints/openai/tts_adapters/audio8_tts.py create mode 100644 vllm_omni/model_executor/models/audio8_tts/__init__.py create mode 100644 vllm_omni/model_executor/models/audio8_tts/audio8_tts_codec_decoder.py create mode 100644 vllm_omni/model_executor/models/audio8_tts/audio8_tts_fast_ar.py create mode 100644 vllm_omni/model_executor/models/audio8_tts/audio8_tts_slow_ar.py create mode 100644 vllm_omni/model_executor/models/audio8_tts/codec.py create mode 100644 vllm_omni/model_executor/models/audio8_tts/codec_utils.py create mode 100644 vllm_omni/model_executor/models/audio8_tts/configuration_audio8_tts.py create mode 100644 vllm_omni/model_executor/models/audio8_tts/pipeline.py create mode 100644 vllm_omni/model_executor/models/audio8_tts/prompt_utils.py create mode 100644 vllm_omni/model_executor/models/audio8_tts/sampling.py create mode 100644 vllm_omni/model_executor/stage_input_processors/audio8_tts.py create mode 100644 vllm_omni/transformers_utils/configs/audio8_tts.py diff --git a/benchmarks/tts/README.md b/benchmarks/tts/README.md index c403e5dd061..918e0727aec 100644 --- a/benchmarks/tts/README.md +++ b/benchmarks/tts/README.md @@ -85,6 +85,25 @@ vllm bench serve --omni \ --save-result --result-dir ./results ``` +#### Set the sample rate for models that are not 24 kHz + +Streamed PCM carries no header, so the harness derives `audio_duration` from the +byte count and a **default 24000 Hz** assumption +(`vllm_omni/metrics/definitions.py:DEFAULT_AUDIO_SAMPLE_RATE`). For a model that +emits another rate, export the real one or every duration-derived metric is +wrong by `real_sr / 24000`: + +```bash +# Audio8 TTS, Fish Speech S2 Pro, Ming-omni-tts, ... all decode at 44.1 kHz +export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=44100 +``` + +Without it, a 44.1 kHz model reports `audio_duration` 1.84x too long, so +`audio_rtf` looks 1.84x *better* than reality and `audio_underrun` / +`Streaming continuity OK rate` are measured against inflated chunk budgets. +A quick sanity check: compare `Mean AUDIO_DURATION (s)` against the actual +length of a saved WAV for the same prompt. + #### Streaming continuity (`audio_underrun`) `audio_underrun` (seconds) is the per-request worst-case buffer deficit diff --git a/benchmarks/tts/model_configs.yaml b/benchmarks/tts/model_configs.yaml index 69fe09b50ab..53c84e66a87 100644 --- a/benchmarks/tts/model_configs.yaml +++ b/benchmarks/tts/model_configs.yaml @@ -17,6 +17,17 @@ models: extra_params: lang: en text_normalization: true + # Audio8 TTS Preview: DualAR, 44.1 kHz. trust_remote_code stays false. + # Bench at 44.1 kHz: export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=44100, otherwise + # the streamed-PCM duration defaults to 24 kHz and RTF looks 1.84x too good. + Audio8/Audio8-TTS-Preview-0.6b: + supported_tasks: [default_voice, voice_clone] + backend: openai-audio-speech + endpoint: /v1/audio/speech + trust_remote_code: false + task_extra_body: + default_voice: {} + voice_clone: {} # -CustomVoice checkpoints lack speaker_encoder weights, so voice_clone is # NOT supported (an attempt raises ValueError from _extract_speaker_embedding diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index de5d5e34665..6b0952eac47 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -97,6 +97,7 @@ th { | `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | — | | `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | — | | `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | ✅︎ | — | +| `Audio8TTSSlowARForConditionalGeneration` | Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | ✅︎ | | | | — | | `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only), SenseNova-U1.5 | `SenseNova/SenseNova-U1-8B-MoT`, `sensenova/SenseNova-U1.5-8B-MoT` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SenseNova/SenseNova-U1.5.md) | | `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — | | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — | diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index da00ea44d88..a8a273026a5 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -14,6 +14,7 @@ list of supported architectures across all modalities, see | Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate | |---|---|---|---|---|---|---| +| Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | dual-AR | ✓ | ✓ | 11 languages | 44.1 kHz | | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz | | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz | | Gepard-1.0 | `nineninesix/gepard-1.0` | single (native AR + NanoCodec) | — (zero-shot; cloning WIP) | — (serving WIP) | zero-shot | 22.05 kHz | @@ -44,6 +45,57 @@ python examples/offline_inference/text_to_speech//end2end.py \ --- +## Audio8 TTS Preview + +0.6B DualAR text-to-speech model (Audio8) with its own 44.1 kHz neural audio +codec: a slow AR transformer predicts one semantic token per codec frame, and a +4-layer fast AR predicts that frame's remaining 9 codebooks. + +### Prerequisites +No extra packages: the codec architecture ships in tree and its weights +(`codec.pth`) come with the checkpoint. + +### Quick start +```bash +python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \ + --text "Welcome to Audio8 TTS, a compact multilingual text to speech model." +``` + +### Voice cloning +```bash +python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \ + --text "Welcome to Audio8 TTS." \ + --ref-audio /path/to/reference.wav \ + --ref-text "The exact transcript of the reference recording." +``` +The reference transcript must match what is actually spoken in the clip; +a mismatch degrades both stability and speaker similarity. + +### Streaming +```bash +python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \ + --text "Welcome to Audio8 TTS." --streaming +``` + +### Notes +- Output: 44.1 kHz mono WAV; the codec emits ~21.5 frames/s (2048 samples per frame). +- Context is 2048 packed text+audio positions, so `max_model_len` is 2048. +- Greedy decoding (`temperature: 0`) is degenerate for this checkpoint — it + ends the utterance almost immediately. Keep the deploy defaults + (`temperature 0.7`, `top_k 50`, `top_p 0.9`), which also match + `generation_config.json`. +- Repetition-Aware Sampling is applied to the semantic token: a repeat inside + the last 10 frames is re-drawn from a flatter distribution instead of being + masked out. +- `--num-prompts 4` smoke-tests concurrency (stage 0 runs `max_num_seqs: 4`). + +### Measured behaviour (1x H20, shared GPU) +`--streaming` on a 58-character English sentence: TTFA ~0.74 s (4-frame first +chunk), 7 chunks, 1.55 s wall for 3.16 s of audio (RTF ~0.49). Treat these as +smoke-test magnitudes, not a benchmark. + +--- + ## CosyVoice3 2-stage TTS pipeline (`talker` + `code2wav`) at 24 kHz. diff --git a/examples/offline_inference/text_to_speech/audio8_tts/end2end.py b/examples/offline_inference/text_to_speech/audio8_tts/end2end.py new file mode 100644 index 00000000000..8ce94f3851f --- /dev/null +++ b/examples/offline_inference/text_to_speech/audio8_tts/end2end.py @@ -0,0 +1,237 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Offline inference demo for Audio8 TTS Preview 0.6B via vLLM Omni. + +Usage: + # Text-only synthesis + python end2end.py --text "Welcome to Audio8 TTS." + + # Zero-shot voice cloning (the transcript must match the reference audio) + python end2end.py --text "Welcome to Audio8 TTS." \ + --ref-audio reference.wav --ref-text "The exact transcript of the reference recording." + + # Streaming (per-chunk TTFA / inter-chunk timings) + python end2end.py --text "Welcome to Audio8 TTS." --streaming +""" + +import asyncio +import os +import time + +import numpy as np +import soundfile as sf +import torch + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +from vllm_omni import AsyncOmni, Omni +from vllm_omni.model_executor.models.audio8_tts.codec_utils import estimate_reference_code_frames +from vllm_omni.model_executor.models.audio8_tts.prompt_utils import ( + build_text_only_prompt_ids, + estimate_voice_clone_prompt_len, + normalize_text, +) +from vllm_omni.utils.tracking_parser import TrackingArgumentParser + +DEFAULT_MODEL = "Audio8/Audio8-TTS-Preview-0.6b" + + +def build_prompt( + text: str, + ref_audio_path: str | None = None, + ref_text: str | None = None, + model_name: str = DEFAULT_MODEL, +) -> dict: + """Build an engine prompt, using the same protocol as online serving.""" + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(model_name) + if ref_audio_path is None and ref_text is None: + prompt_ids, normalized_text = build_text_only_prompt_ids(tokenizer, text) + return { + "prompt_token_ids": prompt_ids, + "additional_information": {"text": [normalized_text]}, + } + + if not ref_audio_path or not ref_text: + raise ValueError("Audio8 TTS voice cloning requires both --ref-audio and --ref-text") + + normalized_text = normalize_text(text) + normalized_ref_text = normalize_text(ref_text, add_default_speaker=True) + ref_audio_wav, ref_audio_sr = sf.read(ref_audio_path, dtype="float32", always_2d=False) + if ref_audio_wav.ndim > 1: + ref_audio_wav = ref_audio_wav.mean(axis=-1) + ref_frames = estimate_reference_code_frames(int(ref_audio_wav.shape[0]), int(ref_audio_sr)) + placeholder_len = estimate_voice_clone_prompt_len(tokenizer, normalized_text, normalized_ref_text, ref_frames) + + # The clone prompt embeds the reference audio's own codec codes, so the + # model builds it in preprocess(); we only reserve the exact length here. + return { + "prompt_token_ids": [1] * placeholder_len, + "additional_information": { + "text": normalized_text, + "ref_text": normalized_ref_text, + "ref_audio_wav": torch.from_numpy(np.asarray(ref_audio_wav, dtype=np.float32)), + "ref_audio_sr": int(ref_audio_sr), + "audio8_structured_voice_clone": True, + }, + } + + +def extract_audio(multimodal_output: dict) -> tuple[torch.Tensor, int]: + """Pull the waveform and sample rate out of a multimodal output. + + The output processor concatenates the per-step delta tensors under + ``model_outputs``; ``audio`` is the older alias. Never use ``a or b`` on + these values -- truthiness on a tensor raises. + """ + audio = multimodal_output.get("model_outputs") + if audio is None: + audio = multimodal_output.get("audio") + if audio is None: + raise ValueError(f"No audio in multimodal_output: {sorted(multimodal_output)}") + if isinstance(audio, list): + audio = torch.cat([torch.as_tensor(chunk).reshape(-1) for chunk in audio], dim=0) + audio = torch.as_tensor(audio).reshape(-1) + + sr_raw = multimodal_output.get("sr") + if isinstance(sr_raw, list): + sr_raw = sr_raw[-1] if sr_raw else None + if sr_raw is None: + raise ValueError("Missing sample rate in multimodal_output") + sample_rate = int(sr_raw.item()) if hasattr(sr_raw, "item") else int(sr_raw) + return audio, sample_rate + + +def _save_wav(output_dir: str, request_id: str, audio: torch.Tensor, sample_rate: int) -> None: + out_wav = os.path.join(output_dir, f"output_{request_id}.wav") + sf.write(out_wav, audio.float().cpu().numpy(), samplerate=sample_rate, format="WAV") + duration = audio.numel() / sample_rate + # print(), not logging: vLLM reconfigures logging on engine start and + # disables loggers created before that, which would swallow these lines. + print(f"Request {request_id}: saved {out_wav} (sr={sample_rate}, {duration:.2f}s)") + + +def _omni_kwargs(args, model_name: str) -> dict: + kwargs = { + "model": model_name, + "log_stats": args.log_stats, + "stage_init_timeout": args.stage_init_timeout, + } + if args.deploy_config: + kwargs["deploy_config"] = args.deploy_config + return kwargs + + +def main(args) -> None: + os.makedirs(args.output_dir, exist_ok=True) + model_name = args.model or DEFAULT_MODEL + prompts = [ + build_prompt(text, args.ref_audio, args.ref_text, model_name) for text in ([args.text] * args.num_prompts) + ] + + omni = Omni(**_omni_kwargs(args, model_name)) + t_start = time.perf_counter() + total_audio_seconds = 0.0 + for stage_outputs in omni.generate(prompts): + request_output = stage_outputs + if request_output is None or not request_output.outputs: + continue + audio, sample_rate = extract_audio(request_output.outputs[0].multimodal_output) + assert audio.numel() > 0, "Audio8 TTS produced an empty waveform" + total_audio_seconds += audio.numel() / sample_rate + _save_wav(args.output_dir, request_output.request_id, audio, sample_rate) + elapsed = time.perf_counter() - t_start + rtf = elapsed / max(total_audio_seconds, 1e-6) + print(f"Total {elapsed * 1000:.1f} ms for {total_audio_seconds:.2f} s of audio (RTF={rtf:.3f})") + + +async def main_streaming(args) -> None: + os.makedirs(args.output_dir, exist_ok=True) + model_name = args.model or DEFAULT_MODEL + prompt = build_prompt(args.text, args.ref_audio, args.ref_text, model_name) + + omni = AsyncOmni(**_omni_kwargs(args, model_name)) + request_id = "0" + chunks: list[torch.Tensor] = [] + sample_rate = None + t_start = time.perf_counter() + t_prev = t_start + chunk_index = 0 + + async for stage_output in omni.generate(prompt, request_id=request_id): + multimodal_output = stage_output.outputs[0].multimodal_output + if not multimodal_output: + continue + audio = multimodal_output.get("model_outputs") + if audio is None: + audio = multimodal_output.get("audio") + if audio is not None: + if isinstance(audio, list): + chunks.extend(torch.as_tensor(chunk).reshape(-1) for chunk in audio) + else: + chunks.append(torch.as_tensor(audio).reshape(-1)) + sr_raw = multimodal_output.get("sr") + if isinstance(sr_raw, list): + sr_raw = sr_raw[-1] if sr_raw else None + if sr_raw is not None: + sample_rate = int(sr_raw.item()) if hasattr(sr_raw, "item") else int(sr_raw) + + if stage_output.finished: + break + now = time.perf_counter() + if chunk_index == 0: + print(f"chunk 0: TTFA={(now - t_start) * 1000:.1f} ms") + else: + print(f"chunk {chunk_index}: inter_chunk={(now - t_prev) * 1000:.1f} ms") + t_prev = now + chunk_index += 1 + + if not chunks or sample_rate is None: + raise RuntimeError("Audio8 TTS streaming produced no audio") + _save_wav(args.output_dir, request_id, torch.cat(chunks, dim=0), sample_rate) + print(f"Streaming done: total={(time.perf_counter() - t_start) * 1000:.1f} ms chunks={chunk_index}") + + +def parse_args(): + parser = TrackingArgumentParser(description="Audio8 TTS Preview offline inference with vLLM Omni") + parser.add_argument("--model", type=str, default=DEFAULT_MODEL, help="Model path or HuggingFace repo ID.") + parser.add_argument( + "--text", + type=str, + default="Welcome to Audio8 TTS, a compact multilingual text to speech model.", + help="Text to synthesize.", + ) + parser.add_argument("--ref-audio", type=str, default=None, help="Reference audio path for voice cloning.") + parser.add_argument( + "--ref-text", + type=str, + default=None, + help="Transcript of --ref-audio; required for voice cloning and must match the recording.", + ) + parser.add_argument("--output-dir", type=str, default="output_audio", help="Directory for output WAV files.") + parser.add_argument( + "--num-prompts", + type=int, + default=1, + help="Duplicate the prompt N times to smoke-test concurrent requests.", + ) + parser.add_argument( + "--deploy-config", + type=str, + default=None, + help="Override the deploy config path; defaults to vllm_omni/deploy/audio8_tts.yaml via model_type.", + ) + parser.add_argument("--streaming", action="store_true", help="Use AsyncOmni and report per-chunk latency.") + parser.add_argument("--log-stats", action="store_true", help="Enable engine stats logging.") + parser.add_argument("--stage-init-timeout", type=int, default=600, help="Per-stage init timeout in seconds.") + return parser.parse_args() + + +if __name__ == "__main__": + parsed_args = parse_args() + if parsed_args.streaming: + asyncio.run(main_streaming(parsed_args)) + else: + main(parsed_args) diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index a4cf2fdb96e..b1221de6308 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -14,6 +14,7 @@ For the full list of supported architectures across all modalities, see | Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo | | --- | --- | --- | --- | --- | --- | +| Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | uploaded audio voice only; no presets | ✓ | | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ | | GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ | | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✓ (`ref_audio` or uploaded `voice`) | `stream=true` response, non-chunk | uploaded audio voice only; no presets | — | @@ -253,6 +254,138 @@ python examples/online_serving/text_to_speech/indextts2/speech_client.py \ --- +## Audio8 TTS Preview + +0.6B DualAR TTS at 44.1 kHz, 11 languages, zero-shot voice cloning. + +### Prerequisites +None beyond the base install: the neural audio codec is implemented in tree and +its weights (`codec.pth`) ship with the checkpoint. + +### Launch +```bash +vllm serve Audio8/Audio8-TTS-Preview-0.6b --omni --port 8092 +# or: +./audio8_tts/run_server.sh +``` +The deploy config auto-loads from `vllm_omni/deploy/audio8_tts.yaml` (HF +`model_type` is `arktts`). Do **not** pass `--trust-remote-code`: vllm-omni +registers its own `arktts` config, and transformers would otherwise prefer the +checkpoint's remote code and bypass it. + +### Text-only synthesis +```bash +curl -X POST http://localhost:8092/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Welcome to Audio8 TTS.", + "response_format": "wav" + }' --output output.wav +``` + +### Voice cloning +```bash +curl -X POST http://localhost:8092/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Welcome to Audio8 TTS.", + "ref_audio": "https://example.com/reference.wav", + "ref_text": "The exact transcript of the reference recording." + }' --output cloned.wav +``` +`ref_audio` also accepts a `data:audio/wav;base64,...` URL. `ref_text` is +mandatory whenever `ref_audio` is present and must match the recording. +Uploading a voice via `POST /v1/audio/voices` lets the server reuse the encoded +reference codes across requests (`voice: ""`). + +### Streaming +```bash +python audio8_tts/speech_client.py --text "Welcome to Audio8 TTS." --stream --output out.pcm +ffplay -f s16le -ar 44100 -ac 1 out.pcm +``` + +### Gradio demo +```bash +./audio8_tts/run_gradio_demo.sh # server + demo +python audio8_tts/gradio_demo.py --api-base http://localhost:8092 # demo only +``` + +### Benchmark + +```bash +export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=44100 # required: harness defaults to 24 kHz +vllm bench serve --omni --host 127.0.0.1 --port 8092 \ + --model Audio8/Audio8-TTS-Preview-0.6b \ + --backend openai-audio-speech --endpoint /v1/audio/speech \ + --dataset-name seed-tts-text --dataset-path benchmarks/build_dataset/seed_tts_smoke \ + --seed-tts-locale en --num-prompts 20 --num-warmups 2 \ + --max-concurrency 1 --request-rate inf \ + --percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun +``` + +Measured on 1x H20 **shared with an unrelated job holding ~80% SM**, so treat +these as a lower bound. 20 prompts from `seed_tts_smoke/en`, mean audio 4.5 s, +deploy defaults: + +| | HF reference (batch=1) | vllm-omni c=1 | c=4 | c=8 | +|---|---|---|---|---| +| RTF (mean) | 1.008 | **0.19** | 0.30 | 0.54 | +| Throughput (req/s) | 0.23 | 1.18 | 2.84 | **2.93** | +| E2E latency (mean, ms) | 4375 | 850 | 1349 | 2361 | +| Time to first packet (mean, ms) | n/a (no streaming) | 63 | 105 | 1124 | +| Underrun p99 / continuity OK | n/a | 0 s / 100% | 0 s / 100% | 0 s / 100% | + +~5.3x lower per-request RTF than the upstream reference implementation at +concurrency 1, ~12.8x aggregate throughput at concurrency 8. + +#### Stage-0 `max_num_seqs`: first-packet vs throughput + +This subsection is operator tuning guidance, not a claim this PR optimizes +anything: the shipped deploy default is unchanged (`max_num_seqs: 4`). The +numbers below are a single run on the same contended H20 as above (one unrelated +job holding ~80% SM, no repeats), so read them as directional, not as measured +speedups. Same setup, varying stage 0's `max_num_seqs` (stage 1 stays at 1): + +| stage-0 `max_num_seqs` | client concurrency | req/s | RTF | TTFP mean / p99 (ms) | underrun p99 (s) | +|---|---|---|---|---|---| +| 4 (default) | 1 | 1.18 | 0.19 | 63 / 67 | 0.00 | +| 4 | 4 | 2.84 | 0.30 | 105 / 148 | 0.00 | +| 4 | 8 | 2.93 | 0.54 | 1124 / 1642 | 0.00 | +| 8 | 1 | 1.16 | 0.19 | 64 / 70 | 0.00 | +| 8 | 4 | 2.88 | 0.30 | 107 / 148 | 0.00 | +| 8 | 8 | **3.72** | 0.40 | **247** / 733 | 0.22 | +| 8 | 16 | 3.91 | 0.65 | 1366 / 2530 | 0.29 | + +Reading of this: + +- Below the limit (`concurrency <= max_num_seqs`) the setting is irrelevant, as + expected — c=1 and c=4 are identical for both. +- At c=8, raising the limit to 8 admits every request instead of queueing half + of them: in this single run, roughly +27% throughput and ~4.5x lower + first-packet latency. +- Throughput saturates around 3.7-3.9 req/s (~17 s of audio per second). Past + c=8 you only buy queueing delay. +- The cost of admitting more streams is a 0.22-0.29 s worst-case buffer deficit + in 1 of 20 requests (`Streaming continuity OK rate` 95%), because stage 1 runs + `max_num_seqs: 1` and the codec window round-robins across streams. Client-side + prebuffering (the bundled Gradio demo holds 2 chunks) absorbs that; raising + stage 1's `max_num_seqs` is the untested knob if you need both. + +The shipped default stays at 4 because every measured point there had zero +underrun. Raise it to 8 if first-packet latency under load matters more than a +rare sub-300 ms gap. + +### Notes +- Output: 44.1 kHz mono; ~21.5 codec frames per second. +- No built-in speaker presets. Omit `voice` for a random timbre, or clone one. +- `max_new_tokens` caps generated codec frames (1 frame ~= 46 ms). +- Context is 2048 packed text+audio positions (`max_model_len: 2048`). +- Do not send `temperature: 0`: greedy decoding is degenerate for this + checkpoint and truncates the utterance. The deploy defaults + (0.7 / top_k 50 / top_p 0.9) mirror `generation_config.json`. + +--- + ## Fish Speech S2 Pro 4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec, which is vendored in diff --git a/examples/online_serving/text_to_speech/audio8_tts/gradio_demo.py b/examples/online_serving/text_to_speech/audio8_tts/gradio_demo.py new file mode 100644 index 00000000000..368ddfaa059 --- /dev/null +++ b/examples/online_serving/text_to_speech/audio8_tts/gradio_demo.py @@ -0,0 +1,258 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Gradio demo for Audio8 TTS Preview via the /v1/audio/speech API. + +Supports: + - text-to-speech synthesis in 11 languages + - zero-shot voice cloning from reference audio (upload, microphone, or URL) + - streaming (progressive PCM chunks) and non-streaming modes + +Usage: + # Start the server first (see run_server.sh), then: + python gradio_demo.py --api-base http://localhost:8092 +""" + +import argparse +import base64 +import io + +try: + import gradio as gr +except ImportError: + raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None +import httpx +import numpy as np +import soundfile as sf + +#: Audio8 TTS decodes at 44.1 kHz; raw PCM must be interpreted at that rate. +PCM_SAMPLE_RATE = 44100 +DEFAULT_API_BASE = "http://localhost:8092" + + +def encode_audio_to_base64(audio_data: tuple) -> str: + """Encode a Gradio ``(sample_rate, numpy_array)`` input as a data URL.""" + sample_rate, audio_np = audio_data + if audio_np.dtype != np.int16: + if audio_np.dtype in (np.float32, np.float64): + audio_np = (np.clip(audio_np, -1.0, 1.0) * 32767).astype(np.int16) + else: + audio_np = audio_np.astype(np.int16) + buf = io.BytesIO() + sf.write(buf, audio_np, sample_rate, format="WAV") + return f"data:audio/wav;base64,{base64.b64encode(buf.getvalue()).decode('utf-8')}" + + +def build_payload( + text: str, + ref_audio: tuple | None, + ref_audio_url: str, + ref_text: str, + response_format: str = "wav", + stream: bool = False, +) -> dict: + """Build the ``/v1/audio/speech`` payload. + + Audio8 TTS has no built-in speakers: omit ``voice`` for a random timbre, or + supply ``ref_audio`` + ``ref_text`` to clone a voice. + """ + if not text or not text.strip(): + raise gr.Error("Please enter text to synthesize.") + + payload: dict = { + "input": text.strip(), + "response_format": "pcm" if stream else response_format, + "stream": stream, + } + if stream: + payload["stream_format"] = "audio" + + ref_url = ref_audio_url.strip() if ref_audio_url else "" + if ref_audio is not None: + payload["ref_audio"] = encode_audio_to_base64(ref_audio) + elif ref_url: + payload["ref_audio"] = ref_url + + if "ref_audio" in payload: + if not ref_text or not ref_text.strip(): + raise gr.Error("Voice cloning requires the transcript of the reference audio.") + payload["ref_text"] = ref_text.strip() + + return payload + + +def generate_speech(api_base: str, text: str, ref_audio, ref_audio_url, ref_text, response_format): + """Non-streaming request; returns the whole waveform.""" + payload = build_payload(text, ref_audio, ref_audio_url, ref_text, response_format, stream=False) + try: + with httpx.Client(timeout=300.0) as client: + resp = client.post( + f"{api_base}/v1/audio/speech", + json=payload, + headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"}, + ) + except httpx.TimeoutException: + raise gr.Error("Request timed out. The server may be busy.") from None + except httpx.ConnectError: + raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?") from None + + if resp.status_code != 200: + raise gr.Error(f"Server error ({resp.status_code}): {resp.text}") + if "application/json" in resp.headers.get("content-type", ""): + raise gr.Error(f"Server error: {resp.text}") + + try: + if response_format == "pcm": + audio_np = np.frombuffer(resp.content, dtype=np.int16).astype(np.float32) / 32767.0 + return (PCM_SAMPLE_RATE, audio_np) + audio_np, sample_rate = sf.read(io.BytesIO(resp.content)) + if audio_np.ndim > 1: + audio_np = audio_np[:, 0] + return (sample_rate, audio_np.astype(np.float32)) + except Exception as e: + raise gr.Error(f"Failed to decode audio: {e}") from None + + +def generate_speech_stream(api_base, text, ref_audio, ref_audio_url, ref_text, response_format, chunk_seconds=0.5): + """Streaming request; yields PCM chunks for progressive playback.""" + payload = build_payload(text, ref_audio, ref_audio_url, ref_text, response_format, stream=True) + min_chunk_samples = int(PCM_SAMPLE_RATE * chunk_seconds) + prebuffer_chunks = 2 + pending: list[np.ndarray] = [] + pending_len = 0 + chunks_yielded = 0 + prebuffer: list[np.ndarray] = [] + leftover = b"" + + try: + with ( + httpx.Client(timeout=300.0) as client, + client.stream( + "POST", + f"{api_base}/v1/audio/speech", + json=payload, + headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"}, + ) as resp, + ): + if resp.status_code != 200: + resp.read() + raise gr.Error(f"Server error ({resp.status_code}): {resp.text}") + for chunk in resp.iter_bytes(): + if not chunk: + continue + raw = leftover + chunk + # s16le: never split a sample across chunk boundaries. + usable = len(raw) - (len(raw) % 2) + leftover = raw[usable:] + if usable == 0: + continue + samples = np.frombuffer(raw[:usable], dtype=np.int16).copy().astype(np.float32) / 32767.0 + pending.append(samples) + pending_len += len(samples) + if pending_len < min_chunk_samples: + continue + audio_chunk = np.concatenate(pending) + pending.clear() + pending_len = 0 + if chunks_yielded < prebuffer_chunks: + prebuffer.append(audio_chunk) + chunks_yielded += 1 + if chunks_yielded == prebuffer_chunks: + yield (PCM_SAMPLE_RATE, np.concatenate(prebuffer)) + prebuffer.clear() + else: + yield (PCM_SAMPLE_RATE, audio_chunk) + remaining = prebuffer + pending + if remaining: + yield (PCM_SAMPLE_RATE, np.concatenate(remaining)) + except httpx.TimeoutException: + raise gr.Error("Request timed out. The server may be busy.") from None + except httpx.ConnectError: + raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?") from None + + +def on_stream_change(stream: bool): + """Streaming only makes sense with raw PCM.""" + if stream: + return gr.update(value="pcm", interactive=False) + return gr.update(value="wav", interactive=True) + + +def build_interface(api_base: str, stream_chunk_seconds: float = 0.5): + with gr.Blocks(title="Audio8 TTS Preview Demo") as demo: + gr.Markdown("# Audio8 TTS Preview 0.6B - Text to Speech") + gr.Markdown(f"**Server:** `{api_base}` | **Model:** Audio8/Audio8-TTS-Preview-0.6b | **Output:** 44.1 kHz") + + with gr.Row(): + with gr.Column(scale=3): + text_input = gr.Textbox(label="Text to Synthesize", placeholder="Enter text here...", lines=4) + + with gr.Accordion("Voice Cloning (optional)", open=False): + gr.Markdown( + "Upload or link a short reference recording (5-30 s) and paste its **exact** transcript. " + "A transcript that does not match the audio degrades both stability and speaker similarity." + ) + ref_audio = gr.Audio(label="Reference Audio", type="numpy", sources=["upload", "microphone"]) + ref_audio_url = gr.Textbox( + label="Reference Audio URL (alternative to upload)", + placeholder="https://example.com/reference.wav", + lines=1, + ) + ref_text = gr.Textbox( + label="Reference Audio Transcript (required for cloning)", + placeholder="Exact transcript of the reference audio...", + lines=2, + ) + + with gr.Row(): + response_format = gr.Dropdown( + choices=["wav", "mp3", "flac", "pcm"], value="wav", label="Audio Format", scale=1 + ) + stream_checkbox = gr.Checkbox( + label="Stream output", value=False, info="Progressive PCM streaming", scale=1 + ) + + generate_btn = gr.Button("Generate Speech", variant="primary", size="lg") + + with gr.Column(scale=2): + audio_output = gr.Audio(label="Generated Audio", interactive=False, streaming=True, autoplay=True) + gr.Markdown( + "### About\n" + "- **Audio8 TTS Preview** — 0.6B DualAR model (slow AR + fast AR codebook decoder)\n" + "- **Languages**: Cantonese, Chinese, Dutch, English, French, German, " + "Italian, Japanese, Korean, Polish, Spanish\n" + "- **Voice cloning**: zero-shot from a reference clip + transcript\n" + "- **44.1 kHz** output, ~21.5 codec frames per second" + ) + + stream_checkbox.change(fn=on_stream_change, inputs=[stream_checkbox], outputs=[response_format]) + all_inputs = [text_input, ref_audio, ref_audio_url, ref_text, response_format] + + def dispatch(stream_enabled, *args): + if stream_enabled: + yield from generate_speech_stream(api_base, *args, chunk_seconds=stream_chunk_seconds) + else: + yield generate_speech(api_base, *args) + + generate_btn.click(fn=dispatch, inputs=[stream_checkbox] + all_inputs, outputs=[audio_output]) + demo.queue() + return demo + + +def main() -> None: + parser = argparse.ArgumentParser(description="Gradio demo for Audio8 TTS Preview") + parser.add_argument("--api-base", default=DEFAULT_API_BASE, help=f"API base URL (default: {DEFAULT_API_BASE})") + parser.add_argument("--host", default="0.0.0.0", help="Gradio host") + parser.add_argument("--port", type=int, default=7861, help="Gradio port") + parser.add_argument("--share", action="store_true", help="Create a public share link") + parser.add_argument("--stream-chunk-seconds", type=float, default=0.5, help="Seconds of audio per yielded chunk") + args = parser.parse_args() + + print(f"Connecting to vLLM server at: {args.api_base}") + build_interface(args.api_base, stream_chunk_seconds=args.stream_chunk_seconds).launch( + server_name=args.host, server_port=args.port, share=args.share + ) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/text_to_speech/audio8_tts/run_gradio_demo.sh b/examples/online_serving/text_to_speech/audio8_tts/run_gradio_demo.sh new file mode 100644 index 00000000000..5588b921e66 --- /dev/null +++ b/examples/online_serving/text_to_speech/audio8_tts/run_gradio_demo.sh @@ -0,0 +1,44 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +# Launch the Audio8 TTS Preview server + Gradio demo together. +# +# Usage: +# ./run_gradio_demo.sh +# CUDA_VISIBLE_DEVICES=0 PORT=8092 GRADIO_PORT=7861 ./run_gradio_demo.sh + +set -e + +MODEL="${MODEL:-Audio8/Audio8-TTS-Preview-0.6b}" +PORT="${PORT:-8092}" +GRADIO_PORT="${GRADIO_PORT:-7861}" +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +echo "Starting Audio8 TTS Preview server (port $PORT)..." +vllm serve "$MODEL" \ + --omni \ + --host 0.0.0.0 \ + --port "$PORT" & +SERVER_PID=$! + +cleanup() { + echo "Stopping server (PID $SERVER_PID)..." + kill $SERVER_PID 2>/dev/null + wait $SERVER_PID 2>/dev/null +} +trap cleanup EXIT + +echo "Waiting for server to start..." +for _ in $(seq 1 120); do + if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then + echo "Server ready." + break + fi + sleep 2 +done + +echo "Starting Gradio demo (port $GRADIO_PORT)..." +python "$SCRIPT_DIR/gradio_demo.py" \ + --api-base "http://localhost:$PORT" \ + --port "$GRADIO_PORT" diff --git a/examples/online_serving/text_to_speech/audio8_tts/run_server.sh b/examples/online_serving/text_to_speech/audio8_tts/run_server.sh new file mode 100644 index 00000000000..a208f0a5c2b --- /dev/null +++ b/examples/online_serving/text_to_speech/audio8_tts/run_server.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +# Launch a vLLM-Omni server for Audio8 TTS Preview 0.6B. +# +# Usage: +# ./run_server.sh +# CUDA_VISIBLE_DEVICES=0 PORT=8092 ./run_server.sh + +set -e + +MODEL="${MODEL:-Audio8/Audio8-TTS-Preview-0.6b}" +PORT="${PORT:-8092}" + +echo "Starting Audio8 TTS Preview server with model: $MODEL" + +# --trust-remote-code is deliberately NOT passed: vllm-omni registers its own +# `arktts` config, and transformers would otherwise prefer the checkpoint's +# remote code and bypass it. +vllm serve "$MODEL" \ + --omni \ + --host 0.0.0.0 \ + --port "$PORT" diff --git a/examples/online_serving/text_to_speech/audio8_tts/speech_client.py b/examples/online_serving/text_to_speech/audio8_tts/speech_client.py new file mode 100644 index 00000000000..d0a2c6f6ab4 --- /dev/null +++ b/examples/online_serving/text_to_speech/audio8_tts/speech_client.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Client for Audio8 TTS Preview via the /v1/audio/speech endpoint. + +Examples: + # Basic TTS + python speech_client.py --text "Welcome to Audio8 TTS." + + # Zero-shot voice cloning (the transcript must match the reference audio) + python speech_client.py --text "Welcome to Audio8 TTS." \ + --ref-audio ref.wav --ref-text "The exact transcript of the reference recording." + + # Streaming PCM output + python speech_client.py --text "Welcome to Audio8 TTS." --stream --output output.pcm +""" + +import argparse +import base64 +import os + +import httpx + +DEFAULT_API_BASE = "http://localhost:8092" +DEFAULT_API_KEY = "EMPTY" +DEFAULT_MODEL = "Audio8/Audio8-TTS-Preview-0.6b" +#: The codec runs at 44.1 kHz, so raw PCM must be played back at that rate. +PCM_SAMPLE_RATE = 44100 + + +def encode_audio_to_base64(audio_path: str) -> str: + """Encode a local audio file as a base64 data URL.""" + if not os.path.exists(audio_path): + raise FileNotFoundError(f"Audio file not found: {audio_path}") + ext = audio_path.lower().rsplit(".", 1)[-1] + mime_type = {"wav": "audio/wav", "mp3": "audio/mpeg", "flac": "audio/flac", "ogg": "audio/ogg"}.get( + ext, "audio/wav" + ) + with open(audio_path, "rb") as handle: + audio_b64 = base64.b64encode(handle.read()).decode("utf-8") + return f"data:{mime_type};base64,{audio_b64}" + + +def build_payload(args) -> dict: + payload = { + "model": args.model, + "input": args.text, + "response_format": args.response_format, + } + if args.ref_audio: + if args.ref_audio.startswith(("http://", "https://")): + payload["ref_audio"] = args.ref_audio + else: + payload["ref_audio"] = encode_audio_to_base64(args.ref_audio) + if args.ref_text: + payload["ref_text"] = args.ref_text + if args.max_new_tokens is not None: + payload["max_new_tokens"] = args.max_new_tokens + if args.stream: + payload["stream"] = True + payload["stream_format"] = "audio" + payload["response_format"] = "pcm" + return payload + + +def run_tts(args) -> None: + payload = build_payload(args) + api_url = f"{args.api_base}/v1/audio/speech" + headers = {"Content-Type": "application/json", "Authorization": f"Bearer {args.api_key}"} + + print(f"Model: {args.model}") + print(f"Text: {args.text}") + if args.ref_audio: + print(f"Voice cloning: ref_audio={args.ref_audio!r} ref_text={args.ref_text!r}") + + if args.stream: + output_path = args.output or "output.pcm" + with ( + httpx.Client(timeout=300.0) as client, + client.stream("POST", api_url, json=payload, headers=headers) as resp, + ): + if resp.status_code != 200: + print(f"Error {resp.status_code}: {resp.read().decode()}") + return + total_bytes = 0 + with open(output_path, "wb") as handle: + for chunk in resp.iter_bytes(): + handle.write(chunk) + total_bytes += len(chunk) + seconds = total_bytes / 2 / PCM_SAMPLE_RATE + print(f"Streamed {total_bytes} bytes (~{seconds:.2f}s of s16le @ {PCM_SAMPLE_RATE} Hz) to {output_path}") + print(f"Play with: ffplay -f s16le -ar {PCM_SAMPLE_RATE} -ac 1 {output_path}") + return + + with httpx.Client(timeout=300.0) as client: + response = client.post(api_url, json=payload, headers=headers) + if response.status_code != 200: + print(f"Error {response.status_code}: {response.text}") + return + if response.headers.get("content-type", "").startswith("application/json"): + print(f"Error: {response.text}") + return + + output_path = args.output or f"output.{args.response_format}" + with open(output_path, "wb") as handle: + handle.write(response.content) + print(f"Audio saved to: {output_path} ({len(response.content)} bytes)") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Audio8 TTS Preview client") + parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL") + parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key") + parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name") + parser.add_argument("--text", required=True, help="Text to synthesize") + parser.add_argument("--ref-audio", default=None, help="Reference audio for voice cloning (path or URL)") + parser.add_argument("--ref-text", default=None, help="Transcript of the reference audio") + parser.add_argument("--max-new-tokens", type=int, default=None, help="Cap the number of generated codec frames") + parser.add_argument("--stream", action="store_true", help="Stream PCM instead of returning a whole file") + parser.add_argument( + "--response-format", + default="wav", + choices=["wav", "mp3", "flac", "pcm", "aac", "opus"], + help="Audio format (default: wav)", + ) + parser.add_argument("--output", "-o", default=None, help="Output file path") + run_tts(parser.parse_args()) + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/offline_inference/test_audio8_tts.py b/tests/e2e/offline_inference/test_audio8_tts.py new file mode 100644 index 00000000000..687fff4ae8b --- /dev/null +++ b/tests/e2e/offline_inference/test_audio8_tts.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""E2E offline tests for Audio8 TTS Preview (DualAR, 44.1 kHz codec). + +The prompt is pre-tokenized here (not built from ``mm_processor_kwargs``) +because Audio8 TTS shares Fish Speech's protocol: text-only prompts are plain +token ids, and voice-clone prompts are a placeholder of the exact final length +whose real embeddings are built model-side from the encoded reference audio. +""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner, OmniRunnerHandler +from tests.helpers.stage_config import get_deploy_config_path + +MODEL = os.environ.get("AUDIO8_TTS_MODEL_PATH", "Audio8/Audio8-TTS-Preview-0.6b") +DEPLOY_CONFIG = get_deploy_config_path("audio8_tts.yaml") +SAMPLE_RATE = 44100 +TEXT = "The weather is nice today, perfect for a walk in the park." + +# Four distinct inputs for the concurrency test: distinct text is required so a +# leaked per-request codec buffer (which would make two requests decode to the +# same audio) actually fails the pairwise-difference assertion. +CONCURRENT_TEXTS = [ + "The weather is nice today, perfect for a walk in the park.", + "She sold seashells by the seashore all through the summer.", + "Our train departs at a quarter past nine tomorrow morning.", + "He planted rows of tomatoes and basil behind the old house.", +] + +# The checkpoint registers its own `arktts` config in vllm-omni, so remote code +# must stay disabled (transformers would otherwise prefer the checkpoint's +# auto_map and bypass the Qwen2-shaped view Qwen2Model needs). +_OMNI_RUNNER_PARAM = (MODEL, DEPLOY_CONFIG, {"trust_remote_code": False}) + +pytestmark = pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True) + + +def _text_only_prompt(text: str = TEXT) -> dict: + from transformers import AutoTokenizer + + from vllm_omni.model_executor.models.audio8_tts.prompt_utils import build_text_only_prompt_ids + + tokenizer = AutoTokenizer.from_pretrained(MODEL) + prompt_ids, normalized_text = build_text_only_prompt_ids(tokenizer, text) + return {"prompt_token_ids": prompt_ids, "additional_information": {"text": [normalized_text]}} + + +# At ``core_model`` (L1/L2) the harness patches every stage to +# ``load_format: dummy``, so the AR runs on random weights: it never emits EOS +# and always generates to ``max_tokens`` (1024 frames ~= 47.5 s). Duration +# bounds are therefore only meaningful once real weights load, at +# ``advanced_model`` / ``full_model``. +_REAL_WEIGHT_RUN_LEVELS = frozenset({"advanced_model", "full_model"}) + + +def _duration_bounds(run_level: str) -> dict[str, float]: + if run_level not in _REAL_WEIGHT_RUN_LEVELS: + return {} + return {"min_duration_s": 0.5, "max_duration_s": 30.0} + + +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.tts +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_text_to_audio_001(omni_runner: OmniRunner, run_level: str) -> None: + """Default deploy, single request, text-only synthesis. + + Deploy Setting: audio8_tts.yaml + Input Modal: text + Output Modal: audio + """ + prompt = _text_only_prompt() + request_config = { + "input": TEXT, + "prompt_token_ids": prompt["prompt_token_ids"], + "additional_information": prompt["additional_information"], + "response_format": "wav", + "expected_sample_rate": SAMPLE_RATE, + **_duration_bounds(run_level), + } + OmniRunnerHandler(omni_runner).send_tokenized_tts_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.tts +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_text_to_audio_concurrent_002(omni_runner: OmniRunner, run_level: str) -> None: + """Four concurrent requests with distinct text: stage 0 runs ``max_num_seqs: 4``. + + Per-request codec state is keyed by request id. The four inputs are + deliberately distinct and the decoded outputs are asserted pairwise + different, so leaked codec state (which would make two requests decode to + identical audio) fails here, and only here. + """ + prompts = [_text_only_prompt(text) for text in CONCURRENT_TEXTS] + request_config = { + "input": CONCURRENT_TEXTS[0], + "prompt_token_ids": [p["prompt_token_ids"] for p in prompts], + "additional_information": [p["additional_information"] for p in prompts], + "response_format": "wav", + "expected_sample_rate": SAMPLE_RATE, + "assert_distinct_outputs": True, + **_duration_bounds(run_level), + } + OmniRunnerHandler(omni_runner).send_tokenized_tts_request(request_config) diff --git a/tests/e2e/online_serving/test_audio8_tts.py b/tests/e2e/online_serving/test_audio8_tts.py new file mode 100644 index 00000000000..1ef725d684b --- /dev/null +++ b/tests/e2e/online_serving/test_audio8_tts.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""E2E online tests for Audio8 TTS Preview via /v1/audio/speech. + +Real model inference (no mocks): text-only synthesis, streaming, and zero-shot +voice cloning. The reference audio is inlined as a base64 data URL because CI +hosts cannot fetch external URLs. +""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +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 = os.environ.get("AUDIO8_TTS_MODEL_PATH", "Audio8/Audio8-TTS-Preview-0.6b") +DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 300.0 +# The codec decodes at 44.1 kHz; raw PCM carries no header, so assertions that +# depend on the rate (HNR pitch search) have to be told explicitly. +SAMPLE_RATE = 44100 +MAX_CONCURRENT = 4 + +# ~0.5 s of 44.1 kHz mono PCM_16 WAV. +_MIN_AUDIO_BYTES = 40_000 + +# Reused vendored asset (3.5 s, Chinese) plus its exact transcript: Audio8 TTS +# requires the reference transcript to match the recording. +REF_AUDIO_URL = get_asset_path("cosyvoice3/zero_shot_prompt.wav", as_data_url=True) +REF_TEXT = "希望你以后能够做的比我还好呦。" + + +def get_prompt(prompt_type: str = "en") -> str: + prompts = { + "en": "The weather is nice today, perfect for a walk in the park.", + "zh": "今天天气很好，非常适合去公园散步。", + } + return prompts.get(prompt_type, prompts["en"]) + + +tts_server_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=get_deploy_config_path("audio8_tts.yaml"), + # No --trust-remote-code: vllm-omni owns the `arktts` config. + server_args=["--disable-log-stats"], + ), + id="audio8_tts", + ) +] + + +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.tts +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_text_to_audio_001(omni_server, openai_client) -> None: + """Baseline smoke: default deploy, non-streaming WAV. + + Deploy Setting: audio8_tts.yaml + Input Modal: text + Output Modal: audio + Input Setting: stream=False + Datasets: few requests (max_num_seqs=4) + """ + request_config = { + "model": omni_server.model, + "input": get_prompt(), + "stream": False, + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "response_format": "wav", + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + openai_client.send_audio_speech_request(request_config, request_num=MAX_CONCURRENT) + + +@pytest.mark.advanced_model +@pytest.mark.tts +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_text_to_audio_streaming_002(omni_server, openai_client) -> None: + """Streaming PCM: exercises the async_chunk path end to end. + + Deploy Setting: audio8_tts.yaml (async_chunk: true) + Input Setting: stream=True, response_format=pcm + """ + request_config = { + "model": omni_server.model, + "input": get_prompt(), + "stream": True, + "stream_format": "audio", + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "response_format": "pcm", + "min_audio_bytes": _MIN_AUDIO_BYTES, + # Raw PCM has no header: without this the harness assumes 24 kHz and the + # HNR pitch search lands on 147-735 Hz instead of 80-400 Hz, scoring + # clean 44.1 kHz speech ~1.9 dB too low. + "expected_sample_rate": SAMPLE_RATE, + } + openai_client.send_audio_speech_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.tts +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_voice_clone_003(omni_server, openai_client) -> None: + """Zero-shot voice cloning: the reference audio is encoded model-side and + its codec codes are spliced into the prompt. + + Input Modal: text + ref_audio + ref_text + """ + request_config = { + "model": omni_server.model, + "input": get_prompt("zh"), + "ref_audio": REF_AUDIO_URL, + "ref_text": REF_TEXT, + "stream": False, + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "response_format": "wav", + "min_audio_bytes": _MIN_AUDIO_BYTES, + } + openai_client.send_audio_speech_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.tts +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_voice_clone_requires_ref_text_004(omni_server, openai_client) -> None: + """``ref_audio`` without ``ref_text`` must be rejected, not silently + synthesized with an untethered voice.""" + request_config = { + "model": omni_server.model, + "input": get_prompt(), + "ref_audio": REF_AUDIO_URL, + "stream": False, + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "response_format": "wav", + "status_code": (400, 422), + "err_message": "ref_text", + } + openai_client.send_audio_speech_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.tts +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_empty_input_is_rejected_005(omni_server, openai_client) -> None: + request_config = { + "model": omni_server.model, + "input": " ", + "stream": False, + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "response_format": "wav", + "status_code": (400, 422), + "err_message": "Input text cannot be empty", + } + openai_client.send_audio_speech_request(request_config) diff --git a/tests/entrypoints/openai/test_audio8_tts_adapter.py b/tests/entrypoints/openai/test_audio8_tts_adapter.py new file mode 100644 index 00000000000..ffd22d9f35e --- /dev/null +++ b/tests/entrypoints/openai/test_audio8_tts_adapter.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Serving-side registration invariants for the Audio8 TTS adapter. + +The engine never names ``audio8_tts`` directly: ``/v1/audio/speech`` resolves the +adapter by matching the deployed stage's ``model_stage`` key against the +``stage_keys`` each adapter declares. That makes the pipeline definition and the +adapter two halves of one contract, with nothing in the type system holding them +together -- rename either side and the endpoint silently stops recognising the +model (falling through to "unknown TTS model") instead of failing at startup. + +These tests pin the contract from both ends. +""" + +import pytest + +from vllm_omni.entrypoints.openai.tts_adapters import ( + all_tts_model_types, + all_tts_stage_keys, + detect_tts_model_type, + resolve_adapter, +) +from vllm_omni.entrypoints.openai.tts_adapters.audio8_tts import Audio8TTSAdapter +from vllm_omni.model_executor.models.audio8_tts.pipeline import AUDIO8_TTS_PIPELINE + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +SLOW_AR_STAGE_KEY = "audio8_tts_slow_ar" + + +def _pipeline_stage_keys() -> list[str]: + return [stage.model_stage for stage in AUDIO8_TTS_PIPELINE.stages] + + +def test_pipeline_stage_key_is_the_one_the_adapter_claims(): + """The shipped pipeline must name the stage the adapter matches on.""" + assert SLOW_AR_STAGE_KEY in _pipeline_stage_keys() + assert SLOW_AR_STAGE_KEY in Audio8TTSAdapter.stage_keys + + +def test_detection_resolves_slow_ar_stage_to_the_audio8_adapter(): + assert detect_tts_model_type(SLOW_AR_STAGE_KEY, "Audio8TTSSlowARForConditionalGeneration") == "audio8_tts" + assert resolve_adapter("audio8_tts") is Audio8TTSAdapter + + +def test_adapter_is_registered_in_the_global_tables(): + """Importing the package must register the adapter, not just define it.""" + assert "audio8_tts" in all_tts_model_types() + assert SLOW_AR_STAGE_KEY in all_tts_stage_keys() + + +def test_codec_decoder_stage_does_not_serve_speech(): + """Only stage 0 answers /v1/audio/speech; the codec stage must not match.""" + codec_stage_key = _pipeline_stage_keys()[1] + assert codec_stage_key != SLOW_AR_STAGE_KEY + assert detect_tts_model_type(codec_stage_key, "Audio8TTSCodecDecoder") != "audio8_tts" + + +def test_apply_sampling_overrides_defaults_tts_local_seed_for_fast_ar(): + """The deploy seed must reach the Fast AR through ``tts_local_seed``. + + Only the Slow AR is seeded via ``sampling_metadata.generators``; the residual + codebooks read ``tts_local_seed`` and otherwise fall back to global RNG, which + makes an unseeded request's audio depend on batch composition. The adapter + defaults it from the stage-0 seed without clobbering an explicit value. + """ + from types import SimpleNamespace + + from vllm.sampling_params import SamplingParams + + adapter = Audio8TTSAdapter.__new__(Audio8TTSAdapter) + request = SimpleNamespace(max_new_tokens=None) + + # Deploy default seed propagates to tts_local_seed. + out = adapter.apply_sampling_overrides([SamplingParams(seed=42)], request) + assert out[0].extra_args["tts_local_seed"] == 42 + + # An explicit tts_local_seed is left intact (setdefault, not overwrite). + pinned = [SamplingParams(seed=42, extra_args={"tts_local_seed": 7})] + assert adapter.apply_sampling_overrides(pinned, request)[0].extra_args["tts_local_seed"] == 7 + + # No seed at all -> nothing injected (non-deterministic by request contract). + out_unseeded = adapter.apply_sampling_overrides([SamplingParams(seed=None)], request) + assert (out_unseeded[0].extra_args or {}).get("tts_local_seed") is None diff --git a/tests/entrypoints/openai_api/test_tts_detection.py b/tests/entrypoints/openai_api/test_tts_detection.py index 3b8d1799764..d54ce36d51b 100644 --- a/tests/entrypoints/openai_api/test_tts_detection.py +++ b/tests/entrypoints/openai_api/test_tts_detection.py @@ -132,6 +132,8 @@ def _legacy_detect(model_stage, model_arch): "audex_thinker", "audex_tta_thinker", "audex_xcodec", + "audio8_tts_codec_decoder", + "audio8_tts_slow_ar", "audio_generation", "audio_tokenizer", "audio_vae", diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py index 65d28109fda..dd25aca7187 100644 --- a/tests/helpers/assertions.py +++ b/tests/helpers/assertions.py @@ -659,7 +659,11 @@ def _compute_pcm_hnr_db(pcm_samples: np.ndarray, sr: int = _PCM_SPEECH_SAMPLE_RA return float(np.mean(hnr_values)) if hnr_values else 0.0 -def _assert_pcm_int16_speech_hnr(audio_bytes: bytes, min_hnr_db: float = _MIN_PCM_SPEECH_HNR_DB) -> None: +def _assert_pcm_int16_speech_hnr( + audio_bytes: bytes, + min_hnr_db: float = _MIN_PCM_SPEECH_HNR_DB, + sr: int = _PCM_SPEECH_SAMPLE_RATE_HZ, +) -> None: """Validate harmonic-to-noise ratio on raw int16 PCM from /v1/audio/speech. min_hnr_db defaults to the global _MIN_PCM_SPEECH_HNR_DB (1.0 dB), @@ -668,12 +672,16 @@ def _assert_pcm_int16_speech_hnr(audio_bytes: bytes, min_hnr_db: float = _MIN_PC intrinsically around -2 dB) can pass a lower per-test threshold via request_config["min_hnr_db"] to keep the catastrophic-failure check while not gating CI on a model-intrinsic property. + + ``sr`` must be the stream's real sample rate: the autocorrelation lag window + is derived from it, so a wrong rate detunes the 80-400 Hz pitch search and + understates HNR for models that do not decode at 24 kHz. """ assert audio_bytes is not None and len(audio_bytes) >= 2, "missing PCM bytes" assert len(audio_bytes) % 2 == 0, "PCM byte length must be aligned to int16" pcm_samples = np.frombuffer(audio_bytes, dtype=np.int16).astype(np.float32) / 32768.0 - hnr = _compute_pcm_hnr_db(pcm_samples) - print(f"PCM speech HNR: {hnr:.2f} dB (threshold: {min_hnr_db} dB)") + hnr = _compute_pcm_hnr_db(pcm_samples, sr=sr) + print(f"PCM speech HNR: {hnr:.2f} dB (threshold: {min_hnr_db} dB, sr={sr})") assert hnr >= min_hnr_db, ( f"Audio distortion detected: HNR={hnr:.2f} dB < {min_hnr_db} dB. " "Voice clone decoder may be losing ref_code speaker context on later chunks." @@ -934,7 +942,15 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any], if run_level in {"advanced_model", "full_model"}: if req_fmt == "pcm" and response.audio_bytes: min_hnr_db = float(request_config.get("min_hnr_db", _MIN_PCM_SPEECH_HNR_DB)) - _assert_pcm_int16_speech_hnr(response.audio_bytes, min_hnr_db=min_hnr_db) + # Raw PCM carries no header, so the HNR pitch search has to be told + # the rate. Defaulting to 24 kHz for a 44.1 kHz model shifts the + # search window to 147-735 Hz and misses the fundamental entirely, + # scoring clean speech ~1.9 dB lower than it is. + _assert_pcm_int16_speech_hnr( + response.audio_bytes, + min_hnr_db=min_hnr_db, + sr=int(request_config.get("expected_sample_rate") or _PCM_SPEECH_SAMPLE_RATE_HZ), + ) transcript = _resolve_audio_transcript(response, request_config, run_level, speech_api=True) if transcript is not None: diff --git a/tests/helpers/client.py b/tests/helpers/client.py index 6a6a9b7ed31..ed5158b92de 100644 --- a/tests/helpers/client.py +++ b/tests/helpers/client.py @@ -12,6 +12,7 @@ import io import json import time +from collections.abc import Mapping from dataclasses import asdict, dataclass from io import BytesIO from pathlib import Path @@ -2123,6 +2124,123 @@ def send_audio_speech_request(self, request_config: dict[str, Any]) -> OmniRespo assert_audio_speech_response(result, request_config, run_level="core_model") return result + def send_tokenized_tts_request(self, request_config: dict[str, Any]) -> OmniResponse: + """Offline TTS for models whose engine prompt is *pre-tokenized*. + + Fish-Speech-style DualAR models (and any model whose prompt embeds codec + codes) cannot go through :meth:`send_audio_speech_request`: the prompt is + built by the caller as ``prompt_token_ids`` plus per-request + ``additional_information``, not from ``mm_processor_kwargs``. + + request_config keys: + - ``prompt_token_ids``: ``list[int]`` or ``list[list[int]]`` (one entry + per request). + - ``additional_information``: dict, or list of dicts matching + ``prompt_token_ids``. + - ``expected_sample_rate``: asserted against the reported ``sr``. + - ``min_duration_s`` / ``max_duration_s``: bounds on decoded audio. + - plus every key ``assert_audio_speech_response`` understands + (``response_format``, ``min_audio_bytes``, ``input``, ...). + """ + token_ids = request_config.get("prompt_token_ids") + if not token_ids: + raise ValueError("request_config must contain 'prompt_token_ids'") + batched_ids = token_ids if isinstance(token_ids[0], (list, tuple)) else [token_ids] + + extra = request_config.get("additional_information") + if extra is None: + batched_extra: list[dict[str, Any] | None] = [None] * len(batched_ids) + elif isinstance(extra, dict): + batched_extra = [extra] * len(batched_ids) + else: + batched_extra = list(extra) + if len(batched_extra) != len(batched_ids): + raise ValueError("'additional_information' must match the number of prompts") + + prompts: list[dict[str, Any]] = [] + for ids, info in zip(batched_ids, batched_extra, strict=True): + prompt: dict[str, Any] = {"prompt_token_ids": list(ids)} + if info is not None: + prompt["additional_information"] = info + prompts.append(prompt) + + start_time = time.perf_counter() + outputs = self.runner.omni.generate(prompts) + e2e_latency = time.perf_counter() - start_time + + assert len(outputs) == len(prompts), f"Expected {len(prompts)} outputs, got {len(outputs)}" + expected_sr = request_config.get("expected_sample_rate") + min_duration = request_config.get("min_duration_s") + max_duration = request_config.get("max_duration_s") + + result = OmniResponse() + wavs: list[torch.Tensor] = [] + for index, stage_output in enumerate(outputs): + assert stage_output.outputs, f"Request {index} produced no output" + multimodal_output = stage_output.outputs[0].multimodal_output + assert isinstance(multimodal_output, Mapping), ( + f"Request {index}: expected a mapping, got {type(multimodal_output)}" + ) + + # Never use ``a or b`` here: truthiness on a tensor raises. + audio = multimodal_output.get("model_outputs") + if audio is None: + audio = multimodal_output.get("audio") + assert audio is not None, f"Request {index}: no audio key in {sorted(multimodal_output)}" + if isinstance(audio, list): + chunks = [torch.as_tensor(chunk).float().cpu().reshape(-1) for chunk in audio if chunk is not None] + assert chunks, f"Request {index}: audio list contained no tensors" + audio = torch.cat(chunks, dim=0) + wav = torch.as_tensor(audio).float().cpu().reshape(-1) + assert wav.numel() > 0, f"Request {index}: decoded audio is empty" + wavs.append(wav) + + sr_raw = multimodal_output.get("sr") + if isinstance(sr_raw, list): + sr_raw = sr_raw[-1] if sr_raw else None + assert sr_raw is not None, f"Request {index}: missing sample rate" + sample_rate = int(sr_raw.item() if hasattr(sr_raw, "item") else sr_raw) + if expected_sr is not None: + assert sample_rate == int(expected_sr), ( + f"Request {index}: expected sample rate {expected_sr}, got {sample_rate}" + ) + + duration_s = wav.numel() / sample_rate + if min_duration is not None: + assert duration_s >= float(min_duration), ( + f"Request {index}: audio too short ({duration_s:.2f}s < {min_duration}s)" + ) + if max_duration is not None: + assert duration_s <= float(max_duration), ( + f"Request {index}: audio too long ({duration_s:.2f}s > {max_duration}s)" + ) + + wav_buf = io.BytesIO() + sf.write(wav_buf, wav.numpy(), samplerate=sample_rate, format="WAV", subtype="PCM_16") + result = OmniResponse( + success=True, + audio_bytes=wav_buf.getvalue(), + audio_format="audio/wav", + e2e_latency=e2e_latency, + ) + assert_audio_speech_response(result, request_config, run_level="core_model") + + if request_config.get("assert_distinct_outputs") and len(wavs) > 1: + # Distinct inputs must decode to distinct audio. Leaked per-request + # codec state shows up as two requests producing (near-)identical + # waveforms; different lengths already prove distinctness. + for i in range(len(wavs)): + for j in range(i + 1, len(wavs)): + a, b = wavs[i], wavs[j] + if a.numel() != b.numel(): + continue + max_diff = (a - b).abs().max().item() + assert max_diff > 1e-3, ( + f"Requests {i} and {j} produced near-identical audio " + f"(max sample diff {max_diff:.2e}); possible codec-state crosstalk" + ) + return result + def start_profile(self, profile_prefix: str | None = None, stages: list[int] | None = None) -> list[Any]: return self.runner.start_profile(profile_prefix=profile_prefix, stages=stages) diff --git a/tests/model_executor/models/audio8_tts/__init__.py b/tests/model_executor/models/audio8_tts/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/model_executor/models/audio8_tts/test_audio8_tts_codec_decoder.py b/tests/model_executor/models/audio8_tts/test_audio8_tts_codec_decoder.py new file mode 100644 index 00000000000..a737a7f4b75 --- /dev/null +++ b/tests/model_executor/models/audio8_tts/test_audio8_tts_codec_decoder.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import functools + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +NUM_CODEBOOKS = 10 +#: Samples the fake codec returns per frame, so trims are easy to assert. +SAMPLES_PER_FRAME = 100 + + +@functools.lru_cache(maxsize=1) +def _decoder_cls(): + from vllm_omni.model_executor.models.audio8_tts.audio8_tts_codec_decoder import Audio8TTSCodecDecoder + + return Audio8TTSCodecDecoder + + +class _FakeCodec: + """Returns ``SAMPLES_PER_FRAME`` ramp samples per frame.""" + + def __init__(self): + self.calls: list[tuple[int, ...]] = [] + + def decode(self, codes_bqf: torch.Tensor) -> torch.Tensor: + self.calls.append(tuple(codes_bqf.shape)) + frames = int(codes_bqf.shape[-1]) + total = frames * SAMPLES_PER_FRAME + return torch.arange(total, dtype=torch.float32).view(1, 1, total) + + +def _make_decoder(codec: _FakeCodec, num_codebooks: int = NUM_CODEBOOKS, splits=None): + decoder = object.__new__(_decoder_cls()) + torch.nn.Module.__init__(decoder) + decoder._codec = codec + decoder._codec_device = torch.device("cpu") + decoder._codec_dtype = torch.float32 + decoder._num_codebooks = num_codebooks + decoder._sample_rate = 44100 + decoder._frame_size = 2048 + decoder._logged_stats = True + decoder._ensure_codec_loaded = lambda: None + if splits is not None: + decoder._split_request_ids = lambda ids, seq_token_counts=None: splits + return decoder + + +def test_left_context_frames_are_trimmed_from_the_emitted_chunk(): + """The decoder must emit *delta* audio: the left-context prefix exists only + so the causal codec has history and must never reach the client twice.""" + codec = _FakeCodec() + decoder = _make_decoder(codec) + codes = torch.arange(NUM_CODEBOOKS * 4, dtype=torch.long).reshape(NUM_CODEBOOKS, 4) + + out = decoder.forward( + input_ids=torch.tensor([0], dtype=torch.long), + runtime_additional_information=[{"codes": {"audio": codes}, "meta": {"left_context_size": 3}}], + ) + + audios = out.multimodal_outputs["model_outputs"] + assert len(audios) == 1 + # 4 frames decoded, 3 of context trimmed => 1 frame of new audio. + assert audios[0].shape[0] == SAMPLES_PER_FRAME + assert codec.calls == [(1, NUM_CODEBOOKS, 4)] + + +def test_tensor_codes_from_runtime_info_win_over_input_ids(): + codec = _FakeCodec() + decoder = _make_decoder(codec, num_codebooks=2) + out = decoder.forward( + input_ids=torch.zeros(64, dtype=torch.long), + runtime_additional_information=[{"codes": {"audio": torch.tensor([[1, 2], [3, 4]], dtype=torch.long)}}], + ) + assert codec.calls == [(1, 2, 2)] + assert out.multimodal_outputs["model_outputs"][0].shape[0] == 2 * SAMPLES_PER_FRAME + + +def test_flat_input_ids_are_reshaped_codebook_major(): + codec = _FakeCodec() + decoder = _make_decoder(codec, num_codebooks=2) + out = decoder.forward(input_ids=torch.arange(6, dtype=torch.long), runtime_additional_information=[{}]) + assert codec.calls == [(1, 2, 3)] + assert out.multimodal_outputs["model_outputs"][0].shape[0] == 3 * SAMPLES_PER_FRAME + + +def test_every_return_path_sets_model_outputs_and_sr(): + """A branch that omits ``model_outputs`` makes the serving layer drop the + chunk silently, so assert the key exists on all the early-exit paths.""" + codec = _FakeCodec() + + empty = _make_decoder(codec).forward(input_ids=torch.zeros(0, dtype=torch.long)) + assert empty.multimodal_outputs["model_outputs"][0].numel() == 0 + assert empty.multimodal_outputs["sr"] + + ragged = _make_decoder(codec, num_codebooks=4).forward( + input_ids=torch.arange(6, dtype=torch.long), # 6 % 4 != 0 + runtime_additional_information=[{}], + ) + assert ragged.multimodal_outputs["model_outputs"][0].numel() == 0 + + codes = torch.zeros((NUM_CODEBOOKS, 2), dtype=torch.long) + no_new = _make_decoder(codec).forward( + input_ids=torch.tensor([0], dtype=torch.long), + runtime_additional_information=[{"codes": {"audio": codes}, "meta": {"left_context_size": 2}}], + ) + assert no_new.multimodal_outputs["model_outputs"][0].numel() == 0 + assert codec.calls == [], "a context-only chunk must not be decoded" + + +def test_mixed_batch_keeps_per_request_alignment(): + """An empty request must not shift the following request's audio into its + slot: the output list index is the request index.""" + codec = _FakeCodec() + splits = [ + torch.empty((0,), dtype=torch.long), + torch.arange(NUM_CODEBOOKS * 2, dtype=torch.long), + ] + decoder = _make_decoder(codec, splits=splits) + out = decoder.forward( + input_ids=torch.arange(NUM_CODEBOOKS * 2, dtype=torch.long), + runtime_additional_information=[{}, {"meta": {"left_context_size": 1}}], + ) + audios = out.multimodal_outputs["model_outputs"] + assert len(audios) == 2 + assert audios[0].numel() == 0 + assert audios[1].shape[0] == SAMPLES_PER_FRAME + + +def test_codes_with_wrong_codebook_count_fall_back_to_input_ids(): + codec = _FakeCodec() + decoder = _make_decoder(codec, num_codebooks=2) + out = decoder.forward( + input_ids=torch.arange(4, dtype=torch.long), + runtime_additional_information=[{"codes": {"audio": torch.zeros((5, 3), dtype=torch.long)}}], + ) + assert codec.calls == [(1, 2, 2)] + assert out.multimodal_outputs["model_outputs"][0].shape[0] == 2 * SAMPLES_PER_FRAME diff --git a/tests/model_executor/models/audio8_tts/test_audio8_tts_prompt_utils.py b/tests/model_executor/models/audio8_tts/test_audio8_tts_prompt_utils.py new file mode 100644 index 00000000000..9047d06f565 --- /dev/null +++ b/tests/model_executor/models/audio8_tts/test_audio8_tts_prompt_utils.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest + +from vllm_omni.model_executor.models.audio8_tts.prompt_utils import ( + build_text_only_prompt_ids, + build_voice_clone_prompt_ids, + estimate_voice_clone_prompt_len, + normalize_text, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _SegmentTokenizer: + """Tokenizer stub that maps each encoded string to one sentinel id. + + Lets the tests assert the *segment sequence* the prompt builder emits, which + is what must stay byte-identical to the reference ``ArkttsProcessor``. + """ + + def __init__(self): + self.vocab: dict[str, int] = {} + + def encode(self, text: str, add_special_tokens: bool = True): + assert add_special_tokens is False, "Audio8 prompts must not add extra special tokens" + return [self.vocab.setdefault(text, 1000 + len(self.vocab))] + + def segment_of(self, token_id: int) -> str: + return next(text for text, tid in self.vocab.items() if tid == token_id) + + +def _segments(tokenizer: _SegmentTokenizer, ids: list[int]) -> list[str]: + return [tokenizer.segment_of(i) for i in ids] + + +def test_text_only_prompt_matches_reference_segment_order(): + tokenizer = _SegmentTokenizer() + ids, normalized = build_text_only_prompt_ids(tokenizer, " Hello world ") + + assert normalized == "Hello world" + assert _segments(tokenizer, ids) == [ + "<|im_start|>system\n", + "convert the provided text to speech", + "<|im_end|>\n", + "<|im_start|>user\n", + "Hello world", + "<|im_end|>\n", + "<|im_start|>assistant\n<|voice|>", + ] + + +def test_voice_clone_prompt_places_reference_codes_between_prefix_and_suffix(): + tokenizer = _SegmentTokenizer() + semantic_ids = [151678, 151679, 151680] + ids, ref_start, text, ref_text = build_voice_clone_prompt_ids( + tokenizer, "Target text", "Reference transcript", semantic_ids + ) + + assert text == "Target text" + # A reference transcript without an explicit tag gets speaker 0. + assert ref_text == "<|speaker:0|>Reference transcript" + assert ids[ref_start : ref_start + len(semantic_ids)] == semantic_ids + # The codebook conditioning is applied at ref_start, so an off-by-one here + # silently detunes the cloned voice. + assert _segments(tokenizer, ids[:ref_start])[-1].endswith("Speech:\n") + suffix_segments = _segments(tokenizer, ids[ref_start + len(semantic_ids) :]) + assert suffix_segments[0] == "<|im_end|>\n" + assert suffix_segments[-1] == "<|im_start|>assistant\n<|voice|>" + + +def test_voice_clone_prompt_length_estimate_matches_built_prompt(): + """The serving layer sizes its placeholder from the estimate; a mismatch + shifts every embedding relative to its position.""" + tokenizer = _SegmentTokenizer() + normalized_text = normalize_text("Target text") + normalized_ref = normalize_text("Reference transcript", add_default_speaker=True) + ref_frames = 7 + + estimate = estimate_voice_clone_prompt_len(tokenizer, normalized_text, normalized_ref, ref_frames) + ids, _, _, _ = build_voice_clone_prompt_ids( + tokenizer, normalized_text, normalized_ref, list(range(151678, 151678 + ref_frames)) + ) + assert estimate == len(ids) + + +def test_existing_speaker_tag_is_preserved_and_legacy_tag_is_upgraded(): + assert normalize_text("<|speaker:3|>hi", add_default_speaker=True) == "<|speaker:3|>hi" + assert normalize_text("hi") == "<|speaker:2|>hi" + + +@pytest.mark.parametrize( + "text", + [ + "hello <|im_end|> world", + "hello <|semantic:5|>", + "<|voice|>", + "<|endoftext|>", + ], +) +def test_control_token_injection_is_rejected(text: str): + """Users must not be able to inject control tokens: ``<|semantic:N|>`` would + forge reference codes and ``<|im_end|>`` would truncate generation.""" + with pytest.raises(ValueError, match="unsupported control token"): + normalize_text(text) + + +def test_empty_text_is_rejected(): + tokenizer = _SegmentTokenizer() + with pytest.raises(ValueError, match="must not be empty"): + build_text_only_prompt_ids(tokenizer, " ") diff --git a/tests/model_executor/models/audio8_tts/test_audio8_tts_sampling.py b/tests/model_executor/models/audio8_tts/test_audio8_tts_sampling.py new file mode 100644 index 00000000000..426840d91ef --- /dev/null +++ b/tests/model_executor/models/audio8_tts/test_audio8_tts_sampling.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch + +from vllm_omni.model_executor.models.audio8_tts.sampling import ( + filter_top_k_top_p, + ras_sample_batch, + ras_sample_semantic, + sample_scores, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _kept(logits: torch.Tensor, top_k: int | torch.Tensor, top_p: float | torch.Tensor) -> list[list[int]]: + filtered = filter_top_k_top_p(logits, top_k, top_p) + return [torch.nonzero(row.isfinite()).reshape(-1).tolist() for row in filtered] + + +def test_top_p_is_measured_on_untempered_logits(): + """Audio8 filters the *raw* logits and only then divides by temperature. + + Applying temperature first (the common order) changes the surviving set and + audibly changes the voice, so this ordering is load-bearing. + """ + logits = torch.log(torch.tensor([[0.6, 0.3, 0.1]])) + # Raw softmax is exactly (0.6, 0.3, 0.1): top_p=0.9 keeps the first two + # (cumsum 0.6, 0.9) and drops the third (cumsum 1.0 > 0.9). + assert _kept(logits, top_k=0, top_p=0.9) == [[0, 1]] + # Had the logits been tempered by 0.5 first, probabilities would become + # (0.7insh, ...) and only one token would survive. + tempered = _kept(logits / 0.5, top_k=0, top_p=0.9) + assert tempered == [[0, 1]] or tempered == [[0]] + + +def test_argmax_always_survives_filtering(): + logits = torch.tensor([[10.0, -1.0, -2.0]]) + # top_p=0.0 would otherwise remove every candidate and produce NaNs. + assert _kept(logits, top_k=1, top_p=0.0) == [[0]] + + +def test_per_row_top_k_and_top_p_are_honoured(): + logits = torch.log(torch.tensor([[0.5, 0.3, 0.15, 0.05], [0.5, 0.3, 0.15, 0.05]])) + kept = _kept(logits, top_k=torch.tensor([1, 3]), top_p=torch.tensor([1.0, 1.0])) + assert kept == [[0], [0, 1, 2]] + + +def test_zero_temperature_row_falls_back_to_argmax(): + logits = torch.tensor([[1.0, 5.0, 2.0], [1.0, 5.0, 2.0]]) + out = sample_scores( + logits, + temperature=torch.tensor([0.0, 1.0]), + top_k=0, + top_p=1.0, + generator=torch.Generator().manual_seed(0), + ) + assert int(out[0]) == 1 + + +def test_ras_substitutes_only_repeated_semantic_ids(): + """RAS re-rolls a repeat from a flatter distribution; the EOS slot and + non-repeats must pass through untouched, or utterances get truncated.""" + num_semantic = 2 + # Row 0 will draw id 0 (a repeat) -> substituted. + # Row 1 will draw id 2 (the EOS slot) -> never substituted even if recent. + logits = torch.tensor( + [ + [10.0, -20.0, -20.0], + [-20.0, -20.0, 10.0], + ] + ) + recent = torch.tensor([[0, -1], [2, -1]]) + out = ras_sample_semantic( + logits, + recent, + temperature=1e-6, # deterministic: argmax within each draw + top_k=0, + top_p=1.0, + ras_temperature=1e-6, + ras_top_p=1.0, + num_semantic_ids=num_semantic, + ) + # Both draws are argmax, so substitution is observable only through which + # branch was taken; assert the EOS row is preserved verbatim. + assert int(out[1]) == 2 + assert int(out[0]) == 0 + + +def test_ras_without_history_returns_the_plain_draw(): + logits = torch.tensor([[10.0, 0.0, -5.0]]) + out = ras_sample_semantic( + logits, + None, + temperature=1e-6, + top_k=0, + top_p=1.0, + ras_temperature=1.0, + ras_top_p=0.9, + num_semantic_ids=2, + ) + assert int(out[0]) == 0 + + +def test_sampling_is_reproducible_for_a_seeded_generator(): + logits = torch.randn(4, 32, generator=torch.Generator().manual_seed(7)) + first = sample_scores(logits, temperature=0.7, top_k=8, top_p=0.9, generator=torch.Generator().manual_seed(1234)) + second = sample_scores(logits, temperature=0.7, top_k=8, top_p=0.9, generator=torch.Generator().manual_seed(1234)) + assert torch.equal(first, second) + + +def test_batch_routes_each_row_to_its_own_generator(): + """Regression for the shared slot-0 generator bug. + + A shared generator over ``[num_reqs, vocab]`` made every row draw from slot + 0's RNG stream, so a request's audio depended on its batch neighbours and + per-request ``seed`` was ignored for slots > 0. Flat logits make the draw + purely noise-driven so the wrong generator changes the token. + """ + # Flat logits: the sampled id is determined only by the generator's noise, + # so routing the generator to the wrong slot changes the result (coincidence ~1/256). + logits = torch.zeros(2, 256) + kw = dict(temperature=1.0, top_k=0, top_p=1.0, ras_temperature=1.0, ras_top_p=0.9, num_semantic_ids=200) + + # (1) Each batched row equals sampling that row alone with its own generator. + batched = ras_sample_batch( + logits, + None, + generators={0: torch.Generator().manual_seed(11), 1: torch.Generator().manual_seed(22)}, + num_reqs=2, + **kw, + ) + row0 = ras_sample_semantic(logits[0:1], None, generator=torch.Generator().manual_seed(11), **kw) + row1 = ras_sample_semantic(logits[1:2], None, generator=torch.Generator().manual_seed(22), **kw) + assert int(batched[0]) == int(row0) + assert int(batched[1]) == int(row1) + + # (2) Row 1 must not change when only slot 0's seed changes. + other = ras_sample_batch( + logits, + None, + generators={0: torch.Generator().manual_seed(9999), 1: torch.Generator().manual_seed(22)}, + num_reqs=2, + **kw, + ) + assert int(other[1]) == int(batched[1]) + + +def test_single_request_batch_keeps_the_fast_path(): + """``num_reqs == 1`` must match a direct single-row draw (no per-row loop).""" + logits = torch.randn(1, 48, generator=torch.Generator().manual_seed(5)) + kw = dict(temperature=0.8, top_k=16, top_p=0.9, ras_temperature=1.0, ras_top_p=0.9, num_semantic_ids=40) + batched = ras_sample_batch(logits, None, generators={0: torch.Generator().manual_seed(99)}, num_reqs=1, **kw) + direct = ras_sample_semantic(logits, None, generator=torch.Generator().manual_seed(99), **kw) + assert torch.equal(batched, direct) diff --git a/tests/model_executor/stage_input_processors/test_audio8_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_audio8_tts_async_chunk.py new file mode 100644 index 00000000000..9190a33f6c6 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_audio8_tts_async_chunk.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from collections import defaultdict +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.stage_input_processors.audio8_tts import ( + extract_last_frame, + slow_ar_to_codec_decoder_async_chunk, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +NUM_CODEBOOKS = 3 + + +def _make_transfer_manager(*, chunk_frames=2, left_context=1, initial_chunk_frames=0): + return SimpleNamespace( + code_prompt_token_ids=defaultdict(list), + put_req_chunk=defaultdict(int), + connector=SimpleNamespace( + config={ + "extra": { + "codec_chunk_frames": chunk_frames, + "codec_left_context_frames": left_context, + "initial_codec_chunk_frames": initial_chunk_frames, + } + } + ), + ) + + +def _make_request(req_id="req", *, finished=False): + return SimpleNamespace( + external_req_id=req_id, + additional_information=None, + is_finished=lambda: finished, + ) + + +def _frame(value: int) -> torch.Tensor: + return torch.full((NUM_CODEBOOKS,), value, dtype=torch.long) + + +def _step(transfer_manager, request, frame_value: int | None, *, finished: bool = False): + multimodal_output = {} if frame_value is None else {"audio_codes": _frame(frame_value).unsqueeze(0)} + return slow_ar_to_codec_decoder_async_chunk(transfer_manager, multimodal_output, request, is_finished=finished) + + +def test_codes_accumulate_across_steps_and_chunk_is_withheld_until_full(): + """Codes must accumulate; emitting per-step would starve the codec of the + context it needs and resetting would truncate the utterance.""" + transfer_manager = _make_transfer_manager(chunk_frames=2, left_context=1) + request = _make_request() + + assert _step(transfer_manager, request, 1) is None + payload = _step(transfer_manager, request, 2) + + assert payload is not None + assert len(transfer_manager.code_prompt_token_ids["req"]) == 2 + # [num_codebooks, frames], codebook-major. + assert tuple(payload.codes.audio.shape) == (NUM_CODEBOOKS, 2) + assert payload.codes.audio[0].tolist() == [1, 2] + # Nothing precedes the first chunk, so there is no left context yet. + assert payload.meta.left_context_size == 0 + + +def test_second_chunk_carries_left_context_and_reports_its_size(): + transfer_manager = _make_transfer_manager(chunk_frames=2, left_context=1) + request = _make_request() + for value in (1, 2, 3): + _step(transfer_manager, request, value) + payload = _step(transfer_manager, request, 4) + + assert payload is not None + # 1 context frame + 2 new frames. + assert tuple(payload.codes.audio.shape) == (NUM_CODEBOOKS, 3) + assert payload.codes.audio[0].tolist() == [2, 3, 4] + assert payload.meta.left_context_size == 1 + + +def test_partial_tail_is_flushed_when_the_request_finishes(): + transfer_manager = _make_transfer_manager(chunk_frames=4, left_context=2) + request = _make_request() + _step(transfer_manager, request, 1) + _step(transfer_manager, request, 2) + payload = _step(transfer_manager, request, 3, finished=True) + + assert payload is not None + assert payload.codes.audio[0].tolist() == [1, 2, 3] + assert bool(payload.meta.finished) is True + + +def test_finished_request_without_any_frame_emits_a_terminal_payload(): + transfer_manager = _make_transfer_manager() + payload = slow_ar_to_codec_decoder_async_chunk( + transfer_manager, None, _make_request(finished=True), is_finished=True + ) + assert payload is not None + assert payload.codes.audio.numel() == 0 + assert bool(payload.meta.finished) is True + + +def test_terminal_call_on_a_boundary_does_not_re_emit_the_last_chunk(): + """A finished call carrying no new frame, arriving when the stream sits + exactly on a chunk boundary, must close the stream rather than re-ship the + last window. Before the guard, the final chunk_frames frames shipped twice.""" + transfer_manager = _make_transfer_manager(chunk_frames=2, left_context=1) + request = _make_request() + for value in (1, 2, 3, 4): # length ends exactly on a chunk boundary + _step(transfer_manager, request, value) + # Separate terminal call, no new frame: the boundary chunk is already sent. + payload = _step(transfer_manager, request, None, finished=True) + assert payload is not None + assert payload.codes.audio.numel() == 0 + assert bool(payload.meta.finished) is True + + +def test_final_frame_on_a_boundary_still_emits_when_it_arrives_finished(): + """If the last frame both lands on a boundary and carries finished, its chunk + has not been shipped yet, so it must still be emitted (the guard must key on + 'no new frame this call', not merely on being finished at a boundary).""" + transfer_manager = _make_transfer_manager(chunk_frames=2, left_context=1) + request = _make_request() + _step(transfer_manager, request, 1) + payload = _step(transfer_manager, request, 2, finished=True) + assert payload is not None + assert payload.codes.audio[0].tolist() == [1, 2] + assert bool(payload.meta.finished) is True + + +def test_initial_chunk_shortens_time_to_first_audio(): + transfer_manager = _make_transfer_manager(chunk_frames=4, left_context=2, initial_chunk_frames=1) + request = _make_request() + payload = _step(transfer_manager, request, 1) + + assert payload is not None + assert payload.codes.audio[0].tolist() == [1] + assert payload.meta.left_context_size == 0 + + +def test_prefill_placeholder_frames_are_not_forwarded(): + """Prefill emits an all-zero placeholder frame per position; forwarding it + would prepend silence and desynchronise the chunk boundaries.""" + assert extract_last_frame({"audio_codes": torch.zeros((3, NUM_CODEBOOKS), dtype=torch.long)}) is None + assert extract_last_frame({}) is None + assert extract_last_frame({"audio_codes": torch.zeros((0, NUM_CODEBOOKS), dtype=torch.long)}) is None + + +def test_explicit_validity_flag_overrides_the_all_zero_heuristic(): + frame = torch.zeros((1, NUM_CODEBOOKS), dtype=torch.long) + out = extract_last_frame({"audio_codes": frame, "audio_code_valid": torch.tensor([True])}) + assert out is not None and out.tolist() == [0] * NUM_CODEBOOKS + + +def test_invalid_chunk_config_fails_loudly(): + transfer_manager = _make_transfer_manager(chunk_frames=0) + request = _make_request() + with pytest.raises(ValueError, match="codec chunk config"): + _step(transfer_manager, request, 1) diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index 5e3cfe865bb..be282da812a 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/tests/worker/test_gpu_ar_model_runner.py @@ -1829,6 +1829,7 @@ def test_in_tree_declarer_inventory_is_exact(self): if _declares_prefer_model_sampler(p.read_text(encoding="utf-8", errors="ignore")) } expected = { + "audio8_tts", "cosyvoice3", "glm_tts", "higgs_audio_v2", diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 5586b956e32..b711f0cfb36 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -39,6 +39,7 @@ AUDEX_TTA_PIPELINE, AUDEX_TTS_PIPELINE, ) +from vllm_omni.model_executor.models.audio8_tts.pipeline import AUDIO8_TTS_PIPELINE from vllm_omni.model_executor.models.aura_omni.pipeline import AURA_OMNI_PIPELINE from vllm_omni.model_executor.models.bagel.pipeline import ( BAGEL_PIPELINE, @@ -172,6 +173,7 @@ "voxtral_tts": VOXTRAL_TTS_PIPELINE, "glm_tts": GLM_TTS_PIPELINE, "fish_qwen3_omni": FISH_SPEECH_PIPELINE, + "arktts": AUDIO8_TTS_PIPELINE, "ming_flash_omni": MING_FLASH_OMNI_PIPELINE, "ming_flash_omni_tts": MING_FLASH_OMNI_TTS_PIPELINE, "ming_flash_omni_thinker_only": MING_FLASH_OMNI_THINKER_ONLY_PIPELINE, diff --git a/vllm_omni/deploy/audio8_tts.yaml b/vllm_omni/deploy/audio8_tts.yaml new file mode 100644 index 00000000000..644d77745f5 --- /dev/null +++ b/vllm_omni/deploy/audio8_tts.yaml @@ -0,0 +1,80 @@ +# Audio8 TTS Preview 0.6B deploy: async-chunk streaming slow_ar -> codec_decoder. +# Model: Audio8/Audio8-TTS-Preview-0.6b (model_type = "arktts"). +# +# trust_remote_code is deliberately OFF: the checkpoint ships an `auto_map`, and +# transformers prefers remote code over registered classes when it is enabled, +# which would bypass vllm-omni's Audio8TTSConfig (the Qwen2-shaped view that +# Qwen2Model needs). The tokenizer is a plain PreTrainedTokenizerFast and needs +# no remote code either. +async_chunk: true +enable_chunked_prefill: false + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + codec_streaming: true + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + # ~21.5 Hz codec; 25 frames ~= 1.16 s of audio. + codec_chunk_frames: 25 + codec_left_context_frames: 25 + initial_codec_chunk_frames: 4 + # float32 keeps the Snake activations / weight-normed convs clean. + audio8_tts_codec_dtype: float32 + +stages: + - stage_id: 0 + # >= 4 so concurrent requests pipeline; a default of 1 starves the codec + # window and produces audible gaps under load. + # + # Measured on 1x H20 (20 prompts, mean audio 4.5 s): at client concurrency 8, + # raising this to 8 buys +27% throughput and 4.5x lower first-packet latency + # (1124 ms -> 247 ms) at the cost of a 0.22 s worst-case buffer deficit in + # 1 of 20 requests. 4 is the default because every measured point here had + # zero underrun; see examples/online_serving/text_to_speech/README.md. + max_num_seqs: 4 + # 0.6B weights + a 2048-token KV cache need very little: sized for a small + # slice of the GPU so the codec stage and other tenants still fit. + gpu_memory_utilization: 0.2 + enforce_eager: false + trust_remote_code: false + enable_prefix_caching: false + async_scheduling: false + # The checkpoint's max_seq_len is 2048 packed text+audio positions; with + # chunked prefill off, max_num_batched_tokens must be >= max_model_len. + max_model_len: 2048 + max_num_batched_tokens: 2048 + devices: "0" + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + # generation_config.json of the released checkpoint. + temperature: 0.7 + top_k: 50 + top_p: 0.9 + max_tokens: 1024 + seed: 42 + repetition_penalty: 1.0 + + - stage_id: 1 + max_num_seqs: 1 + # The codec is ~1.4 GB of weights plus per-chunk activations. + gpu_memory_utilization: 0.1 + enforce_eager: true + trust_remote_code: false + enable_prefix_caching: false + async_scheduling: false + max_model_len: 2048 + max_num_batched_tokens: 2048 + devices: "0" + input_connectors: + from_stage_0: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + repetition_penalty: 1.0 diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 75a8a435e94..88d9f090d39 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -712,6 +712,7 @@ def _apply_uploaded_speaker(self, request: OpenAICreateSpeechRequest) -> str | N if self._tts_model_type in ( "cosyvoice3", "fish_tts", + "audio8_tts", "omnivoice", "moss_tts_nano", "glm_tts", @@ -721,6 +722,7 @@ def _apply_uploaded_speaker(self, request: OpenAICreateSpeechRequest) -> str | N label = { "cosyvoice3": "CosyVoice3", "fish_tts": "Fish Speech", + "audio8_tts": "Audio8 TTS", "omnivoice": "OmniVoice", "moss_tts_nano": "MOSS-TTS-Nano", "higgs_audio_v2": "Higgs-Audio V2", diff --git a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py index a09d97b6875..39882274df4 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py @@ -104,6 +104,7 @@ def tts_entry_stage_archs() -> frozenset[str]: from vllm_omni.entrypoints.openai.tts_adapters import ( # noqa: E402,F401 audex, audex_tta, + audio8_tts, cosyvoice3, covo_audio, dots_tts, diff --git a/vllm_omni/entrypoints/openai/tts_adapters/audio8_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/audio8_tts.py new file mode 100644 index 00000000000..1706a64d3fd --- /dev/null +++ b/vllm_omni/entrypoints/openai/tts_adapters/audio8_tts.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS Preview serving adapter for ``/v1/audio/speech``.""" + +import copy +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch +from vllm.inputs import tokens_input +from vllm.utils.async_utils import make_async + +from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter +from vllm_omni.entrypoints.openai.tts_adapters.base import ( + ARTTSAdapter, + PreparedRequest, + apply_max_new_tokens, + conditioning_cache_salt, +) +from vllm_omni.model_executor.models.audio8_tts.codec_utils import ( + estimate_reference_code_frames, +) +from vllm_omni.model_executor.models.audio8_tts.prompt_utils import ( + build_text_only_prompt_ids, + estimate_voice_clone_prompt_len, + normalize_text, +) + +if TYPE_CHECKING: + from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest + + +@register_tts_adapter +class Audio8TTSAdapter(ARTTSAdapter): + """Audio8 TTS Preview 0.6B: text-only synthesis and zero-shot voice cloning.""" + + stage_keys = frozenset({"audio8_tts_slow_ar"}) + name = "audio8_tts" + + def __init__(self, ctx: Any) -> None: + super().__init__(ctx) + self._cached_tokenizer: Any = None + + def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: + server = self.ctx.server + err = server._apply_uploaded_speaker(request) + if err: + return err + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + + if request.ref_audio is not None: + fmt_err = server._validate_ref_audio_format(request.ref_audio) + if fmt_err: + return fmt_err + if not request.ref_text or not request.ref_text.strip(): + return "Voice cloning requires 'ref_text' (transcript of the reference audio)" + + 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 cannot exceed {self.max_new_tokens_max}" + + return None + + def apply_sampling_overrides( + self, + sampling_params_list: list, + request: "OpenAICreateSpeechRequest", + prompt: dict[str, Any] | None = None, + request_id: str | None = None, + ) -> list: + sampling_params_list = apply_max_new_tokens(sampling_params_list, request) + # Seed the Fast AR (talker_mtp) residual sampling per request. The Slow AR + # is already seeded via sampling_metadata.generators, but the residual + # codebooks read tts_local_seed; without it they fall back to global RNG, + # so an unseeded request's audio would depend on batch composition. Default + # it from the deploy seed; an explicit request.seed still overrides it + # downstream. Mirrors how serving handles qwen3_tts. + if sampling_params_list: + default_seed = sampling_params_list[0].seed + if default_seed is not None: + sampling_params_list = copy.deepcopy(sampling_params_list) + stage0 = sampling_params_list[0] + if stage0.extra_args is None: + stage0.extra_args = {} + stage0.extra_args.setdefault("tts_local_seed", int(default_seed)) + return sampling_params_list + + async def build( + self, + request: "OpenAICreateSpeechRequest", + sampling_params_list: list, + has_inline_ref_audio: bool, + ) -> PreparedRequest: + del sampling_params_list + server = self.ctx.server + ref_audio_data = None + if request.ref_audio is not None: + wav_list, sample_rate = await server._resolve_ref_audio(request.ref_audio) + ref_audio_data = (wav_list, sample_rate) + # Prompt building tokenizes and, for voice clone, allocates tensors; keep + # it off the event loop via the server's single-worker TTS executor. + build_prompt = make_async(self._build_prompt, executor=server._tts_executor) + prompt = await build_prompt( + request, + ref_audio_data=ref_audio_data, + has_inline_ref_audio=has_inline_ref_audio, + ) + tts_params: dict = {} + prompt["cache_salt"] = conditioning_cache_salt(request, tts_params) + return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type=self.name) + + def _tokenizer(self) -> Any: + """The Audio8 tokenizer, loaded once per process.""" + if self._cached_tokenizer is None: + from transformers import AutoTokenizer + + model_name = self.ctx.server.engine_client.model_config.model + self._cached_tokenizer = AutoTokenizer.from_pretrained(model_name) + return self._cached_tokenizer + + def _estimate_prompt_len(self, text: str, ref_text: str, ref_audio: object) -> int: + """Exact clone-prompt length, without encoding the reference audio. + + Deliberately not defensive: the placeholder this sizes must match the + length ``preprocess()`` actually builds, so a wrong guess corrupts the + request instead of failing it. Let errors surface to the caller. + """ + if not isinstance(ref_audio, (list, tuple)) or len(ref_audio) != 2: + raise ValueError("Audio8 TTS reference audio must be a (samples, sample_rate) pair") + wav, sample_rate = ref_audio + ref_frames = estimate_reference_code_frames(len(wav), int(sample_rate)) + return estimate_voice_clone_prompt_len(self._tokenizer(), text, ref_text, ref_frames) + + def _build_prompt( + self, + request: "OpenAICreateSpeechRequest", + ref_audio_data: tuple[list[float], int] | None = None, + *, + has_inline_ref_audio: bool = False, + ) -> dict[str, Any]: + """Build the engine prompt for Audio8 TTS Preview. + + Text-only prompts are tokenized here. Voice cloning cannot be: the + prompt embeds the reference audio's own codec codes, so the model-side + ``preprocess`` builds it and this path only reserves a placeholder of the + exact final length. + """ + server = self.ctx.server + if ref_audio_data is None or not request.ref_text: + prompt_ids, normalized_text = build_text_only_prompt_ids(self._tokenizer(), request.input) + # Scalars are list-wrapped for the text-only path, matching the + # other TTS entrypoints' additional_information shape. + additional_information: dict[str, Any] = {"text": [normalized_text]} + if request.max_new_tokens is not None: + additional_information["max_new_tokens"] = [request.max_new_tokens] + prompt = tokens_input(prompt_token_ids=prompt_ids) + prompt["additional_information"] = additional_information + return prompt + + wav_samples, sample_rate = ref_audio_data + normalized_text = normalize_text(request.input) + normalized_ref_text = normalize_text(request.ref_text, add_default_speaker=True) + placeholder_len = self._estimate_prompt_len(normalized_text, normalized_ref_text, ref_audio_data) + + # Structured clone: scalars (not list-wrapped) because model-side + # preprocess() consumes these fields directly. + additional_information = { + "text": normalized_text, + "ref_text": normalized_ref_text, + "ref_audio_wav": torch.from_numpy(np.asarray(wav_samples, dtype=np.float32)), + "ref_audio_sr": int(sample_rate), + "audio8_structured_voice_clone": True, + } + if request.voice is not None: + voice_lower = request.voice.lower() + if voice_lower in server.uploaded_speakers and not has_inline_ref_audio: + # Lets the model cache this speaker's encoded codes. + additional_information["voice_name"] = voice_lower + additional_information["voice_created_at"] = server._voice_created_at(voice_lower) + if request.max_new_tokens is not None: + additional_information["max_new_tokens"] = request.max_new_tokens + prompt = tokens_input(prompt_token_ids=[1] * placeholder_len) + prompt["additional_information"] = additional_information + return prompt + + +__all__ = ["Audio8TTSAdapter"] diff --git a/vllm_omni/model_executor/models/audio8_tts/__init__.py b/vllm_omni/model_executor/models/audio8_tts/__init__.py new file mode 100644 index 00000000000..a6db0478028 --- /dev/null +++ b/vllm_omni/model_executor/models/audio8_tts/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS Preview (DualAR, 44.1 kHz codec) model implementation.""" diff --git a/vllm_omni/model_executor/models/audio8_tts/audio8_tts_codec_decoder.py b/vllm_omni/model_executor/models/audio8_tts/audio8_tts_codec_decoder.py new file mode 100644 index 00000000000..6f93d31a260 --- /dev/null +++ b/vllm_omni/model_executor/models/audio8_tts/audio8_tts_codec_decoder.py @@ -0,0 +1,284 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS Preview -- codec decoder (Stage 1). + +Consumes frames of ``num_codebooks`` codec codes and decodes them to a 44.1 kHz +waveform. Streaming contract is delta: with ``async_chunk`` on, each chunk is +``[left_context | new_frames]``; the whole window is decoded so the causal +codec keeps continuous context, then only the new tail is returned. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.audio8_tts.codec_utils import load_arktts_codec +from vllm_omni.model_executor.models.audio8_tts.configuration_audio8_tts import ( + ARKTTS_CODEC_FRAME_SIZE, + ARKTTS_CODEC_SAMPLE_RATE, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + +logger = init_logger(__name__) + +_DTYPES = { + "fp16": torch.float16, + "float16": torch.float16, + "half": torch.float16, + "bf16": torch.bfloat16, + "bfloat16": torch.bfloat16, + "fp32": torch.float32, + "float32": torch.float32, + "": torch.float32, +} + + +def _connector_extra_config(vllm_config: VllmConfig) -> dict[str, Any]: + model_config = getattr(vllm_config, "model_config", None) + connector_cfg = getattr(model_config, "stage_connector_config", None) + if isinstance(connector_cfg, dict): + return connector_cfg.get("extra", connector_cfg) + extra = getattr(connector_cfg, "extra", None) + return extra if isinstance(extra, dict) else {} + + +def _resolve_dtype(extra_cfg: dict[str, Any]) -> torch.dtype: + raw = str(extra_cfg.get("audio8_tts_codec_dtype", "float32")).strip().lower() + if raw not in _DTYPES: + raise ValueError(f"Invalid Audio8 TTS codec dtype: {raw!r}") + return _DTYPES[raw] + + +class Audio8TTSCodecDecoder(nn.Module): + """Stage 1: codec codes -> waveform (runs under the generation runner).""" + + input_modalities = "audio" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + del prefix + self.vllm_config = vllm_config + self.model_path = vllm_config.model_config.model + config = vllm_config.model_config.hf_config + + self.have_multimodal_outputs = True + self.has_preprocess = False + self.has_postprocess = False + self.enable_update_additional_information = True + self.requires_raw_input_tokens = True + + self._num_codebooks = int(getattr(config, "num_codebooks", 10)) + self._sample_rate = int(getattr(config, "codec_sample_rate", ARKTTS_CODEC_SAMPLE_RATE)) + self._frame_size = int(getattr(config, "codec_frame_size", ARKTTS_CODEC_FRAME_SIZE)) + self._codec_kwargs = { + "post_n_layer": int(getattr(config, "codec_post_n_layer", 8)), + "post_n_head": int(getattr(config, "codec_post_n_head", 16)), + "post_n_local_heads": int(getattr(config, "codec_post_n_local_heads", 8)), + "post_intermediate_size": int(getattr(config, "codec_post_intermediate_size", 1216)), + } + # float32 by default: the codec's Snake activations and weight-normed + # convs lose audible headroom in bf16. + self._codec_dtype = _resolve_dtype(_connector_extra_config(vllm_config)) + self._codec: nn.Module | None = None + self._codec_device: torch.device | None = None + self._logged_stats = False + + # -------------------- codec -------------------- + + def _ensure_codec_loaded(self) -> None: + if self._codec is not None: + return + device = self.vllm_config.device_config.device + self._codec = load_arktts_codec( + self.model_path, + device=device, + dtype=self._codec_dtype, + role="decode", + **self._codec_kwargs, + ) + self._codec_device = torch.device(device) + + # -------------------- vLLM hooks -------------------- + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + if input_ids.numel() == 0: + return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32) + return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32) + + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> None: + return None + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """No-op: the codec comes from ``codec.pth``, not the main checkpoint.""" + del weights + return set() + + # -------------------- request splitting -------------------- + + def _split_request_ids( + self, + ids: torch.Tensor, + seq_token_counts: list[int] | None = None, + ) -> list[torch.Tensor]: + if seq_token_counts is not None and len(seq_token_counts) > 1: + boundaries = [0] + for count in seq_token_counts: + boundaries.append(boundaries[-1] + count) + total = ids.numel() + return [ids[boundaries[i] : min(boundaries[i + 1], total)] for i in range(len(seq_token_counts))] + if is_forward_context_available(): + slices = get_forward_context().ubatch_slices + if slices is not None and len(slices) > 1 and not any(hasattr(s, "token_slice") for s in slices): + boundaries = [0] + for size in slices: + boundaries.append(boundaries[-1] + size) + return [ids[boundaries[i] : boundaries[i + 1]] for i in range(len(boundaries) - 1)] + return [ids] + + def _codes_from_runtime_info( + self, + info: dict[str, Any] | None, + device: torch.device, + ) -> torch.Tensor | None: + """Prefer the tensor payload the connector carried over input_ids.""" + if not isinstance(info, dict): + return None + codes = info.get("codes", {}).get("audio") + if not isinstance(codes, torch.Tensor) or codes.numel() == 0: + return None + codes = codes.to(device=device, dtype=torch.long, non_blocking=True).contiguous() + if codes.ndim != 2 or int(codes.shape[0]) != self._num_codebooks: + logger.warning( + "Audio8 TTS codec decoder expected codes of shape [%d, frames], got %s", + self._num_codebooks, + tuple(codes.shape), + ) + return None + return codes + + # -------------------- decode -------------------- + + @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 one chunk per request. + + ``input_ids`` is codebook-major flat codes + (``[cb0_f0 ... cb0_fN, cb1_f0 ...]``); tensor codes on the connector + win over ``input_ids``. Every return path must set ``model_outputs`` or + the serving layer silently drops a chunk of audio. + """ + del positions, intermediate_tensors, inputs_embeds + sr_tensor = torch.tensor(self._sample_rate, dtype=torch.int32) + empty = torch.zeros((0,), dtype=torch.float32) + + self._ensure_codec_loaded() + codec = self._codec + assert codec is not None + + if input_ids is None or input_ids.numel() == 0: + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": [empty], "sr": [sr_tensor]}, + ) + + ids = input_ids.reshape(-1).to(dtype=torch.long) + per_request_ids = self._split_request_ids(ids, kwargs.get("seq_token_counts")) + num_req = len(per_request_ids) + + context_frames = [0] * num_req + if runtime_additional_information is not None: + for index, info in enumerate(runtime_additional_information[:num_req]): + meta = info.get("meta", {}) if isinstance(info, dict) else {} + context_frames[index] = int(meta.get("left_context_size", 0) or 0) + + audios: list[torch.Tensor] = [empty] * num_req + for index, request_ids in enumerate(per_request_ids): + codes_qf = ( + self._codes_from_runtime_info(runtime_additional_information[index], ids.device) + if runtime_additional_information is not None and index < len(runtime_additional_information) + else None + ) + if codes_qf is None: + if request_ids.numel() < self._num_codebooks: + continue + if int(request_ids.numel()) % self._num_codebooks != 0: + logger.warning( + "Audio8 TTS codec decoder got %d ids, not divisible by num_codebooks=%d; skipping request", + int(request_ids.numel()), + self._num_codebooks, + ) + continue + codes_qf = request_ids.reshape(self._num_codebooks, -1) + + total_frames = int(codes_qf.shape[1]) + if total_frames <= 0: + continue + new_frames = total_frames - context_frames[index] + if new_frames <= 0: + logger.warning( + "Audio8 TTS codec decoder chunk has no new frames (total=%d, context=%d)", + total_frames, + context_frames[index], + ) + continue + + if not self._logged_stats: + self._logged_stats = True + logger.info( + "Audio8 TTS codec decoder: frames=%d codebooks=%d dtype=%s sr=%d", + total_frames, + self._num_codebooks, + self._codec_dtype, + self._sample_rate, + ) + + with torch.amp.autocast("cuda", enabled=False): + # Move codes to the codec's device explicitly: a mismatch here + # silently falls back to CPU or raises deep inside the decoder. + waveform = codec.decode(codes_qf.unsqueeze(0).to(device=self._codec_device)) + waveform = waveform.reshape(-1).to(dtype=torch.float32) + + # Trim the left-context prefix proportionally: the codec's padding / + # rounding means the decoded length is not exactly + # frames * frame_size. + if context_frames[index] > 0: + cut = int(round(context_frames[index] / total_frames * int(waveform.shape[0]))) + cut = max(0, min(cut, int(waveform.shape[0]))) + waveform = waveform[cut:] + if waveform.numel() > 0: + audios[index] = waveform.contiguous() + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": audios, "sr": [sr_tensor] * num_req}, + ) + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + del kwargs + if isinstance(model_outputs, OmniOutput): + return model_outputs + if not (isinstance(model_outputs, tuple) and len(model_outputs) == 2): + raise TypeError(f"Audio8TTSCodecDecoder expected (audio_tensor, sr), got {type(model_outputs)}") + audio_tensor, sample_rate = model_outputs + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": audio_tensor, "sr": sample_rate}, + ) + + +__all__ = ["Audio8TTSCodecDecoder"] diff --git a/vllm_omni/model_executor/models/audio8_tts/audio8_tts_fast_ar.py b/vllm_omni/model_executor/models/audio8_tts/audio8_tts_fast_ar.py new file mode 100644 index 00000000000..5b4ff969604 --- /dev/null +++ b/vllm_omni/model_executor/models/audio8_tts/audio8_tts_fast_ar.py @@ -0,0 +1,346 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS Fast AR: the residual-codebook predictor. + +Runs once per Slow AR step over ``num_codebooks + 1`` positions: position 0 is +the projected Slow AR hidden state (primes the KV cache, logits discarded), +position 1 is the sampled semantic code (not predicted, only fed in), positions +2.. are the previously sampled residual codes; only codes 1..num_codebooks-1 +are sampled here. + +Not paged: at most 10 tokens per call, SDPA over a per-call pre-allocated KV +cache. Matches the reference ``ArkttsModel._generate_codebooks``. +""" + +from __future__ import annotations + +from collections.abc import Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.config import VllmConfig +from vllm.config.vllm import set_current_vllm_config +from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from .configuration_audio8_tts import Audio8TTSFastARConfig, Audio8TTSSlowARConfig +from .sampling import sample_scores + +logger = init_logger(__name__) + + +class _FastARAttention(nn.Module): + """GQA self-attention over the short codebook sequence. + + ``forward_one`` decodes a single position against a caller-owned KV cache; + there is no batched-prefill path because the codebook loop is inherently + sequential. + """ + + def __init__(self, config: Audio8TTSFastARConfig, *, prefix: str = "") -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self._use_gqa = self.num_kv_heads != self.num_heads + + self.qkv_proj = QKVParallelLinear( + hidden_size=self.hidden_size, + head_size=self.head_dim, + total_num_heads=self.num_heads, + total_num_kv_heads=self.num_kv_heads, + bias=config.attention_qkv_bias, + prefix=f"{prefix}.qkv_proj", + disable_tp=True, + ) + self.o_proj = RowParallelLinear( + input_size=self.num_heads * self.head_dim, + output_size=self.hidden_size, + bias=False, + prefix=f"{prefix}.o_proj", + disable_tp=True, + ) + self.rotary_emb = get_rope( + self.head_dim, + max_position=config.max_position_embeddings, + # Audio8 TTS uses interleaved (GPT-J) RoPE, not NeoX. + is_neox_style=False, + rope_parameters={"rope_theta": config.rope_theta, "rope_type": "default"}, + ) + if config.attention_qk_norm: + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + else: + self.q_norm = None + self.k_norm = None + + def forward_one( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + cache_pos: int, + ) -> torch.Tensor: + bsz = int(hidden_states.shape[0]) + + qkv, _ = self.qkv_proj(hidden_states.reshape(bsz, -1)) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + if self.q_norm is not None: + q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view(q.shape) + k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view(k.shape) + + q, k = self.rotary_emb(position_ids.reshape(-1), q, k) + + q = q.view(bsz, self.num_heads, self.head_dim).unsqueeze(2) + k_cache[:bsz, :, cache_pos, :] = k.view(bsz, self.num_kv_heads, self.head_dim) + v_cache[:bsz, :, cache_pos, :] = v.view(bsz, self.num_kv_heads, self.head_dim) + + attn_out = F.scaled_dot_product_attention( + q, + k_cache[:bsz, :, : cache_pos + 1, :], + v_cache[:bsz, :, : cache_pos + 1, :], + scale=self.scaling, + is_causal=False, + enable_gqa=self._use_gqa, + ) + attn_out = attn_out.transpose(1, 2).reshape(bsz, -1) + output, _ = self.o_proj(attn_out) + return output + + +class _FastARMLP(nn.Module): + """SiLU-gated MLP (``w1``/``w3`` gate-up, ``w2`` down).""" + + def __init__(self, config: Audio8TTSFastARConfig, *, prefix: str = "") -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + input_size=config.hidden_size, + output_sizes=[config.intermediate_size] * 2, + bias=False, + prefix=f"{prefix}.gate_up_proj", + disable_tp=True, + ) + self.down_proj = RowParallelLinear( + input_size=config.intermediate_size, + output_size=config.hidden_size, + bias=False, + prefix=f"{prefix}.down_proj", + disable_tp=True, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + gate, up = gate_up.chunk(2, dim=-1) + x, _ = self.down_proj(F.silu(gate) * up) + return x + + +class _FastARDecoderLayer(nn.Module): + def __init__(self, config: Audio8TTSFastARConfig, *, prefix: str = "") -> None: + super().__init__() + self.self_attn = _FastARAttention(config, prefix=f"{prefix}.self_attn") + self.mlp = _FastARMLP(config, prefix=f"{prefix}.mlp") + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward_one( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + cache_pos: int, + ) -> torch.Tensor: + hidden_states = hidden_states + self.self_attn.forward_one( + self.input_layernorm(hidden_states), position_ids, k_cache, v_cache, cache_pos + ) + return hidden_states + self.mlp(self.post_attention_layernorm(hidden_states)) + + +class Audio8TTSFastAR(nn.Module): + """Residual-codebook predictor for Audio8 TTS. + + ``forward`` returns all ``num_codebooks`` codes for one frame, with code 0 + copied from the already-sampled semantic token. + """ + + def __init__( + self, + *, + vllm_config: VllmConfig, + config: Audio8TTSFastARConfig, + slow_ar_config: Audio8TTSSlowARConfig, + prefix: str = "fast_ar", + ) -> None: + super().__init__() + self._vllm_config = vllm_config + self.config = config + self.slow_ar_config = slow_ar_config + + self.layers = nn.ModuleList( + [_FastARDecoderLayer(config, prefix=f"{prefix}.layers.{i}") for i in range(config.num_hidden_layers)] + ) + self.fast_embeddings = nn.Embedding(config.vocab_size, config.hidden_size) + self.fast_output = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.fast_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if slow_ar_config.hidden_size != config.hidden_size: + self.fast_project_in = nn.Linear(slow_ar_config.hidden_size, config.hidden_size, bias=True) + else: + # 0.6b: dim == fast_dim == 896, and the checkpoint has no + # fast_project_in tensor. + self.fast_project_in = nn.Identity() + + self._num_codebooks = int(config.num_codebooks) + self._fast_dim = int(config.hidden_size) + self._k_cache: torch.Tensor | None = None + self._v_cache: torch.Tensor | None = None + self._pos_ids: torch.Tensor | None = None + + def _ensure_buffers(self, bsz: int, device: torch.device, dtype: torch.dtype) -> None: + max_seq = self._num_codebooks + 1 + if ( + self._k_cache is not None + and self._pos_ids is not None + and self._k_cache.shape[1] >= bsz + and self._k_cache.device == device + and self._k_cache.dtype == dtype + and self._pos_ids.shape[0] >= bsz + ): + return + self._k_cache = torch.empty( + self.config.num_hidden_layers, + bsz, + self.config.num_key_value_heads, + max_seq, + self.config.head_dim, + dtype=dtype, + device=device, + ) + self._v_cache = torch.empty_like(self._k_cache) + self._pos_ids = torch.arange(max_seq, dtype=torch.long, device=device).unsqueeze(0).expand(bsz, -1).contiguous() + + def _forward_one(self, embed: torch.Tensor, position_ids: torch.Tensor, cache_pos: int) -> torch.Tensor: + assert self._k_cache is not None and self._v_cache is not None + hidden = embed + for layer_idx, layer in enumerate(self.layers): + hidden = layer.forward_one( + hidden, + position_ids, + self._k_cache[layer_idx], + self._v_cache[layer_idx], + cache_pos, + ) + return hidden + + @torch.inference_mode() + def forward( + self, + slow_ar_hidden: torch.Tensor, + semantic_token_id: torch.Tensor, + *, + do_sample: bool = True, + temperature: float = 0.7, + top_k: int = 50, + top_p: float = 0.9, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + """Predict one frame of codec codes. + + ``slow_ar_hidden`` must already be post-final-norm when + ``norm_fastlayer_input`` is set (it is for the released checkpoint) -- + vLLM's Qwen2Model returns exactly that. Returns ``[B, num_codebooks]`` + int64 codes; code 0 is the semantic code, codes 1.. come from this module. + """ + bsz = int(slow_ar_hidden.shape[0]) + num_cb = self._num_codebooks + device = slow_ar_hidden.device + dtype = slow_ar_hidden.dtype + + semantic_begin = int(self.slow_ar_config.semantic_begin_id) + codebook_size = int(self.slow_ar_config.codebook_size) + # Non-semantic tokens (EOS) clamp to code 0; those frames are dropped + # downstream because the request finishes on that token. + semantic_code = (semantic_token_id.reshape(bsz) - semantic_begin).clamp(min=0, max=codebook_size - 1) + + codes = torch.empty(bsz, num_cb, dtype=torch.long, device=device) + codes[:, 0] = semantic_code + + self._ensure_buffers(bsz, device, dtype) + pos_ids = self._pos_ids + assert pos_ids is not None + + # Position 0 primes the cache; its logits are unused. + projected = self.fast_project_in(slow_ar_hidden.reshape(bsz, -1)) + self._forward_one(projected, pos_ids[:bsz, 0], 0) + + embed = self.fast_embeddings(semantic_code) + for step in range(1, num_cb): + hidden = self._forward_one(embed, pos_ids[:bsz, step], step) + logits = self.fast_output(self.fast_norm(hidden)).float() + next_ids = sample_scores( + logits, + temperature=temperature, + top_k=top_k, + top_p=top_p, + do_sample=do_sample, + generator=generator, + ) + codes[:, step] = next_ids + if step < num_cb - 1: + embed = self.fast_embeddings(next_ids) + + return codes + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load already-remapped Fast AR weights (see the Slow AR remapper).""" + with set_current_vllm_config(self._vllm_config): + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded: set[str] = set() + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + for name, loaded_weight in weights: + handled = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + mapped = name.replace(weight_name, param_name) + if mapped not in params_dict: + continue + param = params_dict[mapped] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if weight_loader == default_weight_loader: + weight_loader(param, loaded_weight) + else: + weight_loader(param, loaded_weight, shard_id) + loaded.add(mapped) + handled = True + break + if handled or name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded.add(name) + return loaded + + +__all__ = ["Audio8TTSFastAR"] diff --git a/vllm_omni/model_executor/models/audio8_tts/audio8_tts_slow_ar.py b/vllm_omni/model_executor/models/audio8_tts/audio8_tts_slow_ar.py new file mode 100644 index 00000000000..9d2cd884e30 --- /dev/null +++ b/vllm_omni/model_executor/models/audio8_tts/audio8_tts_slow_ar.py @@ -0,0 +1,755 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS Preview -- Slow AR model (Stage 0). + +Text (+ optional reference voice) -> one semantic token per codec frame; the +frame's remaining 9 codebooks are produced by the nested Fast AR. + +Backbone is vLLM's ``Qwen2Model`` (qkv bias, no q/k norm), with three +adjustments: interleaved (GPT-J) RoPE, multi-codebook input embedding at +semantic positions (summed, unlike Fish Speech's ``1/sqrt(Q+1)`` rescale), and +logits masked to the semantic range + ``<|im_end|>`` then Repetition-Aware +sampled. + +Streaming contract is delta: each decode step appends one frame of +``num_codebooks`` codes to ``codes.audio``. +""" + +from __future__ import annotations + +import copy +from collections.abc import Iterable +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +from transformers import AutoTokenizer +from vllm.config import VllmConfig +from vllm.distributed import get_pp_group +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.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.qwen2 import Qwen2Model +from vllm.model_executor.models.utils import PPMissingLayer, maybe_prefix +from vllm.sequence import IntermediateTensors +from vllm.v1.outputs import SamplerOutput + +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.utils.speaker_cache import get_speaker_cache + +from .audio8_tts_fast_ar import Audio8TTSFastAR +from .codec_utils import encode_reference_audio_codes +from .configuration_audio8_tts import ( + Audio8TTSConfig, + Audio8TTSFastARConfig, + Audio8TTSSlowARConfig, +) +from .prompt_utils import build_voice_clone_prompt_ids +from .sampling import SAMPLING_EPS, ras_sample_batch + +logger = init_logger(__name__) + +#: Defaults from ``generation_config.json`` of the released checkpoint. +DEFAULT_TEMPERATURE = 0.7 +DEFAULT_TOP_K = 50 +DEFAULT_TOP_P = 0.9 + + +def _remap_audio8_tts_weights( + weights: Iterable[tuple[str, torch.Tensor]], + *, + q_size: int, + kv_size: int, + fast_q_size: int, + fast_kv_size: int, +) -> Iterable[tuple[str, torch.Tensor]]: + """Rename Audio8 checkpoint tensors to vLLM / Qwen2 names. + + The checkpoint is flat (no ``text_model.`` prefix). Slow AR tensors move + under ``model.``, Fast AR tensors under ``fast_ar.``, and ``wqkv`` splits + into ``{q,k,v}_proj``. See the body for the exact mapping table. + """ + + def rewrite_block(suffix: str) -> str: + for old, new in ( + (".attention.wo.", ".self_attn.o_proj."), + (".attention.q_norm.", ".self_attn.q_norm."), + (".attention.k_norm.", ".self_attn.k_norm."), + (".attention_norm.", ".input_layernorm."), + (".feed_forward.w1.", ".mlp.gate_proj."), + (".feed_forward.w3.", ".mlp.up_proj."), + (".feed_forward.w2.", ".mlp.down_proj."), + (".ffn_norm.", ".post_attention_layernorm."), + ): + suffix = suffix.replace(old, new) + return suffix + + for name, tensor in weights: + # RoPE tables are recomputed by vLLM. + if name.endswith(("freqs_cis", "fast_freqs_cis")) or "rotary_emb.inv_freq" in name: + continue + + if name.startswith("fast_layers."): + suffix = name[len("fast_layers.") :] + if ".attention.wqkv." in suffix: + layer_prefix, _, param = suffix.partition(".attention.wqkv.") + yield f"fast_ar.layers.{layer_prefix}.self_attn.q_proj.{param}", tensor[:fast_q_size] + yield ( + f"fast_ar.layers.{layer_prefix}.self_attn.k_proj.{param}", + tensor[fast_q_size : fast_q_size + fast_kv_size], + ) + yield ( + f"fast_ar.layers.{layer_prefix}.self_attn.v_proj.{param}", + tensor[fast_q_size + fast_kv_size :], + ) + continue + yield f"fast_ar.layers.{rewrite_block(suffix)}", tensor + continue + + if name.startswith(("fast_embeddings.", "fast_output.", "fast_norm.", "fast_project_in.")): + yield f"fast_ar.{name}", tensor + continue + + if name.startswith("layers."): + suffix = name[len("layers.") :] + if ".attention.wqkv." in suffix: + layer_prefix, _, param = suffix.partition(".attention.wqkv.") + yield f"model.layers.{layer_prefix}.self_attn.q_proj.{param}", tensor[:q_size] + yield ( + f"model.layers.{layer_prefix}.self_attn.k_proj.{param}", + tensor[q_size : q_size + kv_size], + ) + yield f"model.layers.{layer_prefix}.self_attn.v_proj.{param}", tensor[q_size + kv_size :] + continue + yield f"model.layers.{rewrite_block(suffix)}", tensor + continue + + if name == "embeddings.weight": + yield "model.embed_tokens.weight", tensor + continue + if name == "norm.weight": + yield "model.norm.weight", tensor + continue + + yield name, tensor + + +class Audio8TTSSlowARForConditionalGeneration(nn.Module): + """Stage 0: text -> semantic tokens + residual codec codes.""" + + prefer_model_sampler = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.vllm_config = vllm_config + self.model_path = vllm_config.model_config.model + + config: Audio8TTSConfig = vllm_config.model_config.hf_config # type: ignore[assignment] + self.config = config + self.text_config: Audio8TTSSlowARConfig = config.get_text_config() + self.fast_ar_config: Audio8TTSFastARConfig = config.fast_ar_config + + self._semantic_begin_id = int(self.text_config.semantic_begin_id) + self._semantic_end_id = int(self.text_config.semantic_end_id) + self._num_semantic_ids = self._semantic_end_id - self._semantic_begin_id + 1 + self._eos_token_id = int(self.text_config.eos_token_id) + self._pad_token_id = int(self.text_config.pad_token_id) + self._codebook_size = int(self.text_config.codebook_size) + self._num_codebooks = int(self.text_config.num_codebooks) + self._ras_window_size = int(config.ras_window_size) + self._ras_temperature = float(config.ras_temperature) + self._ras_top_p = float(config.ras_top_p) + + self.have_multimodal_outputs = True + self.has_preprocess = True + self.has_postprocess = True + self.mtp_hidden_size = int(self.text_config.hidden_size) + self.talker_mtp_output_key = ("codes", "audio") + self.gpu_resident_buffer_keys: set[tuple[str, str]] = {("hidden_states", "last")} + self.talker_mtp_graph_safe = True + + self.model = Qwen2Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")) + self._fix_rope_style() + + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + self.text_config.vocab_size, + self.text_config.hidden_size, + quant_config=vllm_config.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + self.logits_processor = LogitsProcessor(self.text_config.vocab_size) + self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors + + # Summed multi-codebook input embedding table. + self.codebook_embeddings = nn.Embedding( + self._codebook_size * self._num_codebooks, + self.text_config.hidden_size, + ) + + # Fast AR shares the parent's VllmConfig but must not share the + # compilation context: see the Fish Speech note on why this is + # copy.copy and not dataclasses.replace (pydantic validators re-run and + # reject an already-rebound compilation backend). + fast_ar_compilation = copy.copy(vllm_config.compilation_config) + fast_ar_compilation.static_forward_context = {} + self._fast_ar_vllm_config = copy.copy(vllm_config) + self._fast_ar_vllm_config.compilation_config = fast_ar_compilation + from vllm.config.vllm import set_current_vllm_config + + with set_current_vllm_config(self._fast_ar_vllm_config): + self.fast_ar = Audio8TTSFastAR( + vllm_config=self._fast_ar_vllm_config, + config=self.fast_ar_config, + slow_ar_config=self.text_config, + prefix="fast_ar", + ) + + # Constant logits mask: semantic codes plus <|im_end|>. Safe as a + # non-persistent buffer under vLLM (module built on the target device), + # unlike under ``PreTrainedModel.from_pretrained`` on transformers >= 5, + # which meta-inits and then empty_like's every non-persistent buffer. + vocab = int(self.text_config.vocab_size) + allowed = torch.zeros((vocab,), dtype=torch.bool) + allowed[self._semantic_begin_id : min(self._semantic_end_id + 1, vocab)] = True + if self._eos_token_id < vocab: + allowed[self._eos_token_id] = True + self.register_buffer("_semantic_allowed_mask", allowed, persistent=False) + + self._speaker_cache = get_speaker_cache() + self._tokenizer = None + + def _fix_rope_style(self) -> None: + """Rebuild RoPE as interleaved (GPT-J); vLLM's Qwen2 defaults to NeoX. + + This also replaces the reference implementation's ``freqs_cis`` table + entirely: ``get_rope`` computes its cos/sin cache in the constructor, so + there is no checkpoint-absent buffer for a loader to leave uninitialised. + """ + from vllm.model_executor.layers.rotary_embedding import get_rope + + rope_parameters = dict(self.text_config.rope_parameters or {}) + rope_parameters.setdefault("rope_type", "default") + for layer in self.model.layers: + attn = layer.self_attn + attn.rotary_emb = get_rope( + head_size=attn.head_dim, + max_position=self.text_config.max_position_embeddings, + is_neox_style=False, + rope_parameters=rope_parameters, + ) + logger.info("Audio8 TTS: switched %d layers to interleaved RoPE", len(self.model.layers)) + + # -------------------- vLLM hooks -------------------- + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **_: Any, + ) -> torch.Tensor | IntermediateTensors: + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) + + def compute_logits( + self, + hidden_states: torch.Tensor | OmniOutput, + sampling_metadata: Any = None, + ) -> torch.Tensor | None: + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if hidden_states is None: + return None + logits = self.logits_processor(self.lm_head, hidden_states) + if logits is None: + return None + return logits.masked_fill(~self._semantic_allowed_mask, float("-inf")) + + # -------------------- Repetition-Aware Sampling -------------------- + + def _compact_semantic_logits(self, logits: torch.Tensor) -> torch.Tensor: + """Slice the allowed vocabulary into ``[B, num_semantic_ids + 1]``. + + Local ids ``0..num_semantic_ids-1`` are semantic codes; the last column + is ``<|im_end|>``. Sorting 4097 columns instead of the full 155776 + vocabulary is the difference between a cheap and an expensive AR step, + and is exactly equivalent because everything else is already ``-inf``. + """ + semantic = logits[:, self._semantic_begin_id : self._semantic_end_id + 1] + eos = logits[:, self._eos_token_id : self._eos_token_id + 1] + return torch.cat((semantic, eos), dim=-1) + + def _local_to_token_id(self, local_ids: torch.Tensor) -> torch.Tensor: + is_eos = local_ids >= self._num_semantic_ids + return torch.where(is_eos, torch.full_like(local_ids, self._eos_token_id), local_ids + self._semantic_begin_id) + + def _recent_local_ids(self, output_token_ids: list[list[int]], num_reqs: int, device: torch.device): + """Build the ``[B, W]`` RAS window in compact id space. + + Reads the host-side decoded-token history the AR runner already + materialised, so this costs one small H2D copy and never syncs the GPU. + Returns ``None`` when no request has history yet. + """ + window = self._ras_window_size + if window <= 0: + return None + rows: list[list[int]] = [] + any_history = False + for req_idx in range(num_reqs): + history = output_token_ids[req_idx] if req_idx < len(output_token_ids) else [] + recent = [int(token) for token in history[-window:]] + local = [ + token - self._semantic_begin_id if self._semantic_begin_id <= token <= self._semantic_end_id else -1 + for token in recent + ] + any_history = any_history or bool(local) + rows.append([-1] * (window - len(local)) + local) + if not any_history: + return None + return torch.tensor(rows, dtype=torch.long, device=device) + + def sample(self, logits: torch.Tensor, sampling_metadata: Any) -> SamplerOutput | None: + """Sample the semantic token with the reference filter order + RAS. + + Returns ``None`` to fall back to vLLM's sampler when the request mix + needs features RAS does not model (logprobs, penalties, bad words). + """ + if logits is None or logits.numel() == 0: + return None + if sampling_metadata.max_num_logprobs is not None: + return None + if not sampling_metadata.no_penalties: + return None + if sampling_metadata.bad_words_token_ids: + return None + + logits = logits.to(torch.float32) + if sampling_metadata.allowed_token_ids_mask is not None: + num_reqs = int(logits.shape[0]) + logits.masked_fill_(sampling_metadata.allowed_token_ids_mask[:num_reqs], float("-inf")) + for processor in sampling_metadata.logitsprocs.non_argmax_invariant: + logits = processor.apply(logits) + + compact = self._compact_semantic_logits(logits) + num_reqs = int(compact.shape[0]) + device = compact.device + + temperature = sampling_metadata.temperature + if sampling_metadata.all_greedy or temperature is None: + local_ids = compact.argmax(dim=-1) + else: + temperature = temperature[:num_reqs] + top_p = DEFAULT_TOP_P if sampling_metadata.top_p is None else sampling_metadata.top_p[:num_reqs] + top_k = DEFAULT_TOP_K if sampling_metadata.top_k is None else sampling_metadata.top_k[:num_reqs] + recent = self._recent_local_ids(sampling_metadata.output_token_ids, num_reqs, device) + local_ids = ras_sample_batch( + compact, + recent, + temperature=temperature, + top_k=top_k, + top_p=top_p, + ras_temperature=self._ras_temperature, + ras_top_p=self._ras_top_p, + num_semantic_ids=self._num_semantic_ids, + generators=sampling_metadata.generators, + num_reqs=num_reqs, + ) + + token_ids = self._local_to_token_id(local_ids).to(dtype=torch.int32) + return SamplerOutput(sampled_token_ids=token_ids.unsqueeze(-1), logprobs_tensors=None) + + # -------------------- Omni multimodal plumbing -------------------- + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + return model_outputs + + hidden = model_outputs + info_dicts = kwargs.get("model_intermediate_buffer") + if info_dicts is None: + info_dicts = kwargs.get("runtime_additional_information") or [] + + audio_codes_list: list[torch.Tensor] = [] + for info in info_dicts: + if not isinstance(info, dict): + continue + codes = info.get("codes", {}).get("audio") + if isinstance(codes, torch.Tensor): + audio_codes_list.append(codes) + + if not audio_codes_list: + return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) + + audio_codes = torch.cat(audio_codes_list, dim=0) + hidden = hidden[: int(audio_codes.shape[0])] + return OmniOutput(text_hidden_states=hidden, multimodal_outputs={"audio_codes": audio_codes}) + + # -------------------- preprocess / postprocess -------------------- + + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + **info_dict: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + additional_information = info_dict.get("additional_information") + if isinstance(additional_information, dict): + merged: dict[str, Any] = {k: v for k, v in info_dict.items() if k != "additional_information"} + for key, value in additional_information.items(): + merged.setdefault(key, value) + info_dict = merged + + span_len = int(input_ids.shape[0]) + if span_len <= 0: + embeds = input_embeds if input_embeds is not None else self.embed_input_ids(input_ids) + return input_ids, embeds, {} + + if span_len > 1: + return self._preprocess_prefill(input_ids, info_dict, span_len) + return self._preprocess_decode(input_ids, info_dict) + + def _prefill_chunk( + self, + prompt_embeds_buf: torch.Tensor, + start: int, + span_len: int, + device: torch.device, + ) -> torch.Tensor: + """Slice ``span_len`` rows out of the pinned CPU prompt-embed buffer.""" + total = int(prompt_embeds_buf.shape[0]) + begin = max(0, min(start, total)) + end = max(0, min(start + span_len, total)) + take = prompt_embeds_buf[begin:end] + if int(take.shape[0]) < span_len: + pad_rows = span_len - int(take.shape[0]) + pad_embed = self.embed_input_ids( + torch.tensor([self._pad_token_id], device=device, dtype=torch.long) + ).reshape(1, -1) + take = torch.cat([take.to(device=device, dtype=torch.bfloat16), pad_embed.expand(pad_rows, -1)], dim=0) + return take.to(dtype=torch.bfloat16) + return take.to(device=device, dtype=torch.bfloat16, non_blocking=True) + + def _preprocess_prefill( + self, + input_ids: torch.Tensor, + info_dict: dict[str, Any], + span_len: int, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + device = input_ids.device + embed = info_dict.get("embed", {}) + prompt_embeds_buf = embed.get("prefill") + is_first_chunk = not isinstance(prompt_embeds_buf, torch.Tensor) or prompt_embeds_buf.ndim != 2 + + if is_first_chunk: + if bool(info_dict.get("audio8_structured_voice_clone", False)): + prompt_embeds = self._build_voice_clone_prefill_embeds(info_dict) + else: + prompt_embeds = self.embed_input_ids(input_ids.reshape(1, -1).to(torch.long)).squeeze(0) + prompt_embeds_buf = prompt_embeds.detach().to("cpu", dtype=torch.bfloat16).contiguous() + if not prompt_embeds_buf.is_pinned(): + prompt_embeds_buf = prompt_embeds_buf.pin_memory() + offset = 0 + else: + offset = int(info_dict.get("meta", {}).get("prefill_offset", 0) or 0) + + total_prompt_len = int(prompt_embeds_buf.shape[0]) + chunk = self._prefill_chunk(prompt_embeds_buf, offset, span_len, device) + next_offset = min(offset + span_len, total_prompt_len) + + # No codes are emitted during prefill; a zero frame per position keeps + # the codes buffer aligned with the hidden-state span. + zeros = torch.zeros((chunk.shape[0], self._num_codebooks), device=device, dtype=torch.long) + info_update = { + "embed": {"prefill": prompt_embeds_buf if next_offset < total_prompt_len else None}, + "meta": {"prefill_offset": next_offset}, + "codes": {"audio": zeros}, + } + return torch.full_like(input_ids, self._pad_token_id), chunk, info_update + + def _preprocess_decode( + self, + input_ids: torch.Tensor, + info_dict: dict[str, Any], + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + device = input_ids.device + token_embed = self.embed_input_ids(input_ids.reshape(1, 1).to(torch.long)).to( + device=device, dtype=torch.bfloat16 + ) + last_hidden = info_dict.get("hidden_states", {}).get("last") + if not isinstance(last_hidden, torch.Tensor): + # First decode step right after prefill: no Fast AR input yet. + logger.warning("Audio8 TTS preprocess: hidden_states.last missing; emitting text-only embedding") + return input_ids, token_embed.reshape(1, -1), {} + + # Codebook embeddings are added in talker_mtp, using this step's Fast AR + # output; adding them here would use the previous step's codes. + info_update = { + "mtp_inputs": ( + last_hidden.to(device=device, dtype=torch.bfloat16).reshape(1, -1), + torch.zeros(1, self.text_config.hidden_size, device=device, dtype=torch.bfloat16), + ), + } + return input_ids, token_embed.reshape(1, -1), info_update + + def postprocess(self, hidden_states: torch.Tensor, **_: Any) -> dict[str, Any]: + if hidden_states.numel() == 0: + return {} + return {"hidden_states": {"last": hidden_states[-1, :].detach().contiguous()}} + + # -------------------- voice cloning -------------------- + + def _get_tokenizer(self): + if self._tokenizer is None: + self._tokenizer = AutoTokenizer.from_pretrained(self.model_path) + return self._tokenizer + + def _codec_kwargs(self) -> dict[str, int]: + return { + "post_n_layer": int(self.config.codec_post_n_layer), + "post_n_head": int(self.config.codec_post_n_head), + "post_n_local_heads": int(self.config.codec_post_n_local_heads), + "post_intermediate_size": int(self.config.codec_post_intermediate_size), + } + + def _build_voice_clone_prefill_embeds(self, info_dict: dict[str, Any]) -> torch.Tensor: + """Encode the reference audio, build the clone prompt, embed the codes.""" + text = info_dict.get("text") + ref_text = info_dict.get("ref_text") + if not isinstance(text, str) or not isinstance(ref_text, str): + raise ValueError("Audio8 TTS voice cloning requires string 'text' and 'ref_text'") + + device = self.codebook_embeddings.weight.device + cache_key = None + voice_name = info_dict.get("voice_name") + if isinstance(voice_name, str) and voice_name: + cache_key = self._speaker_cache.make_cache_key( + voice_name, + model_type="audio8_tts", + created_at=int(info_dict.get("voice_created_at") or 0), + ) + cached = self._speaker_cache.get(cache_key) + if cached is not None: + logger.debug("Speaker cache HIT for Audio8 TTS speaker '%s'", voice_name) + return self._embed_voice_clone_prompt( + text, ref_text, cached["ref_codes_fq"].to(device=device, dtype=torch.long) + ) + + ref_audio_sr = info_dict.get("ref_audio_sr") + if not isinstance(ref_audio_sr, int): + raise ValueError("Audio8 TTS voice cloning requires integer 'ref_audio_sr'") + ref_audio_wav = info_dict.get("ref_audio_wav") + if ref_audio_wav is None: + raise ValueError("Audio8 TTS voice cloning requires 'ref_audio_wav'") + if not isinstance(ref_audio_wav, torch.Tensor): + ref_audio_wav = torch.from_numpy(np.asarray(ref_audio_wav, dtype=np.float32)) + + ref_codes_fq = encode_reference_audio_codes( + self.model_path, + ref_audio_wav, + ref_audio_sr, + device=device, + **self._codec_kwargs(), + ) + if cache_key is not None: + self._speaker_cache.put(cache_key, {"ref_codes_fq": ref_codes_fq.detach().cpu()}) + logger.debug("Speaker cache STORE for Audio8 TTS speaker '%s'", voice_name) + return self._embed_voice_clone_prompt(text, ref_text, ref_codes_fq) + + def _embed_voice_clone_prompt( + self, + text: str, + ref_text: str, + ref_codes_fq: torch.Tensor, + ) -> torch.Tensor: + """Embed the clone prompt, summing codebooks over the reference frames. + + Args: + ref_codes_fq: ``[frames, num_codebooks]`` reference codec codes. + """ + device = self.codebook_embeddings.weight.device + ref_codes_fq = ref_codes_fq.to(device=device, dtype=torch.long) + semantic_token_ids = (ref_codes_fq[:, 0] + self._semantic_begin_id).tolist() + prompt_ids, ref_start, _, _ = build_voice_clone_prompt_ids( + self._get_tokenizer(), text, ref_text, semantic_token_ids + ) + prompt = torch.tensor(prompt_ids, dtype=torch.long, device=device) + embeds = self.embed_input_ids(prompt.unsqueeze(0)).squeeze(0).to(dtype=torch.bfloat16) + + frames = int(ref_codes_fq.shape[0]) + if frames <= 0: + return embeds + codebooks = min(int(ref_codes_fq.shape[1]), self._num_codebooks) + offsets = (torch.arange(codebooks, device=device, dtype=torch.long) * self._codebook_size).unsqueeze(0) + codes = ref_codes_fq[:, :codebooks].clamp(min=0, max=self._codebook_size - 1) + offsets + codebook_sum = self.codebook_embeddings(codes).sum(dim=1).to(dtype=embeds.dtype) + + result = embeds.clone() + # Audio8 TTS sums the embeddings directly -- no 1/sqrt(Q+1) rescale. + result[ref_start : ref_start + frames] += codebook_sum + return result.to(dtype=torch.bfloat16) + + # -------------------- GPU-side Fast AR fast path -------------------- + + @torch.inference_mode() + def talker_mtp( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor, + last_talker_hidden: torch.Tensor, + text_step: torch.Tensor, + seed: int | None = None, + **kwargs: Any, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Run the Fast AR and fold its codes into this step's input embedding. + + Returns: + ``(inputs_embeds, audio_codes)`` where ``audio_codes`` is + ``[B, num_codebooks]``. + """ + del text_step + bsz = int(input_ids.shape[0]) + device = input_embeds.device + + token_ids = input_ids.reshape(bsz).to(dtype=torch.long, device=device) + past_hidden = last_talker_hidden.reshape(bsz, -1).to(dtype=torch.bfloat16, device=device) + + temperature = kwargs.get("temperature") + temperature = DEFAULT_TEMPERATURE if temperature is None else float(temperature) + do_sample = kwargs.get("do_sample") + generator = kwargs.get("generator") + if generator is None and seed is not None: + generator = torch.Generator(device=device) + generator.manual_seed(int(seed)) + + _top_k = kwargs.get("top_k") + _top_p = kwargs.get("top_p") + audio_codes = self.fast_ar( + slow_ar_hidden=past_hidden, + semantic_token_id=token_ids, + do_sample=(temperature >= SAMPLING_EPS) if do_sample is None else bool(do_sample), + temperature=temperature, + top_k=int(DEFAULT_TOP_K if _top_k is None else _top_k), + top_p=float(DEFAULT_TOP_P if _top_p is None else _top_p), + generator=generator, + ) + + embeds = input_embeds.reshape(bsz, -1) + offsets = (torch.arange(self._num_codebooks, device=device, dtype=torch.long) * self._codebook_size).unsqueeze( + 0 + ) + codebook_sum = self.codebook_embeddings(audio_codes + offsets).sum(dim=1).to(dtype=embeds.dtype) + is_semantic = (token_ids >= self._semantic_begin_id) & (token_ids <= self._semantic_end_id) + embeds = torch.where(is_semantic.unsqueeze(-1), embeds + codebook_sum, embeds) + return embeds, audio_codes.to(dtype=torch.long) + + # -------------------- prompt length estimation -------------------- + + @staticmethod + def estimate_prompt_len_from_additional_information( + additional_information: dict[str, Any] | None, + **kwargs: Any, + ) -> int: + """Upper-bound the text-only prompt length for placeholder allocation.""" + del kwargs + info = additional_information or {} + text = info.get("text", "") + if isinstance(text, list): + text = text[0] if text else "" + # ~1 token per character is a safe bound for CJK; +64 covers the + # chat-template scaffolding. + return max(2, len(str(text)) + 64) + + # -------------------- weight loading -------------------- + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + text_config = self.text_config + fast_config = self.fast_ar_config + remapped = _remap_audio8_tts_weights( + weights, + q_size=text_config.num_attention_heads * text_config.head_dim, + kv_size=text_config.num_key_value_heads * text_config.head_dim, + fast_q_size=fast_config.num_attention_heads * fast_config.head_dim, + fast_kv_size=fast_config.num_key_value_heads * fast_config.head_dim, + ) + + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded_params: set[str] = set() + unexpected: list[str] = [] + + for name, loaded_weight in remapped: + if name == "model.embed_tokens.weight" and self.text_config.tie_word_embeddings: + lm_key = "lm_head.weight" + if lm_key in params_dict: + param = params_dict[lm_key] + loader = getattr(param, "weight_loader", default_weight_loader) + loader(param, loaded_weight) + loaded_params.add(lm_key) + + handled = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + mapped = name.replace(weight_name, param_name) + if mapped not in params_dict: + continue + param = params_dict[mapped] + loader = getattr(param, "weight_loader", default_weight_loader) + if loader == default_weight_loader: + loader(param, loaded_weight) + else: + loader(param, loaded_weight, shard_id) + loaded_params.add(mapped) + handled = True + break + if handled: + continue + + if name in params_dict: + param = params_dict[name] + loader = getattr(param, "weight_loader", default_weight_loader) + loader(param, loaded_weight) + loaded_params.add(name) + else: + unexpected.append(name) + + if unexpected: + raise ValueError( + f"Audio8 TTS Slow AR received {len(unexpected)} unmapped checkpoint tensors, " + f"e.g. {sorted(unexpected)[:5]}. The weight remapper is out of sync with the checkpoint." + ) + missing = sorted(set(params_dict) - loaded_params) + if missing: + raise ValueError(f"Audio8 TTS Slow AR is missing weights for {missing[:5]} ({len(missing)} total)") + logger.info("Loaded %d weights for Audio8TTSSlowARForConditionalGeneration", len(loaded_params)) + + # The reference precomputes RoPE in bf16; keeping fp32 cos/sin here + # shifts the logits enough to trigger early EOS (same failure mode as + # Fish Speech). + truncated = 0 + for module in self.modules(): + cache = getattr(module, "cos_sin_cache", None) + if isinstance(cache, torch.Tensor): + module.cos_sin_cache = cache.to(torch.bfloat16).to(cache.dtype) + truncated += 1 + if truncated: + logger.info("Audio8 TTS: truncated %d RoPE cos_sin_cache buffers to bf16", truncated) + + return loaded_params + + +__all__ = ["Audio8TTSSlowARForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/audio8_tts/codec.py b/vllm_omni/model_executor/models/audio8_tts/codec.py new file mode 100644 index 00000000000..71f43642b33 --- /dev/null +++ b/vllm_omni/model_executor/models/audio8_tts/codec.py @@ -0,0 +1,600 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS neural audio codec (44.1 kHz, 10 codebooks). + +Vendored from the reference ``modeling_arktts_codec.py`` shipped with +``Audio8/Audio8-TTS-Preview-0.6b`` (Apache-2.0). The module tree and attribute +names are reproduced verbatim so ``codec.pth`` loads under ``strict=True``. + +Inference-only deviations: no ``torch.jit.script`` on Snake (it fights +``torch.compile`` / CUDA graphs), RoPE cached in a plain dict instead of a +non-persistent buffer (the pattern that leaves the reference LM with +uninitialised RoPE under ``from_pretrained`` on transformers >= 5), and +:func:`build_arktts_codec` prunes the unused encoder or decoder half per stage. +""" + +from __future__ import annotations + +import math +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from torch.nn.utils import weight_norm as legacy_weight_norm +from torch.nn.utils.parametrizations import weight_norm + +# The codec's own transformer blocks always use base 10000, independent of the +# 1e6 rope_base of the DualAR language model. +_CODEC_ROPE_BASE = 10000.0 + + +def _rope_table(length: int, head_dim: int, base: float, device: torch.device) -> Tensor: + frequencies = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)) + phases = torch.outer(torch.arange(length, device=device), frequencies) + values = torch.polar(torch.ones_like(phases), phases) + return torch.stack((values.real, values.imag), dim=-1).to(torch.bfloat16) + + +def _apply_rope(x: Tensor, values: Tensor) -> Tensor: + shaped = x.float().reshape(*x.shape[:-1], -1, 2) + values = values.view(1, shaped.shape[1], 1, shaped.shape[3], 2) + output = torch.stack( + ( + shaped[..., 0] * values[..., 0] - shaped[..., 1] * values[..., 1], + shaped[..., 1] * values[..., 0] + shaped[..., 0] * values[..., 1], + ), + dim=-1, + ) + return output.flatten(3).to(x.dtype) + + +class ArkttsCodecRMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-5) -> None: + super().__init__() + self.eps = float(eps) + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + output = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps) + return output.to(x.dtype) * self.weight + + +class ArkttsCodecLayerScale(nn.Module): + def __init__(self, dim: int, init_values: float = 1e-2) -> None: + super().__init__() + self.gamma = nn.Parameter(init_values * torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + return x * self.gamma + + +class ArkttsCodecTransformerConfig: + """Plain option holder for the codec's transformer blocks. + + Deliberately not a dataclass / msgspec Struct: it mirrors a reference + dataclass whose only job is to carry constructor arguments, and it is + instantiated from Python literals inside this module only. + """ + + def __init__( + self, + n_layer: int, + n_head: int, + dim: int, + intermediate_size: int, + n_local_heads: int = -1, + head_dim: int = 64, + rope_base: float = _CODEC_ROPE_BASE, + norm_eps: float = 1e-5, + channels_first: bool = True, + ) -> None: + self.n_layer = int(n_layer) + self.n_head = int(n_head) + self.dim = int(dim) + self.intermediate_size = int(intermediate_size) + self.n_local_heads = int(n_head) if n_local_heads == -1 else int(n_local_heads) + self.head_dim = int(head_dim) + self.rope_base = float(rope_base) + self.norm_eps = float(norm_eps) + self.channels_first = bool(channels_first) + + +class ArkttsCodecAttention(nn.Module): + def __init__(self, config: ArkttsCodecTransformerConfig) -> None: + super().__init__() + total = (config.n_head + 2 * config.n_local_heads) * config.head_dim + self.wqkv = nn.Linear(config.dim, total, bias=False) + self.wo = nn.Linear(config.head_dim * config.n_head, config.dim, bias=False) + self.n_head = config.n_head + self.n_local_heads = config.n_local_heads + self.head_dim = config.head_dim + + def forward(self, x: Tensor, rope_values: Tensor, mask: Tensor) -> Tensor: + batch, length, _ = x.shape + query_size = self.n_head * self.head_dim + kv_size = self.n_local_heads * self.head_dim + query, key, value = self.wqkv(x).split((query_size, kv_size, kv_size), dim=-1) + query = query.view(batch, length, self.n_head, self.head_dim) + key = key.view(batch, length, self.n_local_heads, self.head_dim) + value = value.view(batch, length, self.n_local_heads, self.head_dim) + query = _apply_rope(query, rope_values).transpose(1, 2) + key = _apply_rope(key, rope_values).transpose(1, 2) + value = value.transpose(1, 2) + output = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=mask, + enable_gqa=self.n_local_heads != self.n_head, + ) + output = output.transpose(1, 2).contiguous().view(batch, length, query_size) + return self.wo(output) + + +class ArkttsCodecFeedForward(nn.Module): + def __init__(self, config: ArkttsCodecTransformerConfig) -> None: + super().__init__() + self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False) + self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False) + self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False) + + def forward(self, x: Tensor) -> Tensor: + return self.w2(F.silu(self.w1(x)) * self.w3(x)) + + +class ArkttsCodecTransformerBlock(nn.Module): + def __init__(self, config: ArkttsCodecTransformerConfig) -> None: + super().__init__() + self.attention = ArkttsCodecAttention(config) + self.feed_forward = ArkttsCodecFeedForward(config) + self.ffn_norm = ArkttsCodecRMSNorm(config.dim, config.norm_eps) + self.attention_norm = ArkttsCodecRMSNorm(config.dim, config.norm_eps) + self.attention_layer_scale = ArkttsCodecLayerScale(config.dim) + self.ffn_layer_scale = ArkttsCodecLayerScale(config.dim) + + def forward(self, x: Tensor, rope_values: Tensor, mask: Tensor) -> Tensor: + hidden = x + self.attention_layer_scale(self.attention(self.attention_norm(x), rope_values, mask)) + return hidden + self.ffn_layer_scale(self.feed_forward(self.ffn_norm(hidden))) + + +class ArkttsCodecWindowTransformer(nn.Module): + """Causal transformer with a sliding attention window.""" + + def __init__( + self, + config: ArkttsCodecTransformerConfig, + input_dim: int, + window_size: int | None, + ) -> None: + super().__init__() + self.layers = nn.ModuleList([ArkttsCodecTransformerBlock(config) for _ in range(config.n_layer)]) + self.norm = ArkttsCodecRMSNorm(config.dim, config.norm_eps) + self.window_size = window_size + self.channels_first = config.channels_first + self.input_proj = nn.Linear(input_dim, config.dim) if input_dim != config.dim else nn.Identity() + self.output_proj = nn.Linear(config.dim, input_dim) if input_dim != config.dim else nn.Identity() + self.look_ahead_conv = nn.Identity() + self.head_dim = config.head_dim + self.rope_base = config.rope_base + self._mask_cache: dict[tuple[int, str], Tensor] = {} + self._rope_cache: dict[tuple[int, str], Tensor] = {} + + def _mask(self, length: int, device: torch.device) -> Tensor: + key = (length, str(device)) + cached = self._mask_cache.get(key) + if cached is not None: + return cached + row = torch.arange(length, device=device)[:, None] + column = torch.arange(length, device=device)[None, :] + mask = column <= row + if self.window_size is not None: + mask &= column >= (row - self.window_size + 1).clamp_min(0) + mask = mask[None, None] + self._mask_cache[key] = mask + return mask + + def _rope(self, length: int, device: torch.device) -> Tensor: + key = (length, str(device)) + cached = self._rope_cache.get(key) + if cached is not None: + return cached + values = _rope_table(length, self.head_dim, self.rope_base, device) + self._rope_cache[key] = values + return values + + def forward(self, x: Tensor) -> Tensor: + if self.channels_first: + x = x.transpose(1, 2) + x = self.look_ahead_conv(self.input_proj(x)) + length = int(x.shape[1]) + mask = self._mask(length, x.device) + rope_values = self._rope(length, x.device) + for layer in self.layers: + x = layer(x, rope_values, mask) + x = self.output_proj(self.norm(x)) + return x.transpose(1, 2) if self.channels_first else x + + +def _extra_padding(x: Tensor, kernel_size: int, stride: int, padding_total: int = 0) -> int: + length = x.shape[-1] + frames = (length - kernel_size + padding_total) / stride + 1 + ideal = (math.ceil(frames) - 1) * stride + kernel_size - padding_total + return ideal - length + + +class ArkttsCausalConv1d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + dilation: int = 1, + stride: int = 1, + groups: int = 1, + ) -> None: + super().__init__() + self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, groups=groups) + self.stride = stride + self.kernel_size = (kernel_size - 1) * dilation + 1 + self.padding = self.kernel_size - self.stride + + def forward(self, x: Tensor) -> Tensor: + right = _extra_padding(x, self.kernel_size, self.stride, self.padding) + return self.conv(F.pad(x, (self.padding, right))).contiguous() + + def apply_weight_norm(self) -> ArkttsCausalConv1d: + self.conv = weight_norm(self.conv) + return self + + +class ArkttsCausalConvTranspose1d(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + dilation: int = 1, + ) -> None: + super().__init__() + self.conv = nn.ConvTranspose1d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation) + self.stride = stride + self.kernel_size = kernel_size + + def forward(self, x: Tensor) -> Tensor: + x = self.conv(x) + crop = self.kernel_size - self.stride + return x[..., : x.shape[-1] - crop].contiguous() if crop else x.contiguous() + + def apply_weight_norm(self) -> ArkttsCausalConvTranspose1d: + self.conv = weight_norm(self.conv) + return self + + +def _causal_wn_conv(*args: Any, **kwargs: Any) -> ArkttsCausalConv1d: + return ArkttsCausalConv1d(*args, **kwargs).apply_weight_norm() + + +def _causal_wn_transpose(*args: Any, **kwargs: Any) -> ArkttsCausalConvTranspose1d: + return ArkttsCausalConvTranspose1d(*args, **kwargs).apply_weight_norm() + + +class ArkttsSnake1d(nn.Module): + def __init__(self, channels: int) -> None: + super().__init__() + self.alpha = nn.Parameter(torch.ones(1, channels, 1)) + + def forward(self, x: Tensor) -> Tensor: + shape = x.shape + x = x.reshape(shape[0], shape[1], -1) + x = x + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * x).pow(2) + return x.reshape(shape) + + +class ArkttsResidualUnit(nn.Module): + def __init__(self, dim: int, dilation: int) -> None: + super().__init__() + self.block = nn.Sequential( + ArkttsSnake1d(dim), + _causal_wn_conv(dim, dim, kernel_size=7, dilation=dilation), + ArkttsSnake1d(dim), + _causal_wn_conv(dim, dim, kernel_size=1), + ) + + def forward(self, x: Tensor) -> Tensor: + output = self.block(x) + difference = x.shape[-1] - output.shape[-1] + if difference > 0: + x = x[..., :-difference] + return x + output + + +class ArkttsEncoderBlock(nn.Module): + def __init__(self, dim: int, stride: int, transformer_layers: int) -> None: + super().__init__() + modules: list[nn.Module] = [ + ArkttsResidualUnit(dim // 2, 1), + ArkttsResidualUnit(dim // 2, 3), + ArkttsResidualUnit(dim // 2, 9), + ArkttsSnake1d(dim // 2), + _causal_wn_conv(dim // 2, dim, kernel_size=2 * stride, stride=stride), + ] + if transformer_layers: + config = ArkttsCodecTransformerConfig( + n_layer=transformer_layers, + n_head=dim // 64, + dim=dim, + intermediate_size=dim * 3, + ) + modules.append(ArkttsCodecWindowTransformer(config, dim, window_size=512)) + else: + modules.append(nn.Identity()) + self.block = nn.Sequential(*modules) + + def forward(self, x: Tensor) -> Tensor: + return self.block(x) + + +class ArkttsEncoder(nn.Module): + def __init__(self) -> None: + super().__init__() + dim = 64 + modules: list[nn.Module] = [_causal_wn_conv(1, dim, kernel_size=7)] + for stride, transformer_layers in zip((2, 4, 8, 8), (0, 0, 0, 4), strict=True): + dim *= 2 + modules.append(ArkttsEncoderBlock(dim, stride, transformer_layers)) + modules.extend((ArkttsSnake1d(dim), _causal_wn_conv(dim, 1024, kernel_size=3))) + self.block = nn.Sequential(*modules) + + def forward(self, x: Tensor) -> Tensor: + return self.block(x) + + +class ArkttsDecoderBlock(nn.Module): + def __init__(self, input_dim: int, output_dim: int, stride: int) -> None: + super().__init__() + self.block = nn.Sequential( + ArkttsSnake1d(input_dim), + _causal_wn_transpose(input_dim, output_dim, kernel_size=2 * stride, stride=stride), + ArkttsResidualUnit(output_dim, 1), + ArkttsResidualUnit(output_dim, 3), + ArkttsResidualUnit(output_dim, 9), + ) + + def forward(self, x: Tensor) -> Tensor: + return self.block(x) + + +class ArkttsDecoder(nn.Module): + def __init__(self) -> None: + super().__init__() + channels = 1536 + modules: list[nn.Module] = [_causal_wn_conv(1024, channels, kernel_size=7)] + output_dim = channels + for index, stride in enumerate((8, 8, 4, 2)): + input_dim = channels // (2**index) + output_dim = channels // (2 ** (index + 1)) + modules.append(ArkttsDecoderBlock(input_dim, output_dim, stride)) + modules.extend( + ( + ArkttsSnake1d(output_dim), + _causal_wn_conv(output_dim, 1, kernel_size=7), + nn.Tanh(), + ) + ) + self.model = nn.Sequential(*modules) + + def forward(self, x: Tensor) -> Tensor: + return self.model(x) + + +class ArkttsVectorQuantizer(nn.Module): + def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int) -> None: + super().__init__() + self.codebook_size = int(codebook_size) + self.codebook_dim = int(codebook_dim) + self.in_proj = legacy_weight_norm(nn.Conv1d(input_dim, codebook_dim, kernel_size=1)) + self.out_proj = legacy_weight_norm(nn.Conv1d(codebook_dim, input_dim, kernel_size=1)) + self.codebook = nn.Embedding(codebook_size, codebook_dim) + + def decode_code(self, indices: Tensor) -> Tensor: + return F.embedding(indices, self.codebook.weight).transpose(1, 2) + + def decode_latents(self, latents: Tensor) -> tuple[Tensor, Tensor]: + batch, _, length = latents.shape + flattened = F.normalize(latents.transpose(1, 2).reshape(batch * length, -1)) + codebook = F.normalize(self.codebook.weight) + distances = ( + flattened.pow(2).sum(1, keepdim=True) + - 2 * flattened @ codebook.t() + + codebook.pow(2).sum(1, keepdim=True).t() + ) + indices = (-distances).argmax(1).view(batch, length) + return self.decode_code(indices), indices + + def forward(self, z: Tensor) -> tuple[Tensor, Tensor]: + projected = self.in_proj(z) + quantized, indices = self.decode_latents(projected) + return self.out_proj(quantized), indices + + +class ArkttsResidualQuantizer(nn.Module): + def __init__(self, input_dim: int, n_codebooks: int, codebook_size: int, codebook_dim: int) -> None: + super().__init__() + self.n_codebooks = int(n_codebooks) + self.codebook_size = int(codebook_size) + self.quantizers = nn.ModuleList( + [ArkttsVectorQuantizer(input_dim, codebook_size, codebook_dim) for _ in range(n_codebooks)] + ) + + def forward(self, z: Tensor) -> tuple[Tensor, Tensor]: + quantized_sum: Tensor | float = 0.0 + residual = z + codes: list[Tensor] = [] + for quantizer in self.quantizers: + quantized, indices = quantizer(residual) + quantized_sum = quantized_sum + quantized + residual = residual - quantized + codes.append(indices) + return quantized_sum, torch.stack(codes, dim=1) + + def from_codes(self, codes: Tensor) -> Tensor: + output: Tensor | float = 0.0 + for index in range(codes.shape[1]): + projected = self.quantizers[index].decode_code(codes[:, index]) + output = output + self.quantizers[index].out_proj(projected) + return output + + +class ArkttsConvNeXtBlock(nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + self.dwconv = ArkttsCausalConv1d(dim, dim, kernel_size=7, groups=dim) + self.norm = nn.LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, 4 * dim) + self.act = nn.GELU() + self.pwconv2 = nn.Linear(4 * dim, dim) + self.gamma = nn.Parameter(1e-6 * torch.ones(dim)) + + def forward(self, x: Tensor) -> Tensor: + residual = x + x = self.dwconv(x).permute(0, 2, 1) + x = self.pwconv2(self.act(self.pwconv1(self.norm(x)))) + x = (self.gamma * x).permute(0, 2, 1) + return residual + x + + +class ArkttsDownsampleQuantizer(nn.Module): + """1 semantic codebook (4096) + 9 residual codebooks (1024), 4x downsampled.""" + + def __init__( + self, + post_n_layer: int = 8, + post_n_head: int = 16, + post_n_local_heads: int = 8, + post_intermediate_size: int = 1216, + ) -> None: + super().__init__() + self.semantic_quantizer = ArkttsResidualQuantizer(1024, 1, 4096, 8) + self.quantizer = ArkttsResidualQuantizer(1024, 9, 1024, 8) + self.downsample = nn.Sequential( + nn.Sequential(ArkttsCausalConv1d(1024, 1024, kernel_size=2, stride=2), ArkttsConvNeXtBlock(1024)), + nn.Sequential(ArkttsCausalConv1d(1024, 1024, kernel_size=2, stride=2), ArkttsConvNeXtBlock(1024)), + ) + self.upsample = nn.Sequential( + nn.Sequential(ArkttsCausalConvTranspose1d(1024, 1024, kernel_size=2, stride=2), ArkttsConvNeXtBlock(1024)), + nn.Sequential(ArkttsCausalConvTranspose1d(1024, 1024, kernel_size=2, stride=2), ArkttsConvNeXtBlock(1024)), + ) + self.pre_module = ArkttsCodecWindowTransformer( + ArkttsCodecTransformerConfig(n_layer=8, n_head=16, dim=1024, intermediate_size=3072), + 1024, + window_size=128, + ) + self.post_module = ArkttsCodecWindowTransformer( + ArkttsCodecTransformerConfig( + n_layer=post_n_layer, + n_head=post_n_head, + n_local_heads=post_n_local_heads, + dim=1024, + intermediate_size=post_intermediate_size, + ), + 1024, + window_size=128, + ) + self.semantic_predictor_module = nn.Identity() + + def forward(self, z: Tensor) -> tuple[Tensor, Tensor]: + original_length = z.shape[-1] + z = self.pre_module(self.downsample(z)) + semantic, semantic_codes = self.semantic_quantizer(z) + residual, residual_codes = self.quantizer(z - semantic) + z = self.upsample(self.post_module(semantic + residual)) + difference = original_length - z.shape[-1] + if difference > 0: + z = F.pad(z, (difference, 0)) + elif difference < 0: + z = z[..., -difference:] + return z, torch.cat((semantic_codes, residual_codes), dim=1) + + def decode(self, indices: Tensor) -> Tensor: + # Cloned because the clamps below are in-place and this is the caller's + # codes buffer; once per decode chunk, not per AR step. + indices = indices.clone() + indices[:, 0].clamp_(0, self.semantic_quantizer.codebook_size - 1) + indices[:, 1:].clamp_(0, self.quantizer.codebook_size - 1) + semantic = self.semantic_quantizer.from_codes(indices[:, :1]) + residual = self.quantizer.from_codes(indices[:, 1:]) + return self.upsample(self.post_module(semantic + residual)) + + +class ArkttsCodec(nn.Module): + """Audio8 TTS codec: waveform <-> 10 codebooks at ~21.5 frames/s.""" + + sample_rate = 44100 + frame_length = 2048 + + def __init__( + self, + post_n_layer: int = 8, + post_n_head: int = 16, + post_n_local_heads: int = 8, + post_intermediate_size: int = 1216, + ) -> None: + super().__init__() + self.encoder = ArkttsEncoder() + self.quantizer = ArkttsDownsampleQuantizer( + post_n_layer=post_n_layer, + post_n_head=post_n_head, + post_n_local_heads=post_n_local_heads, + post_intermediate_size=post_intermediate_size, + ) + self.decoder = ArkttsDecoder() + + @torch.inference_mode() + def encode(self, audio: Tensor, audio_lengths: Tensor | None = None) -> tuple[Tensor, Tensor]: + """Encode ``[B, 1, samples]`` (or ``[B, samples]``) to ``[B, 10, frames]``.""" + if audio.ndim == 2: + audio = audio[:, None] + if audio.ndim != 3 or audio.shape[1] != 1: + raise ValueError(f"audio must have shape [B, 1, samples], got {tuple(audio.shape)}") + original_length = int(audio.shape[-1]) + right = math.ceil(original_length / self.frame_length) * self.frame_length - original_length + if right: + audio = F.pad(audio, (0, right)) + if audio_lengths is None: + audio_lengths = torch.full((audio.shape[0],), original_length, device=audio.device, dtype=torch.long) + _, codes = self.quantizer(self.encoder(audio)) + code_lengths = torch.ceil(audio_lengths.float() / self.frame_length).long() + max_codes = int(codes.shape[-1]) + code_lengths = code_lengths.clamp_max(max_codes) + frame_index = torch.arange(max_codes, device=codes.device) + valid = frame_index[None, None, :] < code_lengths[:, None, None] + return torch.where(valid, codes, torch.full_like(codes, -1)), code_lengths + + @torch.inference_mode() + def decode(self, codes: Tensor) -> Tensor: + """Decode ``[B, 10, frames]`` to ``[B, 1, samples]``.""" + return self.decoder(self.quantizer.decode(codes.long())) + + +def build_arktts_codec( + *, + post_n_layer: int = 8, + post_n_head: int = 16, + post_n_local_heads: int = 8, + post_intermediate_size: int = 1216, +) -> ArkttsCodec: + """Construct the codec with uninitialised weights, on CPU, in eval mode.""" + codec = ArkttsCodec( + post_n_layer=post_n_layer, + post_n_head=post_n_head, + post_n_local_heads=post_n_local_heads, + post_intermediate_size=post_intermediate_size, + ) + codec.eval() + return codec + + +__all__ = ["ArkttsCodec", "build_arktts_codec"] diff --git a/vllm_omni/model_executor/models/audio8_tts/codec_utils.py b/vllm_omni/model_executor/models/audio8_tts/codec_utils.py new file mode 100644 index 00000000000..64f9c0afb2f --- /dev/null +++ b/vllm_omni/model_executor/models/audio8_tts/codec_utils.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Loading / caching helpers for the Audio8 TTS ``codec.pth`` checkpoint. + +Both stages need the codec but for opposite halves: Stage 0 encodes reference +audio for zero-shot voice cloning, Stage 1 decodes generated codes to waveform. +The unused half is pruned before the module is moved to the device so a stage +never pays GPU memory for weights it cannot reach. +""" + +from __future__ import annotations + +import os +from functools import lru_cache + +import numpy as np +import torch +import torch.nn as nn +from torch.nn.utils.parametrize import remove_parametrizations +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.audio8_tts.codec import ArkttsCodec, build_arktts_codec +from vllm_omni.model_executor.models.audio8_tts.configuration_audio8_tts import ( + ARKTTS_CODEC_FRAME_SIZE, + ARKTTS_CODEC_SAMPLE_RATE, +) + +logger = init_logger(__name__) + +CODEC_FILENAME = "codec.pth" + +_codec_cache: dict[tuple[str, str, str, str], ArkttsCodec] = {} + + +def resolve_codec_path(model_path: str, filename: str = CODEC_FILENAME) -> str: + """Return a local path to ``filename``, downloading from the Hub if needed.""" + local = os.path.join(model_path, filename) + if os.path.exists(local): + return local + + from transformers.utils.hub import cached_file + + cached = cached_file(model_path, filename) + if cached is None or not os.path.exists(cached): + raise FileNotFoundError( + f"{filename} not found for {model_path}; the Audio8 TTS checkpoint must ship its neural audio codec." + ) + return cached + + +def bake_weight_norm(module: nn.Module) -> int: + """Fold weight-norm parametrisations into plain weights (inference only).""" + baked = 0 + for submodule in module.modules(): + parametrizations = getattr(submodule, "parametrizations", None) + if not parametrizations: + continue + for name in list(parametrizations.keys()): + remove_parametrizations(submodule, name, leave_parametrized=True) + baked += 1 + return baked + + +def load_arktts_codec( + model_path: str, + *, + device: torch.device | str = "cpu", + dtype: torch.dtype = torch.float32, + role: str = "decode", + post_n_layer: int = 8, + post_n_head: int = 16, + post_n_local_heads: int = 8, + post_intermediate_size: int = 1216, +) -> ArkttsCodec: + """Load ``codec.pth`` for ``role`` in {"encode", "decode", "both"}. + + Instances are cached per (path, device, dtype, role) because both the Stage-0 + voice-clone path and the Stage-1 decoder may ask for the same codec many + times per request. + """ + if role not in {"encode", "decode", "both"}: + raise ValueError(f"role must be encode/decode/both, got {role!r}") + device = torch.device(device) + cache_key = (model_path, str(device), str(dtype), role) + cached_codec = _codec_cache.get(cache_key) + if cached_codec is not None: + return cached_codec + + codec_path = resolve_codec_path(model_path) + codec = build_arktts_codec( + post_n_layer=post_n_layer, + post_n_head=post_n_head, + post_n_local_heads=post_n_local_heads, + post_intermediate_size=post_intermediate_size, + ) + + state = torch.load(codec_path, map_location="cpu", weights_only=True) + if "state_dict" in state: + state = state["state_dict"] + if any("generator." in key for key in state): + state = {key.replace("generator.", ""): value for key, value in state.items() if "generator." in key} + state = {key: value for key, value in state.items() if not key.endswith(("freqs_cis", "causal_mask"))} + codec.load_state_dict(state, strict=True) + + baked = bake_weight_norm(codec) + if role == "encode": + codec.decoder = None + elif role == "decode": + codec.encoder = None + codec.quantizer.pre_module = None + codec.quantizer.downsample = None + + codec = codec.to(device=device, dtype=dtype) + codec.eval() + _codec_cache[cache_key] = codec + logger.info( + "Loaded Audio8 TTS codec from %s (role=%s, device=%s, dtype=%s, baked_weight_norms=%d)", + codec_path, + role, + device, + dtype, + baked, + ) + return codec + + +@lru_cache(maxsize=8) +def _resample_transform(source_sr: int, target_sr: int, device: torch.device, dtype: torch.dtype): + import torchaudio + + return torchaudio.transforms.Resample(source_sr, target_sr).to(device=device, dtype=dtype) + + +def prepare_reference_waveform( + wav_samples: list[float] | np.ndarray | torch.Tensor, + sample_rate: int, + *, + device: torch.device, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Down-mix to mono and resample reference audio to the codec's rate. + + Returns a 1-D waveform tensor. + """ + if isinstance(wav_samples, torch.Tensor): + wav = wav_samples.detach() + else: + wav = torch.as_tensor(np.asarray(wav_samples)) + wav = wav.to(device=device, dtype=dtype) + + if wav.ndim == 2: + # Accept [channels, samples] and [samples, channels]. + if wav.shape[0] <= 8 and wav.shape[1] > wav.shape[0]: + wav = wav.mean(dim=0) + elif wav.shape[-1] <= 8 and wav.shape[0] > wav.shape[-1]: + wav = wav.mean(dim=-1) + else: + wav = wav.mean(dim=0) + elif wav.ndim > 2: + wav = wav.reshape(-1, wav.shape[-1]).mean(dim=0) + wav = wav.reshape(-1) + + if wav.numel() == 0: + raise ValueError("Reference audio must not be empty") + if int(sample_rate) <= 0: + raise ValueError(f"Reference audio sample rate must be positive, got {sample_rate}") + if int(sample_rate) != ARKTTS_CODEC_SAMPLE_RATE: + transform = _resample_transform(int(sample_rate), ARKTTS_CODEC_SAMPLE_RATE, wav.device, wav.dtype) + wav = transform(wav.unsqueeze(0)).squeeze(0) + return wav.contiguous() + + +@torch.no_grad() +def encode_reference_audio_codes( + model_path: str, + wav_samples: list[float] | np.ndarray | torch.Tensor, + sample_rate: int, + *, + device: torch.device | str, + **codec_kwargs: int, +) -> torch.Tensor: + """Encode reference audio into codec codes. + + Returns: + ``[frames, num_codebooks]`` int64 codes on ``device``. + """ + device = torch.device(device) + codec = load_arktts_codec(model_path, device=device, dtype=torch.float32, role="encode", **codec_kwargs) + wav = prepare_reference_waveform(wav_samples, sample_rate, device=device) + lengths = torch.tensor([wav.numel()], device=device, dtype=torch.long) + codes, code_lengths = codec.encode(wav.reshape(1, 1, -1), lengths) + frames = int(code_lengths[0].item()) + codes_fq = codes[0, :, :frames].transpose(0, 1).to(dtype=torch.long).contiguous() + logger.info( + "Encoded Audio8 TTS reference audio: %d samples @ %d Hz -> frames=%d codebooks=%d", + int(wav.numel()), + int(sample_rate), + int(codes_fq.shape[0]), + int(codes_fq.shape[1]), + ) + return codes_fq + + +def estimate_reference_code_frames(num_samples: int, sample_rate: int) -> int: + """Frames the codec will emit for ``num_samples`` at ``sample_rate``. + + Used by the serving layer to size the prompt placeholder without paying for + a codec encode on the API thread. + """ + if num_samples <= 0 or sample_rate <= 0: + raise ValueError("Reference audio must have a positive length and sample rate") + resampled = max(1, -(-num_samples * ARKTTS_CODEC_SAMPLE_RATE // sample_rate)) + return max(1, -(-resampled // ARKTTS_CODEC_FRAME_SIZE)) + + +__all__ = [ + "CODEC_FILENAME", + "bake_weight_norm", + "encode_reference_audio_codes", + "estimate_reference_code_frames", + "load_arktts_codec", + "prepare_reference_waveform", + "resolve_codec_path", +] diff --git a/vllm_omni/model_executor/models/audio8_tts/configuration_audio8_tts.py b/vllm_omni/model_executor/models/audio8_tts/configuration_audio8_tts.py new file mode 100644 index 00000000000..de44ab2a2c7 --- /dev/null +++ b/vllm_omni/model_executor/models/audio8_tts/configuration_audio8_tts.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Configuration classes for Audio8 TTS Preview (``model_type = "arktts"``). + +Audio8's HuggingFace config is flat and uses DualAR field names (``dim``, +``n_layer``, ``n_head``, ``fast_dim``, ...); this module re-exports them under +Qwen2 attribute names so vLLM's ``Qwen2Model`` consumes the Slow AR half, and +splits the Fast AR half into its own sub-config. The backbone is Qwen2 rather +than Qwen3 (Fish Speech) because Audio8 uses ``qkv_bias=true`` with no q/k +norm -- exactly the Qwen2 attention shape. +""" + +from __future__ import annotations + +from typing import Any + +from transformers import PretrainedConfig + +#: Codec frame stride in waveform samples (encoder rates 2*4*8*8 = 512 times the +#: quantizer's 2*2 downsample). 44100 / 2048 ~= 21.5 frames per second. +ARKTTS_CODEC_FRAME_SIZE = 2048 +ARKTTS_CODEC_SAMPLE_RATE = 44100 +#: Only the first codebook is a 4096-entry semantic codebook; codebooks 1..N-1 +#: are 1024-entry residual codebooks (see ``ArkttsDownsampleQuantizer``). +ARKTTS_SEMANTIC_CODEBOOK_SIZE = 4096 +ARKTTS_RESIDUAL_CODEBOOK_SIZE = 1024 + + +class Audio8TTSSlowARConfig(PretrainedConfig): + """Slow AR config exposed with Qwen2-compatible attribute names.""" + + model_type = "arktts_slow_ar" + + def __init__( + self, + vocab_size: int = 155776, + dim: int = 896, + n_head: int = 14, + n_local_heads: int = 2, + head_dim: int = 64, + n_layer: int = 24, + intermediate_size: int = 4864, + max_seq_len: int = 2048, + rope_base: float = 1_000_000.0, + norm_eps: float = 1e-6, + attention_qkv_bias: bool = True, + attention_qk_norm: bool = False, + tie_word_embeddings: bool = True, + codebook_size: int = ARKTTS_SEMANTIC_CODEBOOK_SIZE, + num_codebooks: int = 10, + semantic_begin_id: int = 151678, + semantic_end_id: int = 155773, + eos_token_id: int = 151645, + pad_token_id: int = 151643, + **kwargs: Any, + ) -> None: + # Audio8 field names -> standard Transformers / Qwen2 names. + self.hidden_size = int(dim) + self.num_attention_heads = int(n_head) + self.num_key_value_heads = int(n_local_heads) + self.head_dim = int(head_dim) + self.num_hidden_layers = int(n_layer) + self.intermediate_size = int(intermediate_size) + self.max_position_embeddings = int(max_seq_len) + self.rms_norm_eps = float(norm_eps) + self.hidden_act = "silu" + # Qwen2Attention hardcodes qkv bias=True / o_proj bias=False, which is + # what Audio8 TTS ships. Keep the flags so a future checkpoint that + # flips them fails loudly in the model rather than loading silently. + self.attention_qkv_bias = bool(attention_qkv_bias) + self.attention_qk_norm = bool(attention_qk_norm) + # Set explicitly (rather than relying on vLLM's set_default_rope_theta) + # so that a checkpoint with a non-default rope_base is honoured. + self.rope_parameters = {"rope_type": "default", "rope_theta": float(rope_base)} + + # Codec / codebook fields. + self.codebook_size = int(codebook_size) + self.num_codebooks = int(num_codebooks) + self.semantic_begin_id = int(semantic_begin_id) + self.semantic_end_id = int(semantic_end_id) + + super().__init__( + vocab_size=int(vocab_size), + tie_word_embeddings=bool(tie_word_embeddings), + eos_token_id=int(eos_token_id), + pad_token_id=int(pad_token_id), + **kwargs, + ) + + +class Audio8TTSFastARConfig(PretrainedConfig): + """Fast AR config: the ``n_fast_layer`` residual-codebook predictor.""" + + model_type = "arktts_fast_ar" + + def __init__( + self, + codebook_size: int = ARKTTS_SEMANTIC_CODEBOOK_SIZE, + num_codebooks: int = 10, + fast_dim: int = 896, + fast_n_head: int = 14, + fast_n_local_heads: int = 2, + fast_head_dim: int = 64, + n_fast_layer: int = 4, + fast_intermediate_size: int = 4864, + fast_attention_qkv_bias: bool = False, + fast_attention_qk_norm: bool = False, + rope_base: float = 1_000_000.0, + norm_eps: float = 1e-6, + **kwargs: Any, + ) -> None: + self.hidden_size = int(fast_dim) + self.num_attention_heads = int(fast_n_head) + self.num_key_value_heads = int(fast_n_local_heads) + self.head_dim = int(fast_head_dim) + self.num_hidden_layers = int(n_fast_layer) + self.intermediate_size = int(fast_intermediate_size) + # The Fast AR sequence is [slow hidden state, code_0, ..., code_{N-1}]. + self.max_position_embeddings = int(num_codebooks) + self.rms_norm_eps = float(norm_eps) + self.hidden_act = "silu" + self.attention_qkv_bias = bool(fast_attention_qkv_bias) + self.attention_qk_norm = bool(fast_attention_qk_norm) + self.rope_theta = float(rope_base) + self.num_codebooks = int(num_codebooks) + + super().__init__(vocab_size=int(codebook_size), **kwargs) + + +class Audio8TTSConfig(PretrainedConfig): + """Top-level Audio8 TTS config (``model_type = "arktts"``). + + Accepts the flat HF checkpoint fields and derives ``text_config`` (Slow AR) + and ``fast_ar_config`` (Fast AR). ``get_text_config()`` returns the Slow AR + config, which is what ``Qwen2Model`` reads. + """ + + model_type = "arktts" + sub_configs = { + "text_config": Audio8TTSSlowARConfig, + "fast_ar_config": Audio8TTSFastARConfig, + } + + def __init__( + self, + text_config: dict | Audio8TTSSlowARConfig | None = None, + fast_ar_config: dict | Audio8TTSFastARConfig | None = None, + vocab_size: int = 155776, + dim: int = 896, + n_head: int = 14, + n_local_heads: int = 2, + head_dim: int = 64, + n_layer: int = 24, + intermediate_size: int = 4864, + max_seq_len: int = 2048, + rope_base: float = 1_000_000.0, + norm_eps: float = 1e-6, + attention_qkv_bias: bool = True, + attention_qk_norm: bool = False, + tie_word_embeddings: bool = True, + codebook_size: int = ARKTTS_SEMANTIC_CODEBOOK_SIZE, + num_codebooks: int = 10, + semantic_begin_id: int = 151678, + semantic_end_id: int = 155773, + n_fast_layer: int = 4, + fast_dim: int = 896, + fast_n_head: int = 14, + fast_n_local_heads: int = 2, + fast_head_dim: int = 64, + fast_intermediate_size: int = 4864, + fast_attention_qkv_bias: bool = False, + fast_attention_qk_norm: bool = False, + norm_fastlayer_input: bool = True, + codec_filename: str = "codec.pth", + codec_sample_rate: int = ARKTTS_CODEC_SAMPLE_RATE, + codec_frame_size: int = ARKTTS_CODEC_FRAME_SIZE, + codec_post_n_layer: int = 8, + codec_post_n_head: int = 16, + codec_post_n_local_heads: int = 8, + codec_post_intermediate_size: int = 1216, + ras_window_size: int = 10, + ras_temperature: float = 1.0, + ras_top_p: float = 0.9, + eos_token_id: int = 151645, + pad_token_id: int = 151643, + **kwargs: Any, + ) -> None: + if isinstance(text_config, dict): + text_config = Audio8TTSSlowARConfig(**text_config) + if isinstance(fast_ar_config, dict): + fast_ar_config = Audio8TTSFastARConfig(**fast_ar_config) + + self.text_config = text_config or Audio8TTSSlowARConfig( + vocab_size=vocab_size, + dim=dim, + n_head=n_head, + n_local_heads=n_local_heads, + head_dim=head_dim, + n_layer=n_layer, + intermediate_size=intermediate_size, + max_seq_len=max_seq_len, + rope_base=rope_base, + norm_eps=norm_eps, + attention_qkv_bias=attention_qkv_bias, + attention_qk_norm=attention_qk_norm, + tie_word_embeddings=tie_word_embeddings, + codebook_size=codebook_size, + num_codebooks=num_codebooks, + semantic_begin_id=semantic_begin_id, + semantic_end_id=semantic_end_id, + eos_token_id=eos_token_id, + pad_token_id=pad_token_id, + ) + self.fast_ar_config = fast_ar_config or Audio8TTSFastARConfig( + codebook_size=codebook_size, + num_codebooks=num_codebooks, + fast_dim=fast_dim, + fast_n_head=fast_n_head, + fast_n_local_heads=fast_n_local_heads, + fast_head_dim=fast_head_dim, + n_fast_layer=n_fast_layer, + fast_intermediate_size=fast_intermediate_size, + fast_attention_qkv_bias=fast_attention_qkv_bias, + fast_attention_qk_norm=fast_attention_qk_norm, + rope_base=rope_base, + norm_eps=norm_eps, + ) + + # Fields the stages read directly off the top-level config. + self.codebook_size = int(codebook_size) + self.num_codebooks = int(num_codebooks) + self.semantic_begin_id = int(semantic_begin_id) + self.semantic_end_id = int(semantic_end_id) + self.norm_fastlayer_input = bool(norm_fastlayer_input) + self.codec_filename = str(codec_filename) + self.codec_sample_rate = int(codec_sample_rate) + self.codec_frame_size = int(codec_frame_size) + self.codec_post_n_layer = int(codec_post_n_layer) + self.codec_post_n_head = int(codec_post_n_head) + self.codec_post_n_local_heads = int(codec_post_n_local_heads) + self.codec_post_intermediate_size = int(codec_post_intermediate_size) + # Repetition-Aware Sampling (RAS): resample a repeated semantic token + # from a flatter distribution instead of masking it. + self.ras_window_size = int(ras_window_size) + self.ras_temperature = float(ras_temperature) + self.ras_top_p = float(ras_top_p) + + super().__init__( + eos_token_id=int(eos_token_id), + pad_token_id=int(pad_token_id), + tie_word_embeddings=bool(tie_word_embeddings), + **kwargs, + ) + + def get_text_config(self, *args: Any, **kwargs: Any) -> Audio8TTSSlowARConfig: + return self.text_config + + +__all__ = [ + "ARKTTS_CODEC_FRAME_SIZE", + "ARKTTS_CODEC_SAMPLE_RATE", + "ARKTTS_RESIDUAL_CODEBOOK_SIZE", + "ARKTTS_SEMANTIC_CODEBOOK_SIZE", + "Audio8TTSConfig", + "Audio8TTSFastARConfig", + "Audio8TTSSlowARConfig", +] diff --git a/vllm_omni/model_executor/models/audio8_tts/pipeline.py b/vllm_omni/model_executor/models/audio8_tts/pipeline.py new file mode 100644 index 00000000000..836ab71d98b --- /dev/null +++ b/vllm_omni/model_executor/models/audio8_tts/pipeline.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS Preview pipeline topology. + +Stage 0: ``audio8_tts_slow_ar`` -- text -> semantic tokens + codec codes. +Stage 1: ``audio8_tts_codec_decoder`` -- codec codes -> 44.1 kHz waveform. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.audio8_tts" + +#: ``<|im_end|>``: the Slow AR ends the utterance with an end-of-turn token. +AUDIO8_TTS_EOS_TOKEN_ID = 151645 + +AUDIO8_TTS_PIPELINE = PipelineConfig( + model_type="arktts", + default_deploy_config_name="audio8_tts.yaml", + model_arch="Audio8TTSSlowARForConditionalGeneration", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="audio8_tts_slow_ar", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + engine_output_type="latent", + async_chunk_process_next_stage_input_func=(f"{_PROC}.slow_ar_to_codec_decoder_async_chunk"), + sampling_constraints={ + "detokenize": False, + "stop_token_ids": [AUDIO8_TTS_EOS_TOKEN_ID], + }, + ), + StagePipelineConfig( + stage_id=1, + model_stage="audio8_tts_codec_decoder", + model_arch="Audio8TTSCodecDecoder", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(0,), + final_output=True, + final_output_type="audio", + engine_output_type="audio", + sampling_constraints={"detokenize": True}, + ), + ), +) diff --git a/vllm_omni/model_executor/models/audio8_tts/prompt_utils.py b/vllm_omni/model_executor/models/audio8_tts/prompt_utils.py new file mode 100644 index 00000000000..f0660ac5c52 --- /dev/null +++ b/vllm_omni/model_executor/models/audio8_tts/prompt_utils.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS prompt construction. + +Reproduces ``ArkttsProcessor._prompt_segments`` from the reference checkpoint, +so tokenisation is byte-identical to HF inference. +""" + +from __future__ import annotations + +from typing import Any + +import regex as re + +SYSTEM_PROMPT_TEXT_ONLY = "convert the provided text to speech" +SYSTEM_PROMPT_CLONE_PREFIX = "convert the provided text to speech reference to the following:\n\nText:\n" +SYSTEM_PROMPT_CLONE_SUFFIX = "\n\nSpeech:\n" + +_SPEAKER_TAG_PATTERN = re.compile(r"<\|speaker:\d+\|>") +_LEGACY_SPEAKER_TAG_PATTERN = re.compile(r"") +_CONTROL_TOKEN_PATTERN = re.compile(r"<\|[^>]+\|>") + + +def clean_text(text: str) -> str: + """Collapse all whitespace runs to single spaces (reference ``_clean_text``).""" + return " ".join(str(text).strip().split()) + + +def normalize_text(text: str, *, add_default_speaker: bool = False) -> str: + """Clean text, normalise speaker tags, reject other control tokens. + + ``<|speaker:N|>`` is allowed; any other ``<|...|>`` token raises so callers + cannot inject ``<|semantic:...|>`` / ``<|im_end|>`` into a prompt. + """ + normalized = clean_text(_LEGACY_SPEAKER_TAG_PATTERN.sub(r"<|speaker:\1|>", text)) + + disallowed = [ + token for token in _CONTROL_TOKEN_PATTERN.findall(normalized) if not _SPEAKER_TAG_PATTERN.fullmatch(token) + ] + if disallowed: + tokens = ", ".join(sorted(set(disallowed))) + raise ValueError(f"Audio8 TTS input contains unsupported control token(s): {tokens}") + + if add_default_speaker and not _SPEAKER_TAG_PATTERN.search(normalized): + normalized = f"<|speaker:0|>{normalized}" + return normalized + + +def _encode(tokenizer: Any, text: str) -> list[int]: + return list(tokenizer.encode(text, add_special_tokens=False)) + + +def _encode_segments(tokenizer: Any, segments: list[str]) -> list[int]: + ids: list[int] = [] + for segment in segments: + ids.extend(_encode(tokenizer, segment)) + return ids + + +def build_text_only_prompt_ids(tokenizer: Any, text: str) -> tuple[list[int], str]: + """Build the no-reference prompt. + + Returns: + ``(prompt_token_ids, normalized_text)``. + """ + normalized = normalize_text(text) + if not normalized: + raise ValueError("Audio8 TTS input text must not be empty") + ids = _encode_segments( + tokenizer, + [ + "<|im_start|>system\n", + SYSTEM_PROMPT_TEXT_ONLY, + "<|im_end|>\n", + "<|im_start|>user\n", + normalized, + "<|im_end|>\n", + "<|im_start|>assistant\n<|voice|>", + ], + ) + return ids, normalized + + +def build_voice_clone_prompt_parts( + tokenizer: Any, + text: str, + ref_text: str, +) -> tuple[list[int], list[int], str, str]: + """Build the clone prompt around the reference-code slot. + + Caller splices ``[code + semantic_begin_id for code in ref_codes[:, 0]]`` + between the returned prefix and suffix. + """ + normalized = normalize_text(text) + if not normalized: + raise ValueError("Audio8 TTS input text must not be empty") + normalized_ref = normalize_text(ref_text, add_default_speaker=True) + if not normalized_ref: + raise ValueError("Audio8 TTS voice cloning requires a non-empty reference transcript") + + prefix = _encode_segments( + tokenizer, + [ + "<|im_start|>system\n", + SYSTEM_PROMPT_CLONE_PREFIX, + normalized_ref, + SYSTEM_PROMPT_CLONE_SUFFIX, + ], + ) + suffix = _encode_segments( + tokenizer, + [ + "<|im_end|>\n", + "<|im_start|>user\n", + normalized, + "<|im_end|>\n", + "<|im_start|>assistant\n<|voice|>", + ], + ) + return prefix, suffix, normalized, normalized_ref + + +def build_voice_clone_prompt_ids( + tokenizer: Any, + text: str, + ref_text: str, + semantic_token_ids: list[int], +) -> tuple[list[int], int, str, str]: + """Build the full clone prompt. + + ``ref_start_index`` is the position of the first reference frame, i.e. + where the residual codebook embeddings must be added. + """ + prefix, suffix, normalized, normalized_ref = build_voice_clone_prompt_parts(tokenizer, text, ref_text) + return prefix + list(semantic_token_ids) + suffix, len(prefix), normalized, normalized_ref + + +def estimate_voice_clone_prompt_len( + tokenizer: Any, + normalized_text: str, + normalized_ref_text: str, + ref_frames: int, +) -> int: + """Exact clone prompt length, without encoding the reference audio.""" + prefix, suffix, _, _ = build_voice_clone_prompt_parts(tokenizer, normalized_text, normalized_ref_text) + return len(prefix) + max(0, int(ref_frames)) + len(suffix) + + +__all__ = [ + "SYSTEM_PROMPT_CLONE_PREFIX", + "SYSTEM_PROMPT_CLONE_SUFFIX", + "SYSTEM_PROMPT_TEXT_ONLY", + "build_text_only_prompt_ids", + "build_voice_clone_prompt_ids", + "build_voice_clone_prompt_parts", + "clean_text", + "estimate_voice_clone_prompt_len", + "normalize_text", +] diff --git a/vllm_omni/model_executor/models/audio8_tts/sampling.py b/vllm_omni/model_executor/models/audio8_tts/sampling.py new file mode 100644 index 00000000000..63840841127 --- /dev/null +++ b/vllm_omni/model_executor/models/audio8_tts/sampling.py @@ -0,0 +1,227 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS sampling primitives, ported 1:1 from the reference checkpoint. + +Filter (top-k/top-p on **raw** logits) -> divide surviving logits by +temperature -> Gumbel-max draw. This is the reverse of the usual filter order +and is audibly different at ``temperature != 1``, which is why this cannot +reuse ``models/common/nucleus_ras_sampling.py``. Per-row scalars are accepted +so mixed batches run in one shot without a GPU->CPU sync. +""" + +from __future__ import annotations + +import torch + +#: Below this temperature a request is treated as greedy, matching vLLM. +SAMPLING_EPS = 1e-5 + + +def _as_column( + value: float | int | torch.Tensor, + batch_size: int, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + """Broadcast a scalar or ``[B]`` tensor to a ``[B, 1]`` tensor.""" + if isinstance(value, torch.Tensor): + return value.to(device=device, dtype=dtype).reshape(-1, 1).expand(batch_size, 1) + return torch.full((batch_size, 1), float(value), device=device, dtype=dtype) + + +def filter_top_k_top_p( + logits: torch.Tensor, + top_k: int | torch.Tensor, + top_p: float | torch.Tensor, +) -> torch.Tensor: + """Mask candidates outside the top-k / top-p set of the raw logits. + + ``top_k <= 0`` and ``top_p >= 1`` disable each filter. The highest-scoring + candidate is always kept, so no row can become all ``-inf``. + """ + if logits.ndim != 2: + raise ValueError(f"logits must be 2-D [B, V], got {tuple(logits.shape)}") + batch_size, vocab = logits.shape + scalar_k = not isinstance(top_k, torch.Tensor) + scalar_p = not isinstance(top_p, torch.Tensor) + if scalar_k and scalar_p and (int(top_k) <= 0 or int(top_k) >= vocab) and float(top_p) >= 1.0: + return logits + + sorted_scores, sorted_indices = torch.sort(logits, descending=True, dim=-1) + cumulative = torch.cumsum(torch.softmax(sorted_scores, dim=-1), dim=-1) + positions = torch.arange(vocab, device=logits.device).unsqueeze(0) + + keep_k = _as_column(top_k, batch_size, logits.device, torch.long) + # top_k <= 0 means "no cut"; represent it as the full vocabulary. + keep_k = torch.where(keep_k <= 0, torch.full_like(keep_k, vocab), keep_k) + threshold_p = _as_column(top_p, batch_size, logits.device, cumulative.dtype) + + remove_sorted = (cumulative > threshold_p) | (positions >= keep_k) + # Always keep the argmax so at least one candidate survives. + remove_sorted[:, 0] = False + remove = torch.zeros_like(remove_sorted).scatter(-1, sorted_indices, remove_sorted) + return logits.masked_fill(remove, float("-inf")) + + +def gumbel_argmax_sample( + scores: torch.Tensor, + generator: torch.Generator | None = None, +) -> torch.Tensor: + """Draw one index per row via Gumbel-max over ``softmax(scores)``.""" + probabilities = torch.softmax(scores, dim=-1) + uniform = torch.rand( + probabilities.shape, + dtype=probabilities.dtype, + device=probabilities.device, + generator=generator, + ) + # -log(u) is Exp(1); argmax(p / Exp(1)) samples from p. + noise = -torch.log(uniform.clamp_min(torch.finfo(probabilities.dtype).tiny)) + return torch.argmax(probabilities / noise, dim=-1) + + +def sample_scores( + logits: torch.Tensor, + *, + temperature: float | torch.Tensor, + top_k: int | torch.Tensor, + top_p: float | torch.Tensor, + do_sample: bool = True, + generator: torch.Generator | None = None, +) -> torch.Tensor: + """Filter -> temper -> draw, in the reference order. + + Rows whose temperature is below :data:`SAMPLING_EPS` fall back to argmax. + """ + filtered = filter_top_k_top_p(logits, top_k, top_p) + greedy = filtered.argmax(dim=-1) + if not do_sample: + return greedy + if not isinstance(temperature, torch.Tensor): + if float(temperature) < SAMPLING_EPS: + return greedy + return gumbel_argmax_sample(filtered / float(temperature), generator=generator) + + temperature_col = _as_column(temperature, filtered.shape[0], filtered.device, filtered.dtype) + safe_temperature = torch.where( + temperature_col < SAMPLING_EPS, + torch.ones_like(temperature_col), + temperature_col, + ) + sampled = gumbel_argmax_sample(filtered / safe_temperature, generator=generator) + return torch.where(temperature_col.reshape(-1) < SAMPLING_EPS, greedy, sampled) + + +def ras_sample_semantic( + logits: torch.Tensor, + recent_ids: torch.Tensor | None, + *, + temperature: float | torch.Tensor, + top_k: int | torch.Tensor, + top_p: float | torch.Tensor, + ras_temperature: float, + ras_top_p: float, + do_sample: bool = True, + num_semantic_ids: int | None = None, + generator: torch.Generator | None = None, +) -> torch.Tensor: + """Repetition-Aware Sampling for the semantic (Slow AR) token. + + Draws with the request's ``temperature`` / ``top_p``; if the drawn id is + already in ``recent_ids`` (pad with a value that cannot be sampled, e.g. + ``-1``), redraws from the flatter RAS distribution instead. The + ``num_semantic_ids`` slot and beyond (EOS) are never substituted, so a + repeat EOS still ends the utterance. + + DualAR RAS *resamples* repeats rather than masking them — masking the + dominant token inflates EOS probability and truncates the utterance. + """ + normal = sample_scores( + logits, + temperature=temperature, + top_k=top_k, + top_p=top_p, + do_sample=do_sample, + generator=generator, + ) + if not do_sample or recent_ids is None or recent_ids.numel() == 0: + return normal + + high = sample_scores( + logits, + temperature=ras_temperature, + top_k=top_k, + top_p=ras_top_p, + do_sample=True, + generator=generator, + ) + repeated = (recent_ids == normal[:, None]).any(dim=1) + if num_semantic_ids is not None: + repeated &= normal < int(num_semantic_ids) + return torch.where(repeated, high, normal) + + +def _row_value(value: object, index: int) -> object: + """Slice a per-row value for request ``index``; scalars pass through.""" + return value[index : index + 1] if isinstance(value, torch.Tensor) else value + + +def ras_sample_batch( + logits: torch.Tensor, + recent_ids: torch.Tensor | None, + *, + temperature: float | torch.Tensor, + top_k: int | torch.Tensor, + top_p: float | torch.Tensor, + ras_temperature: float, + ras_top_p: float, + num_semantic_ids: int | None, + generators: dict, + num_reqs: int, +) -> torch.Tensor: + """Batched RAS that honors each request's own RNG generator. + + vLLM keys ``generators`` by batch slot, so a single shared generator would + leak slot 0's RNG stream to every row and make a prompt's audio depend on + its batch neighbours (undercutting per-request ``seed``). The single-request + case keeps the batched fast path; larger batches loop so each row draws from + its own (possibly unseeded) generator. + """ + common = dict( + ras_temperature=ras_temperature, + ras_top_p=ras_top_p, + num_semantic_ids=num_semantic_ids, + ) + if num_reqs == 1: + return ras_sample_semantic( + logits, + recent_ids, + temperature=temperature, + top_k=top_k, + top_p=top_p, + generator=generators.get(0), + **common, + ) + rows = [ + ras_sample_semantic( + logits[i : i + 1], + None if recent_ids is None else recent_ids[i : i + 1], + temperature=_row_value(temperature, i), + top_k=_row_value(top_k, i), + top_p=_row_value(top_p, i), + generator=generators.get(i), + **common, + ) + for i in range(num_reqs) + ] + return torch.cat(rows, dim=0) + + +__all__ = [ + "SAMPLING_EPS", + "filter_top_k_top_p", + "gumbel_argmax_sample", + "ras_sample_batch", + "ras_sample_semantic", + "sample_scores", +] diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index cb5167597a5..b9373e5f229 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -304,6 +304,17 @@ "gepard_talker", "GepardTalkerForConditionalGeneration", ), + ## audio8_tts (Audio8 TTS Preview 0.6B) + "Audio8TTSSlowARForConditionalGeneration": ( + "audio8_tts", + "audio8_tts_slow_ar", + "Audio8TTSSlowARForConditionalGeneration", + ), + "Audio8TTSCodecDecoder": ( + "audio8_tts", + "audio8_tts_codec_decoder", + "Audio8TTSCodecDecoder", + ), ## VoxCPM2 "VoxCPM2TalkerForConditionalGeneration": ( "voxcpm2", diff --git a/vllm_omni/model_executor/stage_input_processors/audio8_tts.py b/vllm_omni/model_executor/stage_input_processors/audio8_tts.py new file mode 100644 index 00000000000..4b453c13ab7 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/audio8_tts.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS Preview: Slow AR -> codec decoder stage input processor.""" + +from collections.abc import Mapping +from typing import Any + +import torch + +from vllm_omni.data_entry_keys import ( + CodesStruct, + MetaStruct, + OmniPayloadStruct, +) + +#: Chunking defaults; ~21.5 codec frames per second, so 25 frames ~= 1.16 s. +DEFAULT_CHUNK_FRAMES = 25 +DEFAULT_LEFT_CONTEXT_FRAMES = 25 + + +def _connector_extra(transfer_manager: Any) -> dict[str, Any]: + connector = getattr(transfer_manager, "connector", None) + raw_cfg = getattr(connector, "config", {}) or {} + return raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {} + + +def _cfg_int(cfg: dict[str, Any], key: str, default: int) -> int: + value = cfg.get(key, default) + try: + return int(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"Invalid Audio8 TTS integer config {key}={value!r}") from exc + + +def _request_initial_chunk_frames(request: Any, default: int) -> int: + """Per-request override of ``initial_codec_chunk_frames``, if present.""" + additional_information = getattr(request, "additional_information", None) + entries = getattr(additional_information, "entries", None) + if not entries or "initial_codec_chunk_frames" not in entries: + return default + entry = entries["initial_codec_chunk_frames"] + if entry.list_data is not None and len(entry.list_data) == 1: + return int(entry.list_data[0]) + return default + + +def extract_last_frame(multimodal_output: Mapping[str, Any]) -> torch.Tensor | None: + """Return the newest frame of codec codes, or ``None`` if this step had none. + + Prefill steps emit an all-zero placeholder frame, which must not be + forwarded to the decoder. + """ + audio_codes = multimodal_output.get("audio_codes") + if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0: + return None + if audio_codes.ndim == 1: + return audio_codes.to(device="cpu", dtype=torch.long).reshape(-1) + if audio_codes.ndim != 2: + raise ValueError(f"Invalid audio_codes shape for Audio8 TTS async_chunk: {tuple(audio_codes.shape)}") + + frame = audio_codes[-1] + if frame.numel() == 0: + return None + valid = multimodal_output.get("audio_code_valid") + if isinstance(valid, torch.Tensor) and valid.numel() > 0: + is_valid = bool(valid.reshape(-1)[-1].item()) + elif valid is not None: + is_valid = bool(valid) + else: + is_valid = bool(frame.any().item()) + if not is_valid: + return None + return frame.to(device="cpu", dtype=torch.long).reshape(-1) + + +def slow_ar_to_codec_decoder_async_chunk( + transfer_manager: Any, + multimodal_output: dict[str, Any] | None, + request: Any, + is_finished: bool = False, +) -> OmniPayloadStruct | None: + """Emit ``[left_context | new_frames]`` chunks as the Slow AR produces codes. + + Codes accumulate across AR steps and are shipped as a + ``[num_codebooks, frames]`` tensor. ``meta.left_context_size`` tells the + decoder how many leading frames to drop, keeping chunk boundaries + artefact-free. Returns ``None`` when the current chunk is not full yet. + """ + request_id = request.external_req_id + finished = bool(is_finished or request.is_finished()) + cfg = _connector_extra(transfer_manager) + + appended_this_call = False + if isinstance(multimodal_output, Mapping): + frame = extract_last_frame(multimodal_output) + if frame is not None: + transfer_manager.code_prompt_token_ids[request_id].append(frame.detach()) + appended_this_call = True + elif not finished: + return None + + chunk_frames = _cfg_int(cfg, "codec_chunk_frames", DEFAULT_CHUNK_FRAMES) + left_context_frames = _cfg_int(cfg, "codec_left_context_frames", DEFAULT_LEFT_CONTEXT_FRAMES) + initial_chunk_frames = _request_initial_chunk_frames(request, _cfg_int(cfg, "initial_codec_chunk_frames", 0)) + if chunk_frames <= 0 or left_context_frames < 0 or initial_chunk_frames < 0: + raise ValueError( + "Invalid Audio8 TTS codec chunk config: " + f"codec_chunk_frames={chunk_frames}, codec_left_context_frames={left_context_frames}, " + f"initial_codec_chunk_frames={initial_chunk_frames}" + ) + initial_chunk_frames = min(initial_chunk_frames, chunk_frames) + + frames = transfer_manager.code_prompt_token_ids[request_id] + length = len(frames) + if length <= 0: + if finished: + # Still emit a terminal payload so the decoder can close the stream. + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), + meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)), + ) + return None + + # A smaller first chunk trades a little quality for time-to-first-audio. + if initial_chunk_frames > 0 and length <= chunk_frames: + already_sent = transfer_manager.put_req_chunk[request_id] * initial_chunk_frames + pending = length - already_sent + if pending <= 0: + return None + if pending < initial_chunk_frames and not finished: + return None + new_frames = min(pending, initial_chunk_frames) + left_context_size = max(0, length - new_frames) + window = frames[:length] + else: + initial_coverage = ( + (chunk_frames // initial_chunk_frames) * initial_chunk_frames if initial_chunk_frames > 0 else 0 + ) + pending = (length - initial_coverage) % chunk_frames + if pending != 0 and not finished: + return None + if pending == 0 and finished and not appended_this_call: + # This boundary's chunk was already shipped when the stream first + # reached it; a terminal call carrying no new frame must close the + # stream, not re-send that window (otherwise the last chunk_frames + # frames ship twice). + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), + meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)), + ) + new_frames = pending if pending != 0 else chunk_frames + end_index = min(length, left_context_frames + new_frames) + left_context_size = max(0, end_index - new_frames) + window = frames[-end_index:] + + # Codebook-major [num_codebooks, frames]; the tensor payload avoids + # expanding codec indices into Python ints across the connector boundary. + codes_qf = torch.stack(window, dim=0).transpose(0, 1).contiguous() + return OmniPayloadStruct( + codes=CodesStruct(audio=codes_qf), + meta=MetaStruct( + left_context_size=left_context_size, + finished=torch.tensor(finished, dtype=torch.bool), + ), + ) + + +__all__ = [ + "DEFAULT_CHUNK_FRAMES", + "DEFAULT_LEFT_CONTEXT_FRAMES", + "extract_last_frame", + "slow_ar_to_codec_decoder_async_chunk", +] diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index 42d6d8d455f..6beb9101866 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.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 """ Custom model configs that cannot be expressed via HuggingFace Transformers alone, following the same pattern as vllm.transformers_utils.configs. @@ -10,6 +10,9 @@ import importlib _CLASS_TO_MODULE: dict[str, str] = { + "Audio8TTSConfig": "vllm_omni.transformers_utils.configs.audio8_tts", + "Audio8TTSSlowARConfig": "vllm_omni.transformers_utils.configs.audio8_tts", + "Audio8TTSFastARConfig": "vllm_omni.transformers_utils.configs.audio8_tts", "HiggsAudioV3Config": "vllm_omni.transformers_utils.configs.higgs_audio_v3", "Mammothmoda2Config": "vllm_omni.transformers_utils.configs.mammoth_moda2", "Mammothmoda2Qwen2_5_VLConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2", @@ -35,6 +38,9 @@ } __all__ = [ + "Audio8TTSConfig", + "Audio8TTSFastARConfig", + "Audio8TTSSlowARConfig", "HiggsAudioV3Config", "Mammothmoda2Config", "Mammothmoda2Qwen2_5_VLConfig", @@ -75,6 +81,7 @@ def __dir__(): # Eagerly import all config modules so their AutoConfig.register() side-effects # run as soon as `vllm_omni.transformers_utils.configs` is imported. +from vllm_omni.transformers_utils.configs import audio8_tts as _audio8_tts # noqa: F401, E402 from vllm_omni.transformers_utils.configs import cosyvoice3 as _cosyvoice3 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import dots_tts as _dots_tts # noqa: F401, E402 from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech # noqa: F401, E402 diff --git a/vllm_omni/transformers_utils/configs/audio8_tts.py b/vllm_omni/transformers_utils/configs/audio8_tts.py new file mode 100644 index 00000000000..42e92041978 --- /dev/null +++ b/vllm_omni/transformers_utils/configs/audio8_tts.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio8 TTS config registration with transformers AutoConfig. + +Registers ``model_type = "arktts"`` (and the two sub-config types) so +``AutoConfig.from_pretrained("Audio8/Audio8-TTS-Preview-0.6b")`` resolves to +vllm-omni's Qwen2-shaped config instead of the checkpoint's remote code. + +This only wins when ``trust_remote_code`` is **off** -- transformers prefers a +checkpoint's ``auto_map`` over registered classes otherwise -- which is why +``deploy/audio8_tts.yaml`` sets ``trust_remote_code: false``. +""" + +from transformers import AutoConfig + +from vllm_omni.model_executor.models.audio8_tts.configuration_audio8_tts import ( + Audio8TTSConfig, + Audio8TTSFastARConfig, + Audio8TTSSlowARConfig, +) + +AutoConfig.register("arktts", Audio8TTSConfig) +AutoConfig.register("arktts_slow_ar", Audio8TTSSlowARConfig) +AutoConfig.register("arktts_fast_ar", Audio8TTSFastARConfig) + +__all__ = [ + "Audio8TTSConfig", + "Audio8TTSFastARConfig", + "Audio8TTSSlowARConfig", +] From 81c1c460b3eb12083a02abe0b8d85247964d4526 Mon Sep 17 00:00:00 2001 From: zhengjia Date: Tue, 8 Sep 2026 22:15:48 -0400 Subject: [PATCH 14/93] [Bugfix][Frontend] Accept the msgpack-numpy package's numpy markers on the OpenPI endpoint (#6051) Signed-off-by: zjli2013 Co-authored-by: Cursor --- .../openai_api/test_openpi_connection.py | 54 ++++++++++++++++++- vllm_omni/entrypoints/openpi/connection.py | 16 +++++- 2 files changed, 67 insertions(+), 3 deletions(-) diff --git a/tests/entrypoints/openai_api/test_openpi_connection.py b/tests/entrypoints/openai_api/test_openpi_connection.py index 398ae649415..19ec23b6de2 100644 --- a/tests/entrypoints/openai_api/test_openpi_connection.py +++ b/tests/entrypoints/openai_api/test_openpi_connection.py @@ -68,7 +68,7 @@ def test_pack_and_unpack_round_trip_numpy_values(): assert decoded["nested"][0]["done"] == np.bool_(True) -def test_unpack_accepts_msgpack_numpy_marker_dicts(): +def test_unpack_accepts_vllm_native_marker_dicts(): action = np.asarray([[1.0, 2.0]], dtype=np.float32) payload = { b"actions": { @@ -88,6 +88,58 @@ def test_unpack_accepts_msgpack_numpy_marker_dicts(): np.testing.assert_allclose(decoded[b"actions"], np.asarray([[1.0, 0.0]], dtype=np.float32)) +def test_unpack_accepts_msgpack_numpy_marker_dicts(): + action = np.asarray([[1.0, 2.0]], dtype=np.float32) + payload = { + b"actions": { + b"nd": True, + b"type": action.dtype.str, + b"kind": b"", + b"shape": action.shape, + b"data": action.tobytes(), + } + } + + decoded = openpi_connection._unpack_numpy(payload) + + np.testing.assert_allclose(decoded[b"actions"], action) + assert decoded[b"actions"].dtype == np.float32 + + +def test_unpack_rejects_msgpack_numpy_structured_array_markers(): + values = np.zeros(2, dtype=[("a", " {nd: true, type, kind, shape, data} scalar -> {nd: false, type, kind, data} + +`kind` is required on the `nd` markers because it is what separates them from a +plain user mapping, but its value is dialect-specific: the vLLM-native markers +carry the dtype kind character while the `msgpack-numpy` package leaves it empty +for plain dtypes and sets "V" for structured ones. Both spellings are accepted. +Scalars packed by `msgpack-numpy` omit `kind` entirely and are therefore left as +mappings. """ from __future__ import annotations @@ -83,9 +90,14 @@ def _decode_vllm_numpy_marker(obj: dict[Any, Any]) -> Any: if nd is _MISSING or dtype is _MISSING or kind is _MISSING or data is _MISSING: return _MISSING - dtype_obj = np.dtype(_decode_marker_text(dtype)) kind_text = _decode_marker_text(kind) - if dtype_obj.kind != kind_text: + # Structured markers carry a dtype descriptor list in `type`, so reject them + # before `np.dtype()` sees something it cannot parse. + if kind_text == "V": + raise ValueError("Unsupported dtype: structured arrays") + + dtype_obj = np.dtype(_decode_marker_text(dtype)) + if kind_text not in ("", dtype_obj.kind): raise ValueError(f"NumPy dtype marker kind mismatch: {dtype_obj.kind!r} != {kind_text!r}") if dtype_obj.kind in ("V", "O", "c"): raise ValueError(f"Unsupported dtype: {dtype_obj}") From 2bae550c04318269d6f672a670da4d0776307248 Mon Sep 17 00:00:00 2001 From: Rakesh Kariya <83279947+rk9595@users.noreply.github.com> Date: Wed, 9 Sep 2026 07:45:59 +0530 Subject: [PATCH 15/93] [Frontend] Opt-in WebSocket TTS split_granularity and session seed (#7046) Signed-off-by: rk9595 Signed-off-by: Rakesh Kariya Co-authored-by: Cursor Co-authored-by: Claude Opus 5 --- docs/serving/speech_api.md | 39 ++-- .../online_serving/text_to_speech/README.md | 10 +- .../qwen3_tts/streaming_speech_client.py | 19 +- .../openai_api/test_serving_speech_stream.py | 161 +++++++++++++- .../openai_api/test_speech_text_splitter.py | 127 +++++++++++ .../entrypoints/openai/protocol/audio.py | 16 ++ .../openai/serving_speech_stream.py | 82 ++++--- .../openai/speech_text_splitter.py | 204 ++++++++++++++++++ 8 files changed, 608 insertions(+), 50 deletions(-) create mode 100644 tests/entrypoints/openai_api/test_speech_text_splitter.py create mode 100644 vllm_omni/entrypoints/openai/speech_text_splitter.py diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index a515e655768..f0658f278ab 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -279,10 +279,15 @@ curl -X POST http://localhost:8091/v1/audio/voices \ ## Streaming Text Input (WebSocket) -The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally and generates audio per sentence as boundaries are detected. +The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally. +By default (`split_granularity=none`) it buffers until `input.done` and +synthesizes that flush as **one** TTS request, which keeps long-form timbre +stable. Set `split_granularity` to `sentence` or `clause` to emit a request +at each detected boundary (lower time-to-first-audio for STT/LLM pipelines). -> Note: text input is always streamed incrementally. Audio output remains sentence-scoped: -> use `stream_audio=false` for one binary frame per sentence, or `stream_audio=true` for one or more PCM chunks per sentence. +> Note: `stream_audio` only changes how **audio bytes** are framed (one WAV/PCM +> payload vs chunked PCM). Text segmentation is controlled separately by +> `split_granularity`. ### WebSocket Protocol @@ -314,14 +319,16 @@ upstream LLM) pays the WebSocket handshake once instead of once per utterance. - The session config is sticky. Send `input.text` again straight after `session.done` to reuse it, or send another `session.config` first to change - voice, format, or reference audio. A `session.config` sent while text is - still buffered is rejected so no pending input is silently dropped. -- An utterance is the flush unit, not a linguistic one: it is whatever text was - buffered when `input.done` arrived, of any length, synthesized as one request. - `utterance_index` counts those flushes across the connection, so it tells you - which `input.done` a frame belongs to. `sentence_index` counts within one - flush and so pairs with `total_sentences`, which means every utterance reports - `sentence_index: 0` of `total_sentences: 1` (or `0` for an empty buffer). + voice, format, or reference audio. A `session.config` sent in the middle of + an utterance is rejected so no pending input is silently dropped and so a + split utterance cannot end up half in one voice and half in another. +- An utterance is the flush unit, not a linguistic one: it is one `input.done` + cycle. `utterance_index` counts those flushes across the connection, so it + tells you which `input.done` a frame belongs to. `sentence_index` counts the + TTS requests inside one flush and so pairs with `total_sentences`: with the + default `split_granularity=none` that is always `sentence_index: 0` of + `total_sentences: 1` (or `0` for an empty buffer), while `sentence` or + `clause` counts the linguistic units actually synthesized. - End the connection with `session.close`, or by closing the socket. An idle connection is still closed after the server's idle timeout, which now also applies to the gap between utterances. @@ -332,7 +339,15 @@ All REST API parameters are supported, plus: | Parameter | Type | Default | Description | | ----------- | ------ | --------- | ------------- | -| `stream_audio` | bool | false | Stream one or more PCM chunks for the buffered input over WebSocket | +| `stream_audio` | bool | false | Stream one or more PCM chunks for each TTS request over WebSocket | +| `split_granularity` | string | `"none"` | `"none"`: one request per `input.done`. `"sentence"`: split on `.!?` plus CJK `。！？…`, Indic danda `।॥`, and Arabic `؟`. `"clause"`: also split on `,;，；،؛`. | +| `seed` | integer | null | Forwarded to the speech engine for this session | + +ASCII punctuation only ends a unit when whitespace or `input.done` follows it, +and decimals (`3.14`), thousands separators (`1,000`), abbreviations (`Dr.`, +`e.g.`) and initials (`J. R.`) are not treated as boundaries. A punctuation run +and any closing quote or bracket stay with the unit they close, so `Wait...` +and `He said "Hello."` are one request each. ```bash DELETE /v1/audio/voices/{name} diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index b1221de6308..5dbe583146b 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -796,13 +796,21 @@ Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and ### Streaming WebSocket -The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`: +The `/v1/audio/speech/stream` endpoint accepts text incrementally and, by default, synthesizes the buffered text as one continuous request on `input.done`: ```bash python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine." python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1 ``` +For per-sentence audio (including Indic danda `।`) before `input.done`: + +```bash +python qwen3_tts/streaming_speech_client.py \ + --text "नमस्ते। कैसे हो?" \ + --split-granularity sentence +``` + `input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection: ```bash diff --git a/examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py b/examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py index e141a4c485f..793a5a6d53c 100644 --- a/examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py +++ b/examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py @@ -24,6 +24,11 @@ --text "Hello world. How are you? I am fine." \ --simulate-stt --stt-delay 0.1 + # Opt into per-sentence synthesis (Indic danda, CJK, Latin) + python streaming_speech_client.py \ + --text "नमस्ते। कैसे हो?" \ + --split-granularity sentence + # Receive JSON sidecar chunks with word-level timestamps python streaming_speech_client.py \ --text "Hello world. How are you?" \ @@ -73,8 +78,9 @@ def frame_basename(msg: dict) -> str: """Name a file after the utterance and sentence a frame belongs to. - Every utterance is sentence 0, so the utterance index is what keeps the - files of one connection from overwriting each other. + Every utterance uses `sentence_index` for units inside that flush. + With default `split_granularity=none` that index stays 0; sentence mode + uses 0, 1, ... so files from one connection do not overwrite each other. """ return f"utterance_{msg['utterance_index']:03d}_sentence_{msg['sentence_index']:03d}" @@ -368,6 +374,13 @@ def main(): ) parser.add_argument("--speed", type=float, default=1.0, help="Playback speed (0.25-4.0)") parser.add_argument("--max-new-tokens", type=int, default=None, help="Max tokens") + parser.add_argument("--seed", type=int, default=None, help="Sampling seed forwarded to the server") + parser.add_argument( + "--split-granularity", + default=None, + choices=["none", "sentence", "clause"], + help="Text split mode (default: server none = one request per input.done)", + ) # Base task options parser.add_argument("--ref-audio", default=None, help="Reference audio") @@ -409,6 +422,8 @@ def main(): "max_new_tokens", "ref_audio", "ref_text", + "seed", + "split_granularity", ]: val = getattr(args, key.replace("-", "_"), None) if val is not None: diff --git a/tests/entrypoints/openai_api/test_serving_speech_stream.py b/tests/entrypoints/openai_api/test_serving_speech_stream.py index 7c553af8f56..6414ab4ff8b 100644 --- a/tests/entrypoints/openai_api/test_serving_speech_stream.py +++ b/tests/entrypoints/openai_api/test_serving_speech_stream.py @@ -197,7 +197,7 @@ def test_session_config_rejected_while_input_is_buffered(self, mocker: MockerFix error = ws.receive_json() assert error["type"] == "error" - assert "while input is buffered" in error["message"] + assert "while an utterance is in progress" in error["message"] # The buffered text survives the rejected reconfiguration. ws.send_json({"type": "input.done"}) @@ -742,6 +742,8 @@ async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_ra config.speaker_embedding = None config.stream_audio = True config.word_timestamps = False + config.seed = None + config.non_streaming_mode = None with pytest.raises(WebSocketDisconnect): asyncio.run( @@ -758,6 +760,163 @@ async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_ra assert websocket.send_json.await_count == 2 +class TestWebSocketSentenceSplitting: + def test_sentence_granularity_emits_one_request_per_sentence(self, mocker: MockerFixture): + app, speech_service = _build_test_app(mocker=mocker) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json( + { + "type": "session.config", + "voice": "Vivian", + "split_granularity": "sentence", + } + ) + ws.send_json({"type": "input.text", "text": "Hello world. How are you? "}) + ws.send_json({"type": "input.done"}) + + first = ws.receive_json() + assert first["sentence_index"] == 0 + assert first["sentence_text"] == "Hello world." + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + + second = ws.receive_json() + assert second["sentence_index"] == 1 + assert second["sentence_text"] == "How are you?" + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + + assert ws.receive_json() == { + "type": "session.done", + "utterance_index": 0, + "total_sentences": 2, + } + + assert speech_service._generate_audio_bytes.await_count == 2 + assert [call.args[0].input for call in speech_service._generate_audio_bytes.await_args_list] == [ + "Hello world.", + "How are you?", + ] + + def test_sentence_granularity_emits_before_input_done(self, mocker: MockerFixture): + app, speech_service = _build_test_app(mocker=mocker) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json( + { + "type": "session.config", + "voice": "Vivian", + "split_granularity": "sentence", + } + ) + ws.send_json({"type": "input.text", "text": "First sentence. "}) + + start = ws.receive_json() + assert start["sentence_text"] == "First sentence." + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + assert speech_service._generate_audio_bytes.await_count == 1 + + ws.send_json({"type": "input.done"}) + assert ws.receive_json() == { + "type": "session.done", + "utterance_index": 0, + "total_sentences": 1, + } + + def test_indic_danda_splits_without_latin_period(self, mocker: MockerFixture): + app, speech_service = _build_test_app(mocker=mocker) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json( + { + "type": "session.config", + "voice": "Vivian", + "language": "Auto", + "split_granularity": "sentence", + } + ) + ws.send_json({"type": "input.text", "text": "नमस्ते। कैसे हो?"}) + ws.send_json({"type": "input.done"}) + + first = ws.receive_json() + assert first["sentence_text"] == "नमस्ते।" + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + + second = ws.receive_json() + assert second["sentence_text"] == "कैसे हो?" + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + assert ws.receive_json()["total_sentences"] == 2 + + def test_session_config_rejected_after_a_split_unit_was_emitted(self, mocker: MockerFixture): + """The splitter buffer is empty here, but the utterance is still open.""" + app, speech_service = _build_test_app(mocker=mocker) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json( + { + "type": "session.config", + "voice": "Vivian", + "split_granularity": "sentence", + } + ) + ws.send_json({"type": "input.text", "text": "First sentence. "}) + ws.receive_json() + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + + ws.send_json({"type": "session.config", "voice": "Serena"}) + error = ws.receive_json() + assert error["type"] == "error" + assert "while an utterance is in progress" in error["message"] + + ws.send_json({"type": "input.text", "text": "Second sentence. "}) + assert ws.receive_json()["sentence_index"] == 1 + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + + ws.send_json({"type": "input.done"}) + assert ws.receive_json() == { + "type": "session.done", + "utterance_index": 0, + "total_sentences": 2, + } + + # Reconfiguration is allowed again once the utterance closed. + ws.send_json({"type": "session.config", "voice": "Serena"}) + ws.send_json({"type": "input.text", "text": "Third."}) + ws.send_json({"type": "input.done"}) + ws.receive_json() + ws.receive_bytes() + ws.receive_json() + ws.receive_json() + + voices = [call.args[0].voice for call in speech_service._generate_audio_bytes.await_args_list] + assert voices == ["Vivian", "Vivian", "Serena"] + + def test_seed_is_forwarded_to_speech_request(self, mocker: MockerFixture): + app, speech_service = _build_test_app(mocker=mocker) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json({"type": "session.config", "voice": "Vivian", "seed": 42}) + ws.send_json({"type": "input.text", "text": "Hello."}) + ws.send_json({"type": "input.done"}) + ws.receive_json() + ws.receive_bytes() + ws.receive_json() + ws.receive_json() + + assert speech_service._generate_audio_bytes.await_args_list[0].args[0].seed == 42 + + class TestGeneratePcmChunksContract: """Guard: _generate_pcm_chunks must exist on OmniOpenAIServingSpeech. diff --git a/tests/entrypoints/openai_api/test_speech_text_splitter.py b/tests/entrypoints/openai_api/test_speech_text_splitter.py new file mode 100644 index 00000000000..ca2eb76d87e --- /dev/null +++ b/tests/entrypoints/openai_api/test_speech_text_splitter.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest + +from vllm_omni.entrypoints.openai import speech_text_splitter +from vllm_omni.entrypoints.openai.speech_text_splitter import SpeechTextSplitter, extract_complete_units + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestExtractCompleteUnits: + def test_ascii_sentence_needs_whitespace_or_flush(self): + units, remainder, _ = extract_complete_units("Hello world.", frozenset(".!?。"), flush=False) + assert units == [] + assert remainder == "Hello world." + + units, remainder, _ = extract_complete_units("Hello world.", frozenset(".!?。"), flush=True) + assert units == ["Hello world."] + assert remainder == "" + + def test_ascii_sentence_with_trailing_space(self): + units, remainder, _ = extract_complete_units("Hello world. More", frozenset(".!?"), flush=False) + assert units == ["Hello world."] + assert remainder == "More" + + def test_decimal_is_not_a_sentence_boundary(self): + units, remainder, _ = extract_complete_units("It costs 3.14 dollars. ", frozenset(".!?"), flush=False) + assert units == ["It costs 3.14 dollars."] + assert remainder == "" + + def test_thousands_separator_is_not_a_clause_boundary(self): + units, _, _ = extract_complete_units("It is 1,000 rupees, sir. ", frozenset(".!?,;"), flush=False) + assert units == ["It is 1,000 rupees,", "sir."] + + def test_punctuation_run_stays_one_unit(self): + units, _, _ = extract_complete_units("Wait... Really?! Yes.", frozenset(".!?"), flush=True) + assert units == ["Wait...", "Really?!", "Yes."] + + def test_closing_quote_stays_with_its_sentence(self): + units, _, _ = extract_complete_units('He said "Hello." Bye.', frozenset(".!?"), flush=True) + assert units == ['He said "Hello."', "Bye."] + + def test_abbreviations_and_initials_do_not_split(self): + units, _, _ = extract_complete_units("Dr. Smith, e.g. J. R. Tolkien. Next.", frozenset(".!?"), flush=True) + assert units == ["Dr. Smith, e.g. J. R. Tolkien.", "Next."] + + def test_long_token_before_period_is_not_an_abbreviation(self): + units, _, _ = extract_complete_units("We landed in Washington. Next.", frozenset(".!?"), flush=True) + assert units == ["We landed in Washington.", "Next."] + + def test_abbreviation_lookback_is_bounded(self, monkeypatch): + # One WebSocket frame can reach ~128 KiB; an unbounded lookback walked + # the whole prefix at every `.`, quadratic in the frame size. + checks = 0 + is_token_char = speech_text_splitter._is_token_char + + def counting_is_token_char(ch): + nonlocal checks + checks += 1 + return is_token_char(ch) + + monkeypatch.setattr(speech_text_splitter, "_is_token_char", counting_is_token_char) + text = "a." * 20000 + units, _, _ = extract_complete_units(text, frozenset(".!?"), flush=True) + assert units == [text] + assert checks <= len(text) * (speech_text_splitter._MAX_ABBREVIATION_LEN + 1) + + def test_cjk_period_completes_once_the_next_char_is_known(self): + # The trailing 。 waits: a closing delimiter could still follow it. + units, remainder, _ = extract_complete_units("你好。世界。", frozenset(".!?。！？"), flush=False) + assert units == ["你好。"] + assert remainder == "世界。" + + units, remainder, _ = extract_complete_units("你好。世界。", frozenset(".!?。！？"), flush=True) + assert units == ["你好。", "世界。"] + assert remainder == "" + + def test_indic_danda_and_question_mark(self): + units, remainder, _ = extract_complete_units("नमस्ते। कैसे हो?", frozenset(".!?।॥؟"), flush=True) + assert units == ["नमस्ते।", "कैसे हो?"] + assert remainder == "" + + def test_scan_cursor_skips_already_examined_text(self): + units, remainder, scan = extract_complete_units("no terminator yet", frozenset(".!?"), flush=False) + assert units == [] + assert remainder == "no terminator yet" + assert scan == len(remainder) + + units, remainder, scan = extract_complete_units( + "no terminator yet. Next", frozenset(".!?"), flush=False, scan_from=scan + ) + assert units == ["no terminator yet."] + assert remainder == "Next" + + +class TestSpeechTextSplitter: + def test_none_buffers_until_flush(self): + splitter = SpeechTextSplitter("none") + assert splitter.feed("Hello world. ") == [] + assert splitter.flush() == ["Hello world."] + assert splitter.flush() == [] + + def test_sentence_mode_incremental_words(self): + splitter = SpeechTextSplitter("sentence") + assert splitter.feed("Hello ") == [] + assert splitter.feed("world. ") == ["Hello world."] + assert splitter.feed("How are you?") == [] + assert splitter.flush() == ["How are you?"] + + def test_sentence_mode_resumes_scan_across_feeds(self): + splitter = SpeechTextSplitter("sentence") + assert splitter.feed("It costs 3") == [] + assert splitter.feed(".") == [] + assert splitter.feed("14 dollars. ") == ["It costs 3.14 dollars."] + + def test_clause_mode_splits_on_comma(self): + splitter = SpeechTextSplitter("clause") + assert splitter.feed("Hello, world. ") == ["Hello,", "world."] + + def test_clause_mode_splits_on_arabic_semicolon(self): + splitter = SpeechTextSplitter("clause") + assert splitter.feed("مرحبا؛ كيف حالك") == ["مرحبا؛"] + + def test_empty_flush(self): + splitter = SpeechTextSplitter("sentence") + assert splitter.flush() == [] diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index f204ae209ae..2446114df17 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -578,6 +578,22 @@ class StreamingSpeechSessionConfig(BaseModel): "frames (existing behavior)." ), ) + seed: int | None = Field( + default=None, + ge=_INT64_MIN, + le=_INT64_MAX, + description="Random seed forwarded to /v1/audio/speech for this session.", + ) + split_granularity: Literal["none", "sentence", "clause"] = Field( + default="none", + description=( + "How incoming input.text is segmented before TTS. 'none' (default) " + "buffers until input.done and runs one request, matching the " + "long-form timbre-continuity path. 'sentence' emits a request at " + "each sentence boundary (Latin .!? plus CJK/Indic/Arabic marks). " + "'clause' also splits on commas/semicolons for lower TTFA." + ), + ) @model_validator(mode="after") def validate_streaming_constraints(self) -> "StreamingSpeechSessionConfig": diff --git a/vllm_omni/entrypoints/openai/serving_speech_stream.py b/vllm_omni/entrypoints/openai/serving_speech_stream.py index 0f883f17920..32e9fa40d01 100644 --- a/vllm_omni/entrypoints/openai/serving_speech_stream.py +++ b/vllm_omni/entrypoints/openai/serving_speech_stream.py @@ -3,8 +3,10 @@ """WebSocket handler for streaming text input TTS. -Accepts text incrementally via WebSocket, buffers it until input.done, and -generates audio once for the buffered input using the existing TTS pipeline. +Accepts text incrementally via WebSocket. By default (split_granularity=none) +it buffers until input.done and generates audio once for the buffered input. +Opting into split_granularity=sentence or clause emits a TTS request at each +detected boundary so incremental STT/LLM clients can hear audio before flush. input.done is a flush, not a close: it ends the current utterance and the connection stays open, so the next utterance reuses the same connection @@ -13,9 +15,10 @@ The session config is sticky across flushes and can be replaced by sending another session.config between utterances. -"Utterance" here names the flush unit rather than any linguistic unit: it is -whatever text the client had buffered when it sent input.done, from a word to -several paragraphs, synthesized as a single request. +"Utterance" here names the flush unit. With split_granularity=none it is +whatever text was buffered when input.done arrived, synthesized as one +request. With sentence/clause splitting it is still one input.done cycle, +but sentence_index counts linguistic units inside that flush. Protocol: Client -> Server: @@ -32,11 +35,10 @@ {"type": "audio.done", "utterance_index": 0, "sentence_index": 0} {"type": "session.done", "utterance_index": 0, "total_sentences": N} {"type": "error", "message": "..."} - # session.done ends the flushed utterance, not the connection. An - # utterance is just the flush unit: whatever text was buffered when - # input.done arrived, of any length. utterance_index counts those - # flushes across the connection, while sentence_index counts within - # one of them and so pairs with total_sentences. + # session.done ends the flushed utterance, not the connection. + # utterance_index counts input.done flushes. sentence_index counts + # TTS requests inside one flush (always 0 of 1 when split_granularity + # is none; 0..N-1 when sentence/clause splitting is enabled). Server -> Client (when word_timestamps=true): {"type": "audio.start", "utterance_index": 0, "sentence_index": 0, @@ -64,6 +66,7 @@ StreamingSpeechSessionConfig, ) from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech +from vllm_omni.entrypoints.openai.speech_text_splitter import SpeechTextSplitter from vllm_omni.utils.forced_aligner import extract_word_timestamps logger = init_logger(__name__) @@ -79,11 +82,10 @@ class OmniStreamingSpeechHandler: """Handles WebSocket sessions for streaming text-input TTS. - A connection carries one or more utterances. Text arrives incrementally, - is buffered until input.done, and audio is generated once for the - buffered input using the existing OmniOpenAIServingSpeech pipeline. The - connection outlives each utterance so a client can keep synthesizing on - it without reconnecting. + A connection carries one or more utterances. Text arrives incrementally. + With split_granularity=none it is buffered until input.done and synthesized + as one request. sentence/clause modes emit a request at each boundary, + including before input.done. The connection outlives each utterance. Args: speech_service: The existing TTS serving instance (reused for @@ -105,17 +107,20 @@ def __init__( async def handle_session(self, websocket: WebSocket) -> None: """Main loop for a single WebSocket connection. - Serves any number of utterances: text is buffered until input.done, - which flushes it as one TTS request and then leaves the connection - open for the next one. Rejecting a message is only fatal before the - first valid session.config; afterwards the error is reported and the - connection survives. + Serves any number of utterances. input.done always ends a flush; + split_granularity controls whether that flush is one TTS request or + several sentence/clause requests. """ await websocket.accept() config: StreamingSpeechSessionConfig | None = None - text_parts: list[str] = [] + splitter = SpeechTextSplitter("none") utterance_index = 0 + sentence_index = 0 + # An utterance is in progress from the first input.text until the next + # input.done. Buffered text is not enough to detect it: sentence/clause + # splitting can leave the buffer empty after emitting a unit. + utterance_open = False try: while True: @@ -140,10 +145,10 @@ async def handle_session(self, websocket: WebSocket) -> None: msg_type = msg.get("type") if msg_type == "session.config": - if text_parts: + if utterance_open: await self._send_error( websocket, - "session.config cannot be applied while input is buffered; send input.done first", + "session.config cannot be applied while an utterance is in progress; send input.done first", ) continue new_config = await self._build_config(websocket, msg) @@ -152,6 +157,7 @@ async def handle_session(self, websocket: WebSocket) -> None: return # Error already sent, connection closing continue # Keep serving with the previous config config = new_config + splitter = SpeechTextSplitter(config.split_granularity) elif config is None: await self._send_error( @@ -165,32 +171,39 @@ async def handle_session(self, websocket: WebSocket) -> None: if not isinstance(text, str): await self._send_error(websocket, "input.text requires a string value") continue - text_parts.append(text) + if text: + utterance_open = True + for unit in splitter.feed(text): + await self._generate_and_send( + websocket, + config, + unit, + utterance_index=utterance_index, + sentence_index=sentence_index, + ) + sentence_index += 1 elif msg_type == "input.done": - full_text = "".join(text_parts).strip() - text_parts.clear() - total_sentences = 0 - if full_text: - # However long the buffered text is, the pipeline takes - # it as one request, so every flush is sentence 0 of 1. + for unit in splitter.flush(): await self._generate_and_send( websocket, config, - full_text, + unit, utterance_index=utterance_index, - sentence_index=0, + sentence_index=sentence_index, ) - total_sentences = 1 + sentence_index += 1 await websocket.send_json( { "type": "session.done", "utterance_index": utterance_index, - "total_sentences": total_sentences, + "total_sentences": sentence_index, } ) utterance_index += 1 + sentence_index = 0 + utterance_open = False elif msg_type == "session.close": await websocket.close() @@ -310,6 +323,7 @@ async def _generate_and_send( speaker_embedding=config.speaker_embedding, stream=config.stream_audio, word_timestamps=config.word_timestamps, + seed=config.seed, ) start_payload = { diff --git a/vllm_omni/entrypoints/openai/speech_text_splitter.py b/vllm_omni/entrypoints/openai/speech_text_splitter.py new file mode 100644 index 00000000000..1f25a8dba2e --- /dev/null +++ b/vllm_omni/entrypoints/openai/speech_text_splitter.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Incremental text splitting for WebSocket TTS input. + +Default serving still synthesizes one request per ``input.done`` (see +``split_granularity='none'``). Opting into ``sentence`` or ``clause`` restores +per-boundary requests so incremental STT/LLM clients can start audio before +the full utterance arrives. + +Terminators cover Latin, CJK, Indic danda, and Arabic question/semicolon marks +so Hindi and similar scripts are not stuck waiting for ``.`` / ``?``. +""" + +from __future__ import annotations + +# ASCII punctuation is ambiguous (abbreviations, decimals, thousands +# separators), so it only closes a unit when whitespace or a flush follows. +# Script-specific marks below are unambiguous and close a unit as soon as the +# character after them is known. +_ASCII_SENTENCE_TERMINATORS = frozenset(".!?") +_ASCII_CLAUSE_TERMINATORS = _ASCII_SENTENCE_TERMINATORS | frozenset(",;") +_SCRIPT_SENTENCE_TERMINATORS = frozenset( + { + "。", + "！", + "？", + "…", + "।", # Devanagari danda + "॥", # Devanagari double danda + "؟", # Arabic question mark + } +) +_SCRIPT_CLAUSE_TERMINATORS = _SCRIPT_SENTENCE_TERMINATORS | { + "，", + "；", + "،", # Arabic comma + "؛", # Arabic semicolon +} + +_SENTENCE_TERMINATORS = _ASCII_SENTENCE_TERMINATORS | _SCRIPT_SENTENCE_TERMINATORS +_CLAUSE_TERMINATORS = _ASCII_CLAUSE_TERMINATORS | _SCRIPT_CLAUSE_TERMINATORS + +# Trailing characters that belong to the unit they close, so `Wait...` and +# `He said "Hello."` stay one TTS request instead of several. +_CLOSING_DELIMITERS = frozenset("\"'”’»)]}』」）】") +_RUN_CHARS = _CLAUSE_TERMINATORS | _CLOSING_DELIMITERS + +# Abbreviations that end in `.` mid-sentence. Single letters are covered +# separately so initials such as `J. R. R.` do not split either. +_ABBREVIATIONS = frozenset( + { + "dr", + "mr", + "mrs", + "ms", + "prof", + "sr", + "jr", + "st", + "mt", + "fig", + "no", + "vs", + "etc", + "approx", + "dept", + "inc", + "ltd", + "co", + "al", + "e.g", + "i.e", + "u.s", + "u.k", + "a.m", + "p.m", + } +) +_MAX_ABBREVIATION_LEN = max(len(abbrev) for abbrev in _ABBREVIATIONS) + + +def _is_token_char(ch: str) -> bool: + return ch.isalnum() or ch == "." + + +def _is_numeric_separator(buffer: str, index: int) -> bool: + """True for the `.`/`,` inside `3.14` or `1,000`.""" + if buffer[index] not in ".,": + return False + if index == 0 or not buffer[index - 1].isdigit(): + return False + return index + 1 < len(buffer) and buffer[index + 1].isdigit() + + +def _is_abbreviation(buffer: str, index: int) -> bool: + """True for the `.` of `Dr.`, `e.g.`, or an initial like `J.`.""" + if buffer[index] != ".": + return False + # Bounded so a frame of `a.a.a...` cannot walk the whole prefix at every + # `.`. Every abbreviation fits inside the limit, so hitting it can only + # rule a token out, never produce a false positive. + limit = max(0, index - _MAX_ABBREVIATION_LEN) + start = index + while start > limit and _is_token_char(buffer[start - 1]): + start -= 1 + if start > 0 and _is_token_char(buffer[start - 1]): + return False + token = buffer[start:index] + if not token: + return False + if len(token) == 1 and token.isalpha(): + return True + return token.lower().strip(".") in _ABBREVIATIONS + + +def extract_complete_units( + buffer: str, + terminators: frozenset[str], + *, + flush: bool, + scan_from: int = 0, +) -> tuple[list[str], str, int]: + """Split ``buffer`` into complete units. + + Returns ``(units, remainder, next_scan)``. ``next_scan`` is an offset into + ``remainder``: everything before it has been examined already, so a caller + appending more text can pass it back as ``scan_from`` instead of rescanning + the whole buffer on every ``input.text`` message. + """ + units: list[str] = [] + last_split = 0 + i = max(scan_from, 0) + length = len(buffer) + while i < length: + ch = buffer[i] + if ch not in terminators or _is_numeric_separator(buffer, i) or _is_abbreviation(buffer, i): + i += 1 + continue + + # Absorb the punctuation/closing-delimiter run so `Wait...` emits once. + run_end = i + 1 + while run_end < length and buffer[run_end] in _RUN_CHARS: + run_end += 1 + + if run_end >= length and not flush: + # The run may still grow; re-examine this terminator next feed. + return units, buffer[last_split:], i - last_split + + ascii_terminator = ch in _ASCII_CLAUSE_TERMINATORS + if run_end < length and ascii_terminator and not buffer[run_end].isspace(): + i = run_end + continue + + piece = buffer[last_split:run_end].strip() + if piece: + units.append(piece) + i = run_end + while i < length and buffer[i].isspace(): + i += 1 + last_split = i + + remainder = buffer[last_split:] + if flush: + tail = remainder.strip() + if tail: + units.append(tail) + return units, "", 0 + return units, remainder, len(remainder) + + +class SpeechTextSplitter: + """Stateful splitter used by one WebSocket utterance.""" + + def __init__(self, granularity: str = "none") -> None: + if granularity not in ("none", "sentence", "clause"): + raise ValueError(f"Unsupported split_granularity: {granularity!r}") + self.granularity = granularity + self._buf = "" + self._scan = 0 + + def _terminators(self) -> frozenset[str] | None: + if self.granularity == "none": + return None + if self.granularity == "clause": + return _CLAUSE_TERMINATORS + return _SENTENCE_TERMINATORS + + def feed(self, text: str) -> list[str]: + self._buf += text + terminators = self._terminators() + if terminators is None: + return [] + units, self._buf, self._scan = extract_complete_units(self._buf, terminators, flush=False, scan_from=self._scan) + return units + + def flush(self) -> list[str]: + terminators = self._terminators() + if terminators is None: + piece = self._buf.strip() + self._buf = "" + self._scan = 0 + return [piece] if piece else [] + units, self._buf, self._scan = extract_complete_units(self._buf, terminators, flush=True, scan_from=self._scan) + return units From 43b697541a93259700605af7441814a90c600585 Mon Sep 17 00:00:00 2001 From: Ziming Wang <125807850+ZenAlexa@users.noreply.github.com> Date: Wed, 9 Sep 2026 11:31:42 +0800 Subject: [PATCH 16/93] [Bugfix][Frontend] Clear the P0 multimodal cache through the renderer (#7003) Signed-off-by: ZenAlexa --- .../test_async_omni_pause_sleep_routing.py | 21 +++++++++++++++++++ vllm_omni/entrypoints/async_omni.py | 16 +++----------- 2 files changed, 24 insertions(+), 13 deletions(-) diff --git a/tests/entrypoints/test_async_omni_pause_sleep_routing.py b/tests/entrypoints/test_async_omni_pause_sleep_routing.py index 9600b990d23..38ec4016384 100644 --- a/tests/entrypoints/test_async_omni_pause_sleep_routing.py +++ b/tests/entrypoints/test_async_omni_pause_sleep_routing.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Unit tests for AR EngineCore vs diffusion worker pause/sleep routing.""" from __future__ import annotations @@ -7,6 +10,8 @@ from unittest.mock import AsyncMock, call import pytest +from vllm.renderers import BaseRenderer +from vllm.v1.engine.input_processor import InputProcessor from vllm_omni.diffusion.data import CuMemTag, OmniACK from vllm_omni.entrypoints.async_omni import AsyncOmni @@ -37,6 +42,22 @@ def _make_omni(*, stage_types: list[str]) -> AsyncOmni: return omni +@pytest.mark.cpu +def test_reset_mm_cache_routes_through_renderer(mocker): + async def run() -> None: + omni = object.__new__(AsyncOmni) + renderer = mocker.create_autospec(BaseRenderer, instance=True) + input_processor = mocker.create_autospec(InputProcessor, instance=True) + input_processor.renderer = renderer + omni.input_processor = input_processor + + await omni.reset_mm_cache() + + renderer.clear_mm_cache_async.assert_awaited_once_with() + + asyncio.run(run()) + + @pytest.mark.cpu def test_split_stage_ids_by_type(): omni = _make_omni(stage_types=["llm", "diffusion", "llm"]) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 4026907e1ff..948f7e53a82 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -1416,19 +1416,9 @@ async def reset_mm_cache(self) -> None: ``EngineCore.sleep(level>=1)`` already clears the P1 receiver cache. Clearing P0 avoids hash-only follow-up requests after that reset. """ - processor = getattr(self, "input_processor", None) - if processor is None: - processor = getattr(self.engine, "input_processor", None) - cache = getattr(processor, "mm_processor_cache", None) - if cache is None: - logger.debug("[AsyncOmni] reset_mm_cache: no frontend mm_processor_cache") - return - for name in ("clear", "reset", "clear_cache"): - fn = getattr(cache, name, None) - if callable(fn): - fn() - return - logger.debug("[AsyncOmni] reset_mm_cache: cache has no clear/reset method") + renderer = self.renderer + if renderer is not None: + await renderer.clear_mm_cache_async() async def reset_encoder_cache(self) -> None: """Reset the encoder cache for all stages. From 358b7d79c1994fc17c62cb644b601e584c77aa10 Mon Sep 17 00:00:00 2001 From: Jim Ban <77719403+BANANASJIM@users.noreply.github.com> Date: Tue, 8 Sep 2026 20:31:54 -0700 Subject: [PATCH 17/93] [Bugfix][Frontend] Enforce image pixel limits for video input references (#6963) Signed-off-by: BANANASJIM --- .../openai_api/test_video_api_utils.py | 50 ++++++++++++++++++- .../openai_api/test_video_server.py | 34 ++++++++++++- vllm_omni/entrypoints/openai/api_server.py | 23 +++++++-- .../entrypoints/openai/video_api_utils.py | 25 +++++++++- 4 files changed, 124 insertions(+), 8 deletions(-) diff --git a/tests/entrypoints/openai_api/test_video_api_utils.py b/tests/entrypoints/openai_api/test_video_api_utils.py index 9f643d8caa9..bf395b8136a 100644 --- a/tests/entrypoints/openai_api/test_video_api_utils.py +++ b/tests/entrypoints/openai_api/test_video_api_utils.py @@ -23,9 +23,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -def _png_bytes() -> bytes: +def _png_bytes(size: tuple[int, int] = (2, 1)) -> bytes: buffer = BytesIO() - Image.new("RGB", (2, 1), color=(12, 34, 56)).save(buffer, format="PNG") + Image.new("RGB", size, color=(12, 34, 56)).save(buffer, format="PNG") return buffer.getvalue() @@ -122,6 +122,52 @@ def _unexpected_http_client(*args, **kwargs): assert image.mode == "RGB" +def test_decode_image_bytes_rejects_image_over_pixel_limit_before_conversion(monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_IMAGE_PIXELS", 100) + + def _unexpected_convert(*args, **kwargs): + pytest.fail("over-limit images must be rejected before conversion") + + monkeypatch.setattr(Image.Image, "convert", _unexpected_convert) + + with pytest.raises( + InvalidInputReferenceError, + match=r"Image dimensions 20x20 \(400 pixels\) exceed the maximum of 100 pixels", + ): + video_api_utils._decode_image_bytes(_png_bytes((20, 20)), source="image reference") + + +def test_decode_image_bytes_maps_pillow_pixel_limit_error(monkeypatch): + monkeypatch.setattr(Image, "MAX_IMAGE_PIXELS", 100) + + with pytest.raises(InvalidInputReferenceError, match="decoder pixel limit") as exc_info: + video_api_utils._decode_image_bytes(_png_bytes((20, 20)), source="image reference") + + assert isinstance(exc_info.value.__cause__, Image.DecompressionBombError) + + +@pytest.mark.parametrize("max_pixels", [0, 400]) +def test_decode_image_bytes_allows_disabled_or_inclusive_pixel_limit(monkeypatch, max_pixels): + monkeypatch.setattr(envs, "VLLM_MAX_IMAGE_PIXELS", max_pixels) + + image = video_api_utils._decode_image_bytes(_png_bytes((20, 20)), source="image reference") + + assert image.size == (20, 20) + assert image.mode == "RGB" + + +@pytest.mark.asyncio +async def test_decode_input_reference_preserves_image_pixel_limit_error(monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_IMAGE_PIXELS", 100) + + with pytest.raises(InvalidInputReferenceError, match="VLLM_MAX_IMAGE_PIXELS"): + await video_api_utils.decode_input_reference( + image_reference=None, + video_reference=None, + input_reference_bytes=_png_bytes((20, 20)), + ) + + def _install_fake_video_mux(monkeypatch, mux_calls): def _fake_mux_video_audio_bytes(frames, audio, fps, audio_sample_rate, video_codec_options=None): mux_calls.append( diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index b13dfe42e5d..1b18a9ae518 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -1630,7 +1630,8 @@ def test_generic_video_model_rejects_mixed_image_and_video_references(test_clien assert "does not support mixed image and video" in response.json()["detail"].lower() -def test_h3_multipart_rejects_bmp_image_reference(test_client): +def test_h3_multipart_rejects_bmp_image_reference(test_client, monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_IMAGE_PIXELS", 100) image = Image.new("RGB", (64, 64), color="blue") image_buffer = io.BytesIO() image.save(image_buffer, format="BMP") @@ -1646,6 +1647,37 @@ def test_h3_multipart_rejects_bmp_image_reference(test_client): assert "must use jpg" in response.json()["detail"].lower() +@pytest.mark.parametrize("field", ["input_reference", "input_references"]) +def test_h3_multipart_rejects_image_over_pixel_limit(field, test_client, monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_IMAGE_PIXELS", 100) + test_client.app.state.openai_serving_video._engine_client.model_class_name = "MiniMaxH3Pipeline" + + response = test_client.post( + "/v1/videos/sync", + data={"prompt": "reject oversized image", "extra_params": '{"task":"ref2va"}'}, + files=[(field, ("reference.png", _make_test_image_bytes((20, 20)), "image/png"))], + ) + + assert response.status_code == 400 + assert "VLLM_MAX_IMAGE_PIXELS" in response.json()["detail"] + + +@pytest.mark.parametrize("field", ["input_reference", "input_references"]) +def test_h3_multipart_maps_pillow_pixel_limit_error(field, test_client, monkeypatch): + monkeypatch.setattr(envs, "VLLM_MAX_IMAGE_PIXELS", 0) + monkeypatch.setattr(Image, "MAX_IMAGE_PIXELS", 100) + test_client.app.state.openai_serving_video._engine_client.model_class_name = "MiniMaxH3Pipeline" + + response = test_client.post( + "/v1/videos/sync", + data={"prompt": "reject decoder bomb", "extra_params": '{"task":"ref2va"}'}, + files=[(field, ("reference.png", _make_test_image_bytes((20, 20)), "image/png"))], + ) + + assert response.status_code == 400 + assert "decoder pixel limit" in response.json()["detail"] + + @pytest.mark.asyncio async def test_h3_upload_limit_checks_declared_size_before_read(): class OversizedUpload: diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index f97207bbc9c..d976eb18170 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -161,6 +161,8 @@ from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request from vllm_omni.entrypoints.openai.video_api_utils import ( VideoFrames, + _decode_image_bytes, + _validate_image_pixel_limit, decode_audio_url, decode_input_reference, ) @@ -3438,6 +3440,15 @@ def _validate_minimax_h3_image_payload( try: with Image.open(io.BytesIO(payload)) as image: image_format = str(image.format or "").lower() + if image_format in MINIMAX_H3_REFERENCE_IMAGE_FORMATS: + _validate_image_pixel_limit(image) + except InvalidInputReferenceError as exc: + raise HTTPException(HTTPStatus.BAD_REQUEST.value, detail=str(exc)) from exc + except Image.DecompressionBombError as exc: + raise HTTPException( + HTTPStatus.BAD_REQUEST.value, + detail=f"Invalid uploaded image reference: {filename}; image exceeds the decoder pixel limit.", + ) from exc except (OSError, ValueError) as exc: if allow_non_image: return @@ -3478,10 +3489,14 @@ async def _persist_uploaded_media_references( if kind == "image": try: _validate_minimax_h3_image_payload(payload, filename=upload.filename) - with Image.open(io.BytesIO(payload)) as image: - images.append(image.convert("RGB")) - except (OSError, ValueError) as exc: - raise HTTPException(400, detail=f"Invalid uploaded image reference: {upload.filename}") from exc + images.append( + _decode_image_bytes( + payload, + source=f"uploaded image reference: {upload.filename}", + ) + ) + except InvalidInputReferenceError as exc: + raise HTTPException(HTTPStatus.BAD_REQUEST.value, detail=str(exc)) from exc continue suffix = Path(upload.filename or "").suffix.lower() if kind == "video" and suffix and suffix not in MINIMAX_H3_REFERENCE_VIDEO_SUFFIXES: diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index 09a3d284d83..35a39b2eacc 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -71,6 +71,21 @@ def __init__( self.source_path = source_path +class _ImagePixelLimitError(InvalidInputReferenceError): + """An image exceeded a configured or decoder-enforced pixel limit.""" + + +def _validate_image_pixel_limit(image: Image.Image) -> None: + width, height = image.size + max_pixels = envs.VLLM_MAX_IMAGE_PIXELS + if max_pixels > 0 and width * height > max_pixels: + raise _ImagePixelLimitError( + f"Image dimensions {width}x{height} ({width * height} pixels) exceed " + f"the maximum of {max_pixels} pixels. Set " + f"VLLM_MAX_IMAGE_PIXELS to increase this limit." + ) + + def positive_float(value: Any) -> float | None: if value is None: return None @@ -87,7 +102,13 @@ def positive_float(value: Any) -> float | None: def _decode_image_bytes(image_bytes: bytes, *, source: str) -> Image.Image: try: - return Image.open(BytesIO(image_bytes)).convert("RGB") + with Image.open(BytesIO(image_bytes)) as image: + _validate_image_pixel_limit(image) + return image.convert("RGB") + except _ImagePixelLimitError: + raise + except Image.DecompressionBombError as exc: + raise _ImagePixelLimitError(f"Invalid {source}: image exceeds the decoder pixel limit.") from exc except (UnidentifiedImageError, OSError, ValueError) as exc: raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid image.") from exc @@ -158,6 +179,8 @@ def _decode_media_bytes( ) -> Image.Image | VideoFrames: try: return _decode_image_bytes(media_bytes, source=source) + except _ImagePixelLimitError: + raise except InvalidInputReferenceError: try: return _decode_video_bytes( From f2e30f46b6a3e9bfb00468d7f2596852bffd37de Mon Sep 17 00:00:00 2001 From: Allen Wu <85376543+EchoHayate@users.noreply.github.com> Date: Wed, 9 Sep 2026 11:32:05 +0800 Subject: [PATCH 18/93] [Bugfix][TTS] Isolate shared Higgs v3 reference encode from request cancellation (#7076) Signed-off-by: Allen Wu Co-authored-by: TRAE CLI --- .../openai_api/test_higgs_audio_v3_adapter.py | 129 ++++++++++++++++++ .../openai/tts_adapters/higgs_audio_v3.py | 18 ++- 2 files changed, 140 insertions(+), 7 deletions(-) create mode 100644 tests/entrypoints/openai_api/test_higgs_audio_v3_adapter.py diff --git a/tests/entrypoints/openai_api/test_higgs_audio_v3_adapter.py b/tests/entrypoints/openai_api/test_higgs_audio_v3_adapter.py new file mode 100644 index 00000000000..7992af25fb4 --- /dev/null +++ b/tests/entrypoints/openai_api/test_higgs_audio_v3_adapter.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import asyncio +import gc +import threading +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext +from vllm_omni.entrypoints.openai.tts_adapters.higgs_audio_v3 import HiggsAudioV3Adapter + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_ref_code_inflight_survives_request_cancellation(): + adapter = HiggsAudioV3Adapter(SpeechServingContext(server=SimpleNamespace())) + encode_started = threading.Event() + release_encode = threading.Event() + calls = 0 + expected_codes = torch.arange(16, dtype=torch.long).reshape(2, 8) + 1 + + def encode_reference_audio(_wav, _sr): + nonlocal calls + calls += 1 + encode_started.set() + assert release_encode.wait(timeout=5) + return torch.arange(16, dtype=torch.long).reshape(2, 8) + + def apply_delay_pattern(codes): + return codes + 1 + + async def resolve(): + return await adapter._resolve_higgs_audio_v3_ref_codes( + "ref-a", + object(), + 24000, + encode_reference_audio, + apply_delay_pattern, + ) + + async def run(): + creator = asyncio.create_task(resolve()) + try: + assert await asyncio.wait_for(asyncio.to_thread(encode_started.wait), timeout=5) + + cancelled_waiter = asyncio.create_task(resolve()) + survivor = asyncio.create_task(resolve()) + await asyncio.sleep(0) + + cancelled_waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled_waiter + + creator.cancel() + with pytest.raises(asyncio.CancelledError): + await creator + + release_encode.set() + survivor_result = await asyncio.wait_for(survivor, timeout=5) + await asyncio.sleep(0) + return survivor_result + finally: + release_encode.set() + + codes, cache_hit, inflight_wait = asyncio.run(run()) + + assert calls == 1 + assert cache_hit is False + assert inflight_wait is True + assert torch.equal(codes, expected_codes) + assert "ref-a" not in adapter._higgs_audio_v3_ref_code_inflight + cached = adapter._get_higgs_audio_v3_ref_codes("ref-a") + assert cached is not None + assert torch.equal(cached, expected_codes) + + +def test_ref_code_inflight_consumes_failure_after_all_callers_cancel(): + adapter = HiggsAudioV3Adapter(SpeechServingContext(server=SimpleNamespace())) + encode_started = threading.Event() + release_encode = threading.Event() + loop_errors = [] + + def encode_reference_audio(_wav, _sr): + encode_started.set() + assert release_encode.wait(timeout=5) + raise RuntimeError("encode failed") + + async def resolve(): + return await adapter._resolve_higgs_audio_v3_ref_codes( + "ref-a", + object(), + 24000, + encode_reference_audio, + lambda codes: codes, + ) + + async def run(): + loop = asyncio.get_running_loop() + loop.set_exception_handler(lambda _loop, context: loop_errors.append(context)) + creator = asyncio.create_task(resolve()) + try: + assert await asyncio.wait_for(asyncio.to_thread(encode_started.wait), timeout=5) + waiter = asyncio.create_task(resolve()) + await asyncio.sleep(0) + + creator.cancel() + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await creator + with pytest.raises(asyncio.CancelledError): + await waiter + + release_encode.set() + for _ in range(100): + if "ref-a" not in adapter._higgs_audio_v3_ref_code_inflight: + break + await asyncio.sleep(0.01) + gc.collect() + await asyncio.sleep(0) + finally: + release_encode.set() + + asyncio.run(run()) + + assert "ref-a" not in adapter._higgs_audio_v3_ref_code_inflight + assert not loop_errors diff --git a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py index eb00fe78440..1e64f7a6373 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py @@ -4,7 +4,7 @@ import asyncio from collections import OrderedDict -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import numpy as np import torch @@ -62,7 +62,7 @@ async def _build_higgs_audio_v3_params(self, request: "OpenAICreateSpeechRequest encode_reference_audio, ) - wav_list, sr, cache_key = await self._resolve_ref_audio(request.ref_audio) + wav_list, sr, cache_key = await self._resolve_ref_audio(cast(str, request.ref_audio)) artifact_key = self._get_resolved_ref_audio_artifact_key(cache_key) wav = np.asarray(wav_list, dtype=np.float32) ref_codes_delayed, cache_hit, inflight_wait = await self._resolve_higgs_audio_v3_ref_codes( @@ -109,7 +109,7 @@ async def _resolve_higgs_audio_v3_ref_codes( task = self._higgs_audio_v3_ref_code_inflight.get(artifact_key) if task is not None: - return (await task).clone(), False, True + return (await asyncio.shield(task)).clone(), False, True async def _encode_and_cache() -> torch.Tensor: ref_codes_raw = await asyncio.to_thread(encode_reference_audio, wav, sr) @@ -120,12 +120,16 @@ async def _encode_and_cache() -> torch.Tensor: task = asyncio.create_task(_encode_and_cache()) self._higgs_audio_v3_ref_code_inflight[artifact_key] = task - try: - return (await task).clone(), False, False - finally: - if self._higgs_audio_v3_ref_code_inflight.get(artifact_key) is task: + + def _retire(t: asyncio.Task[torch.Tensor]) -> None: + if not t.cancelled(): + t.exception() + if self._higgs_audio_v3_ref_code_inflight.get(artifact_key) is t: self._higgs_audio_v3_ref_code_inflight.pop(artifact_key, None) + task.add_done_callback(_retire) + return (await asyncio.shield(task)).clone(), False, False + def _get_higgs_audio_v3_ref_codes(self, artifact_key: str | None) -> torch.Tensor | None: if not artifact_key: return None From 968f3e379c292d395e600dfe5dfbb1e21f905d90 Mon Sep 17 00:00:00 2001 From: xutianle <24210290017@m.fudan.edu.cn> Date: Wed, 9 Sep 2026 11:33:38 +0800 Subject: [PATCH 19/93] [Bugfix][CosyVoice3] Resolve hash snapshot pipeline (#6896) Signed-off-by: xutianle --- tests/config/test_config_factory.py | 38 +++++++++++++++++++++++++++++ vllm_omni/deploy/cosyvoice3.yaml | 1 + 2 files changed, 39 insertions(+) diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 6cb328bc5a8..51c7d549e28 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -932,6 +932,44 @@ class FakeConfig(PretrainedConfig): assert isinstance(omni_config, VllmOmniConfig) assert omni_config.stage_by_id(0).diffusion_config.model == "fake/model" + def test_cosyvoice3_deploy_resolves_hash_snapshot_without_hf_metadata(self, tmp_path): + snapshot = tmp_path / "models--FunAudioLLM--Fun-CosyVoice3-0.5B-2512" / "snapshots" / ("a" * 40) + snapshot.mkdir(parents=True) + (snapshot / "config.json").write_text("{}\n", encoding="utf-8") + model = str(snapshot) + deploy_path = get_deploy_config_path("cosyvoice3.yaml") + + assert StageConfigFactory.try_infer_model_type(model, trust_remote_code=False) is None + assert StageConfigFactory.get_pipeline_config(model, trust_remote_code=False) is None + + omni_config = StageConfigFactory.create_from_model( + model, + trust_remote_code=False, + cli_overrides={}, + deploy_config_path=deploy_path, + ) + legacy_configs, _ = StageConfigFactory.create_legacy_stage_configs_from_model( + model, + trust_remote_code=False, + cli_overrides={}, + deploy_config_path=deploy_path, + ) + + assert omni_config is not None + assert omni_config.pipeline_config.model_type == "cosyvoice3" + assert [ + (stage.stage_id, stage.stage_type, stage.worker_type, stage.model_stage) + for stage in omni_config.stage_configs + ] == [ + (0, StageType.LLM, "ar", "cosyvoice3_talker"), + (1, StageType.LLM, "generation", "cosyvoice3_code2wav"), + ] + assert all(stage.model_config.model == model for stage in omni_config.stage_configs) + assert all(stage.stage_type is not StageType.DIFFUSION for stage in omni_config.stage_configs) + + assert legacy_configs is not None + assert [stage.stage_type for stage in legacy_configs] == [StageType.LLM, StageType.LLM] + def test_pipeline_registration(self, clean_pipeline_registry): """Ensure that we can register and create a custom pipeline config.""" new_model_type = "new_model_type" diff --git a/vllm_omni/deploy/cosyvoice3.yaml b/vllm_omni/deploy/cosyvoice3.yaml index 551b42af60c..09fd83290d5 100644 --- a/vllm_omni/deploy/cosyvoice3.yaml +++ b/vllm_omni/deploy/cosyvoice3.yaml @@ -23,6 +23,7 @@ # On a slower GPU class, lower max_num_seqs (e.g. 4) and/or codec_chunk_frames # first. Raising connector_get_sleep_s also relaxes the consumer's polling # cadence, but trades against TTFP — keep it as a last resort. +pipeline: cosyvoice3 async_chunk: true connectors: connector_of_shared_memory: From 09b720b0d2b3528ba17d70900165e1e6508bfa76 Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Wed, 9 Sep 2026 12:23:29 +0800 Subject: [PATCH 20/93] [CI] Skip Qwen3-Omni Server VAD multi-turn realtime test (#7279) (#7314) Signed-off-by: wangyu <410167048@qq.com> --- .../entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py index 199bea0edf4..8c80b964928 100644 --- a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py +++ b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py @@ -408,6 +408,7 @@ def test_long_audio_response_completes_async_chunk(self, omni_server) -> None: @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", realtime_server_vad_server_params, indirect=True) + @pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/7279") def test_server_vad_multi_turn_without_client_commit( self, cached_silero_vad_artifact: str, From e0d08797944277090699da03e29e8a3c0739e0f4 Mon Sep 17 00:00:00 2001 From: andyluo7 <43718156+andyluo7@users.noreply.github.com> Date: Tue, 8 Sep 2026 22:48:08 -0700 Subject: [PATCH 21/93] [Bugfix][Magi2] Allow import without an active Triton driver (#7239) Signed-off-by: andyluo7 --- vllm_omni/diffusion/models/magi2/mh_moe.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/vllm_omni/diffusion/models/magi2/mh_moe.py b/vllm_omni/diffusion/models/magi2/mh_moe.py index 6419cfae890..9e1e0788a7e 100644 --- a/vllm_omni/diffusion/models/magi2/mh_moe.py +++ b/vllm_omni/diffusion/models/magi2/mh_moe.py @@ -21,7 +21,7 @@ import torch import torch.nn as nn import torch.nn.functional as F -from vllm.triton_utils import tl, triton +from vllm.triton_utils import HAS_TRITON, tl, triton from vllm_omni.platforms import current_omni_platform @@ -133,9 +133,10 @@ def torch_mh_moe_forward( return output -_SWIGLU7_ALPHA = tl.constexpr(1.702) -_SWIGLU7_LIMIT = tl.constexpr(7.0) -_SWIGLU7_BIAS = tl.constexpr(1.0) +# vLLM's no-Triton placeholder exposes ``tl.constexpr`` as ``None``. +_SWIGLU7_ALPHA = tl.constexpr(1.702) if HAS_TRITON else 1.702 +_SWIGLU7_LIMIT = tl.constexpr(7.0) if HAS_TRITON else 7.0 +_SWIGLU7_BIAS = tl.constexpr(1.0) if HAS_TRITON else 1.0 @triton.jit @@ -198,7 +199,7 @@ def _mh_moe_kernel( block_t: tl.constexpr, block_dh: tl.constexpr, block_de: tl.constexpr, - acc_dtype: tl.constexpr = tl.float32, + acc_dtype: tl.constexpr, deterministic: tl.constexpr = False, ): tile_id = tl.program_id(0) From 5b002727b00421f2967f996cfd788950d024f932 Mon Sep 17 00:00:00 2001 From: NATURE Date: Wed, 9 Sep 2026 13:49:43 +0800 Subject: [PATCH 22/93] [Core] Split Omni connector model runner mixin (#6903) Signed-off-by: natureofnature --- tests/worker/test_omni_connector_mixin.py | 66 +- tests/worker/test_payload_span.py | 166 -- .../omni_connectors/model_runner/__init__.py | 3 + .../omni_connector_payload_transport.py | 1322 ++++++++++ .../model_runner/omni_connector_runtime.py | 1061 ++++++++ .../omni_connector_model_runner_mixin.py | 2350 +---------------- vllm_omni/worker/payload_span.py | 64 - 7 files changed, 2463 insertions(+), 2569 deletions(-) delete mode 100644 tests/worker/test_payload_span.py create mode 100644 vllm_omni/distributed/omni_connectors/model_runner/__init__.py create mode 100644 vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py create mode 100644 vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py delete mode 100644 vllm_omni/worker/payload_span.py diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index af472fd95ee..8366e177abd 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.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 """Unit tests for OmniConnectorModelRunnerMixin. These tests use a mock connector (in-memory dict store) and do not require @@ -376,6 +376,19 @@ def test_finished_load_reqs_flow_to_chunk_ready(self): class TestLoadCustomFuncSelection(unittest.TestCase): + def test_uses_validator_override_from_public_mixin(self): + config = SimpleNamespace( + async_chunk=True, + custom_process_next_stage_input_func=f"{__name__}._make_request", + ) + + with patch.object(MixinHost, "_is_connector_payload_builder", return_value=False) as validator: + selected_path, func = MixinHost._load_custom_func(config) + + validator.assert_called_once_with(_make_request) + assert selected_path is None + assert func is None + def test_skips_non_payload_stage_input_processors_for_full_payload_mode(self): incompatible_paths = [ "vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav", @@ -861,7 +874,7 @@ def test_rank0_only_polls_connector_for_tp_full_payload(self): host._omni_connector.get.return_value = connector_result tp_group = _FakeTPGroup(world_size=2, rank_in_group=0) - with patch("vllm_omni.worker.omni_connector_model_runner_mixin.get_tp_group", return_value=tp_group): + with patch.object(host, "_get_local_tp_group", return_value=tp_group): made_progress = host._poll_single_request("r1") self.assertTrue(made_progress) @@ -882,7 +895,7 @@ def test_tp_follower_skips_connector_poll_for_full_payload(self): host._get_req_chunk["r1"] = 0 tp_group = _FakeTPGroup(world_size=2, rank_in_group=1) - with patch("vllm_omni.worker.omni_connector_model_runner_mixin.get_tp_group", return_value=tp_group): + with patch.object(host, "_get_local_tp_group", return_value=tp_group): made_progress = host._poll_single_request("r1") self.assertFalse(made_progress) @@ -900,7 +913,7 @@ def test_recv_full_payload_inputs_broadcasts_tp_leader_results_to_followers(self payload = {"tok": [10], "finished": torch.tensor(True)} tp_group = _FakeTPGroup(world_size=2, rank_in_group=1, follower_result={"r1": payload}) - with patch("vllm_omni.worker.omni_connector_model_runner_mixin.get_tp_group", return_value=tp_group): + with patch.object(host, "_get_local_tp_group", return_value=tp_group): results = host.recv_full_payload_inputs(scheduler_output=None) self.assertEqual(results, {"r1": payload}) @@ -936,7 +949,7 @@ def test_rank0_only_polls_connector_for_tp_async_chunk(self): host._omni_connector.get.return_value = (payload, 123) tp_group = _FakeTPGroup(world_size=2, rank_in_group=0) - with patch("vllm_omni.worker.omni_connector_model_runner_mixin.get_tp_group", return_value=tp_group): + with patch.object(host, "_get_local_tp_group", return_value=tp_group): made_progress = host._poll_single_request("r1") self.assertTrue(made_progress) @@ -951,7 +964,7 @@ def test_tp_follower_skips_connector_poll_for_async_chunk(self): host = self._make_host(rank=1) tp_group = _FakeTPGroup(world_size=2, rank_in_group=1) - with patch("vllm_omni.worker.omni_connector_model_runner_mixin.get_tp_group", return_value=tp_group): + with patch.object(host, "_get_local_tp_group", return_value=tp_group): made_progress = host._poll_single_request("r1") self.assertFalse(made_progress) @@ -976,7 +989,7 @@ def test_get_output_broadcasts_tp_async_chunk_payloads_to_followers(self): } tp_group = _FakeTPGroup(world_size=2, rank_in_group=1, follower_result=packet) - with patch("vllm_omni.worker.omni_connector_model_runner_mixin.get_tp_group", return_value=tp_group): + with patch.object(host, "_get_local_tp_group", return_value=tp_group): output = host.get_omni_connector_output() self.assertEqual(output.chunk_ready_req_ids, {"r1"}) @@ -1049,6 +1062,45 @@ def test_cleanup_removes_kv_state(self): class TestAsyncPayloadLifecycle(unittest.TestCase): """Regression tests for async payload delivery lifecycle.""" + def test_accumulate_payload_concatenates_chunks(self): + host = MixinHost() + host._send_side_request_payload = {} + first = host._accumulate_payload( + "r1", + { + "embed": {"decode": torch.tensor([[1.0, 2.0]])}, + "ids": {"output": [1]}, + "meta": {"finished": False}, + }, + ) + merged = host._accumulate_payload( + "r1", + { + "embed": {"decode": torch.tensor([[3.0, 4.0], [5.0, 6.0]])}, + "ids": {"output": [2, 3]}, + "meta": {"finished": True}, + }, + ) + torch.testing.assert_close(merged["embed"]["decode"], torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])) + self.assertEqual(merged["ids"]["output"], [1, 2, 3]) + self.assertIs(merged["meta"]["finished"], True) + self.assertEqual(first["embed"]["decode"].shape, (1, 2)) + self.assertEqual(first["ids"]["output"], [1]) + self.assertIs(first["meta"]["finished"], False) + + def test_accumulate_payload_replaces_override_keys(self): + host = MixinHost() + host._send_side_request_payload = {} + host._accumulate_payload("r1", {"embed": {"decode": torch.ones(2, 2)}, "ids": {"output": [1, 2]}}) + payload = { + "embed": {"decode": torch.zeros(1, 2)}, + "ids": {"output": [3]}, + "meta": {"override_keys": [["embed", "decode"], ["ids", "output"]]}, + } + merged = host._accumulate_payload("r1", payload) + torch.testing.assert_close(merged["embed"]["decode"], payload["embed"]["decode"]) + self.assertEqual(merged["ids"]["output"], [3]) + def test_send_side_request_payload_not_cleared_before_payload_is_consumable(self): host = MixinHost() host.init_omni_connectors( diff --git a/tests/worker/test_payload_span.py b/tests/worker/test_payload_span.py deleted file mode 100644 index 5eb4d8319f8..00000000000 --- a/tests/worker/test_payload_span.py +++ /dev/null @@ -1,166 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Characterization tests for ``vllm_omni.worker.payload_span``. - -These pin the CURRENT behaviour of the thinker-decode span helpers so the later -worker refactors that touch span -plumbing have a regression net. Pure CPU — no model, no CUDA. -""" - -import pytest -import torch - -from vllm_omni.worker.payload_span import ( - get_tensor_span, - get_tensor_span_row, - merge_tensor_spans, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def _rows(n: int, *, start_val: int = 0, hidden: int = 1) -> torch.Tensor: - """A ``[n, hidden]`` tensor whose first column counts up from ``start_val``. - - ``[n, hidden]`` mirrors the connector async-chunk shape (tokens x hidden). - """ - base = torch.arange(start_val, start_val + n, dtype=torch.float32).unsqueeze(1) - return base.repeat(1, hidden) - - -# --------------------------------------------------------------------------- # -# get_tensor_span # -# --------------------------------------------------------------------------- # -def test_get_tensor_span_valid_returns_triplet(): - tensor = _rows(3, hidden=4) - payload = {"emb": tensor, "start": 5, "end": 8} - - span = get_tensor_span(payload, tensor_key="emb", start_key="start", end_key="end") - - assert span is not None - got_tensor, start, end = span - assert got_tensor is tensor - assert (start, end) == (5, 8) - - -def test_get_tensor_span_non_tensor_returns_none(): - payload = {"emb": None, "start": 0, "end": 0} - assert get_tensor_span(payload, tensor_key="emb", start_key="start", end_key="end") is None - - -def test_get_tensor_span_non_int_bounds_returns_none(): - tensor = _rows(2) - # start is a float, not an int -> rejected. - payload = {"emb": tensor, "start": 0.0, "end": 2} - assert get_tensor_span(payload, tensor_key="emb", start_key="start", end_key="end") is None - - -def test_get_tensor_span_negative_start_returns_none(): - tensor = _rows(1) - payload = {"emb": tensor, "start": -1, "end": 0} - assert get_tensor_span(payload, tensor_key="emb", start_key="start", end_key="end") is None - - -def test_get_tensor_span_end_before_start_returns_none(): - tensor = _rows(2) - payload = {"emb": tensor, "start": 5, "end": 3} - assert get_tensor_span(payload, tensor_key="emb", start_key="start", end_key="end") is None - - -def test_get_tensor_span_length_mismatch_returns_none(): - tensor = _rows(3) # tensor.shape[0] == 3 - payload = {"emb": tensor, "start": 0, "end": 2} # end - start == 2 - assert get_tensor_span(payload, tensor_key="emb", start_key="start", end_key="end") is None - - -# --------------------------------------------------------------------------- # -# merge_tensor_spans # -# --------------------------------------------------------------------------- # -def test_merge_none_operand_returns_none(): - span = (_rows(2), 0, 2) - assert merge_tensor_spans(None, span) is None - assert merge_tensor_spans(span, None) is None - - -def test_merge_adjacent_spans_concatenates(): - existing = (_rows(3, start_val=0), 0, 3) - incoming = (_rows(2, start_val=3), 3, 5) - - merged = merge_tensor_spans(existing, incoming) - - assert merged is not None - tensor, start, end = merged - assert (start, end) == (0, 5) - assert tensor.shape[0] == 5 - assert tensor[:, 0].tolist() == [0.0, 1.0, 2.0, 3.0, 4.0] - - -def test_merge_overlapping_spans_trims_incoming(): - # existing covers [0, 3); incoming covers [2, 5) -> 1 row of overlap. - existing = (_rows(3, start_val=0), 0, 3) - incoming = (_rows(3, start_val=100), 2, 5) - - merged = merge_tensor_spans(existing, incoming) - - assert merged is not None - tensor, start, end = merged - assert (start, end) == (0, 5) - # First 3 rows from existing, then incoming's tail (overlap=1 row trimmed). - assert tensor[:, 0].tolist() == [0.0, 1.0, 2.0, 101.0, 102.0] - - -def test_merge_fully_overlapping_incoming_keeps_existing(): - # incoming [1, 3) is fully covered by existing [0, 3): overlap >= incoming len. - existing = (_rows(3, start_val=0), 0, 3) - incoming = (_rows(2, start_val=100), 1, 3) - - merged = merge_tensor_spans(existing, incoming) - - assert merged is not None - tensor, start, end = merged - assert (start, end) == (0, 3) - assert tensor is existing[0] # returned unchanged - - -def test_merge_non_contiguous_spans_returns_none(): - # gap between existing end (3) and incoming start (5). - existing = (_rows(3, start_val=0), 0, 3) - incoming = (_rows(2, start_val=100), 5, 7) - - assert merge_tensor_spans(existing, incoming) is None - - -def test_merge_casts_incoming_to_existing_dtype_and_device(): - existing = (_rows(2, start_val=0).to(torch.float32), 0, 2) - incoming = (torch.arange(2, 4, dtype=torch.int64).unsqueeze(1), 2, 4) - - merged = merge_tensor_spans(existing, incoming) - - assert merged is not None - tensor, _, _ = merged - assert tensor.dtype == torch.float32 - assert tensor.device == existing[0].device - - -# --------------------------------------------------------------------------- # -# get_tensor_span_row # -# --------------------------------------------------------------------------- # -def test_get_tensor_span_row_in_range(): - span = (_rows(3, start_val=10), 5, 8) # rows map to indices 5, 6, 7 - row = get_tensor_span_row(span, 6) - assert row is not None - assert row[0].item() == 11.0 # second row (start_val 10 + 1) - - -def test_get_tensor_span_row_none_span_returns_none(): - assert get_tensor_span_row(None, 0) is None - - -def test_get_tensor_span_row_below_start_returns_none(): - span = (_rows(3), 5, 8) - assert get_tensor_span_row(span, 4) is None - - -def test_get_tensor_span_row_at_or_above_end_returns_none(): - span = (_rows(3), 5, 8) - assert get_tensor_span_row(span, 8) is None diff --git a/vllm_omni/distributed/omni_connectors/model_runner/__init__.py b/vllm_omni/distributed/omni_connectors/model_runner/__init__.py new file mode 100644 index 00000000000..5cfa003116a --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/model_runner/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Internal model-runner components for Omni connector transport.""" 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 new file mode 100644 index 00000000000..25515cfaf06 --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py @@ -0,0 +1,1322 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Omni connector payload transport for model runners.""" + +from __future__ import annotations + +import importlib +import inspect +from collections import deque +from typing import TYPE_CHECKING, Any + +import torch +from vllm.distributed.parallel_state import get_tp_group + +from vllm_omni.data_entry_keys import OmniPayload +from vllm_omni.distributed.omni_connectors.model_runner.omni_connector_runtime import ( + _OmniConnectorRuntimeMixin, + logger, + should_accumulate_full_payload_output, +) +from vllm_omni.outputs import OmniConnectorOutput + +if TYPE_CHECKING: + from vllm_omni.distributed.omni_connectors.connectors.base import ( + OmniConnectorBase, + ) + + +class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): + """Own payload caching, full/chunk transport, and connector I/O.""" + + # ------------------------------------------------------------------ # + # Local payload cache (RFC §2.4 – Model Runner ownership) + # ------------------------------------------------------------------ # + + def put_local_stage_payload(self, req_id: str, payload: OmniPayload) -> None: + """Store a full stage payload in the local cache.""" + self._local_stage_payload_cache[req_id] = payload + + def get_local_stage_payload(self, req_id: str) -> OmniPayload | None: + """Read a stage payload without removing it.""" + return self._local_stage_payload_cache.get(req_id) + + def pop_local_stage_payload(self, req_id: str) -> OmniPayload | None: + """Remove and return a stage payload (consume after use).""" + return self._local_stage_payload_cache.pop(req_id, None) + + def put_local_request_metadata(self, req_id: str, metadata: dict[str, Any]) -> None: + """Store lightweight scheduling metadata for a request.""" + self._local_request_metadata[req_id] = metadata + + def get_local_request_metadata(self, req_id: str) -> dict[str, Any] | None: + """Retrieve scheduling metadata for a request.""" + return self._local_request_metadata.get(req_id) + + # ------------------------------------------------------------------ # + # Scheduling metadata extraction + # ------------------------------------------------------------------ # + + @classmethod + def _extract_scheduling_metadata(cls, payload: OmniPayload) -> dict[str, Any]: + """Extract only the fields the scheduler needs from a full payload.""" + extracted: dict[str, Any] = {} + meta = payload.get("meta") if isinstance(payload, dict) else None + meta = meta if isinstance(meta, dict) else {} + + if "next_stage_prompt_len" in meta: + extracted["next_stage_prompt_len"] = meta["next_stage_prompt_len"] + elif "next_stage_prompt_len" in payload: + logger.warning_once( + "legacy flat 'next_stage_prompt_len' key in payload; expected 'meta.next_stage_prompt_len'" + ) + extracted["next_stage_prompt_len"] = payload["next_stage_prompt_len"] + + audio_codes = cls._payload_audio_codes(payload) + if audio_codes is not None: + extracted["code_predictor_codes"] = audio_codes + + if "left_context_size" in meta: + extracted["left_context_size"] = meta["left_context_size"] + elif "left_context_size" in payload: + logger.warning_once("legacy flat 'left_context_size' key in payload; expected 'meta.left_context_size'") + + return extracted + + _NON_CONSUMABLE_PAYLOAD_KEYS: set[tuple[str, str]] = { + ("meta", "finished"), + ("meta", "override_keys"), + ("meta", "next_stage_prompt_len"), + ("meta", "left_context_size"), + ("ids", "output"), + ("embed", "decode_token_start"), + ("embed", "decode_token_end"), + } + + @staticmethod + def _payload_value_has_content(value: Any) -> bool: + if value is None: + return False + if isinstance(value, torch.Tensor): + return value.numel() > 0 + if isinstance(value, (list, tuple, dict, set)): + return len(value) > 0 + return True + + @staticmethod + def _payload_finished(payload: Any) -> bool: + if not isinstance(payload, dict): + return False + if "finished" in payload: + logger.warning_once("legacy flat 'finished' key in payload; expected 'meta.finished'") + meta = payload.get("meta") + if not isinstance(meta, dict) or "finished" not in meta: + return False + flag = meta["finished"] + if isinstance(flag, torch.Tensor): + return flag.numel() == 1 and bool(flag.item()) + return bool(flag) + + @staticmethod + def _payload_audio_codes(payload: Any) -> Any: + if not isinstance(payload, dict): + return None + if "code_predictor_codes" in payload: + logger.warning_once("legacy flat 'code_predictor_codes' key in payload; expected 'codes.audio'") + codes = payload.get("codes") + if isinstance(codes, dict): + return codes.get("audio") + return None + + @classmethod + def _payload_is_consumable(cls, payload: OmniPayload | None) -> bool: + """Return True when an async payload can drive a real forward step. + + Metadata-only wake-ups should not transition WAITING_FOR_CHUNK requests + back to schedulable state. In particular, a widened token horizon without + any newly visible thinker decode embeds should not force a placeholder-only + talker decode step. + """ + if not isinstance(payload, dict) or not payload: + return False + + embed = payload.get("embed") + if isinstance(embed, dict): + decode_embeddings = embed.get("decode") + if isinstance(decode_embeddings, torch.Tensor): + if decode_embeddings.ndim == 0: + return True + return decode_embeddings.numel() > 0 and decode_embeddings.shape[0] > 0 + + audio_codes = cls._payload_audio_codes(payload) + if audio_codes is not None: + if isinstance(audio_codes, torch.Tensor): + return audio_codes.numel() > 0 + if hasattr(audio_codes, "__len__"): + return len(audio_codes) > 0 + return True + + for key, value in payload.items(): + if isinstance(value, dict): + for sk, sv in value.items(): + if (key, sk) in cls._NON_CONSUMABLE_PAYLOAD_KEYS: + continue + if cls._payload_value_has_content(sv): + return True + continue + if cls._payload_value_has_content(value): + return True + return False + + @staticmethod + def _get_local_tp_group() -> Any | None: + """Return the local TP group when tensor parallelism is initialized.""" + try: + return get_tp_group() + except Exception: + return None + + def _recv_ordinary_stage_result( + self, + connector: OmniConnectorBase, + from_stage: str, + to_stage: str, + connector_get_key: str, + ) -> Any: + """Receive one ordinary non-KV stage payload on the local leader rank only.""" + tp_group = self._get_local_tp_group() + if tp_group is None or getattr(tp_group, "world_size", 1) <= 1: + return connector.get(from_stage, to_stage, connector_get_key) + if not self.is_data_transfer_rank(): + return None + return connector.get(from_stage, to_stage, connector_get_key) + + def _recv_full_payload_result( + self, + connector: OmniConnectorBase, + from_stage: str, + to_stage: str, + connector_get_key: str, + ) -> Any: + """Receive one full-payload transfer on the local leader rank only.""" + return self._recv_ordinary_stage_result( + connector, + from_stage, + to_stage, + connector_get_key, + ) + + def _recv_async_chunk_result( + self, + connector: OmniConnectorBase, + from_stage: str, + to_stage: str, + connector_get_key: str, + ) -> Any: + """Receive one ordinary async chunk on the local leader rank only.""" + return self._recv_ordinary_stage_result( + connector, + from_stage, + to_stage, + connector_get_key, + ) + + @staticmethod + def _snapshot_payload(payload: Any) -> Any: + if isinstance(payload, dict): + return dict(payload) + return payload + + def _broadcast_tp_payload_packet(self, packet: Any) -> Any: + """Broadcast one ordinary payload packet from TP rank 0 when TP is active.""" + tp_group = self._get_local_tp_group() + if tp_group is None or getattr(tp_group, "world_size", 1) <= 1: + return packet + leader_packet = packet if self.is_data_transfer_rank() else None + return tp_group.broadcast_object(leader_packet, src=0) + + def _apply_staged_payloads_locked(self, staged_payloads: dict[str, Any]) -> None: + for req_id, payload in staged_payloads.items(): + self._local_stage_payload_cache[req_id] = self._snapshot_payload(payload) + + def _collect_full_payload_results_locked(self) -> dict[str, Any] | None: + if not self._full_payload_pending_broadcast_req_ids: + return None + results: dict[str, Any] = {} + missing_req_ids: list[str] = [] + for req_id in tuple(self._full_payload_pending_broadcast_req_ids): + payload = self._local_stage_payload_cache.get(req_id) + if payload is None: + missing_req_ids.append(req_id) + continue + results[req_id] = self._snapshot_payload(payload) + self._full_payload_pending_broadcast_req_ids.discard(req_id) + if missing_req_ids: + logger.warning( + "[Stage-%s] _collect_full_payload_results_locked: " + "pending full-payload reqs missing from local cache: %s", + self._stage_id, + missing_req_ids, + ) + return results or None + + def _collect_async_chunk_fanout_packet_locked(self) -> dict[str, Any] | None: + payload_req_ids = set(self._async_chunk_updated_req_ids) + payload_req_ids.update(self._finished_load_reqs) + payload_req_ids.update(self._chunk_finished_req_ids) + payload_req_ids.update(self._local_request_metadata) + if not ( + payload_req_ids or self._finished_load_reqs or self._chunk_finished_req_ids or self._local_request_metadata + ): + return None + + staged_payloads = { + req_id: self._snapshot_payload(self._local_stage_payload_cache[req_id]) + for req_id in payload_req_ids + if req_id in self._local_stage_payload_cache + } + packet = { + "staged_payloads": staged_payloads, + "request_metadata": dict(self._local_request_metadata), + "newly_finished": set(self._finished_load_reqs), + "chunk_finished": set(self._chunk_finished_req_ids), + } + + self._async_chunk_updated_req_ids.clear() + self._finished_load_reqs.clear() + self._chunk_finished_req_ids.clear() + self._local_request_metadata.clear() + + for req_id in packet["chunk_finished"]: + if req_id not in self._local_stage_payload_cache: + continue + ext_req_id = self._request_ids_mapping.get(req_id, req_id) + self._send_side_request_payload.pop(ext_req_id, None) + if ext_req_id != req_id: + self._send_side_request_payload.pop(req_id, None) + + return packet + + def _apply_async_chunk_fanout_packet(self, packet: dict[str, Any]) -> None: + staged_payloads = packet.get("staged_payloads", {}) + chunk_finished = set(packet.get("chunk_finished", ())) + with self._lock: + self._apply_staged_payloads_locked(staged_payloads) + for req_id in chunk_finished: + self._pending_load_reqs.pop(req_id, None) + self._chunk_stream_completed.add(req_id) + + # Output aggregation + # ------------------------------------------------------------------ # + + def get_omni_connector_output(self) -> OmniConnectorOutput: + """Collect and reset transfer results for this execute_model cycle. + + ``request_metadata`` carries only lightweight scheduling metadata. + Full payloads remain owned by the Model Runner local cache for all + paths. + """ + if not hasattr(self, "_lock"): + return OmniConnectorOutput() + + tp_group = self._get_local_tp_group() + if self._async_chunk and tp_group is not None and getattr(tp_group, "world_size", 1) > 1: + if self.is_data_transfer_rank(): + with self._lock: + fanout_packet = self._collect_async_chunk_fanout_packet_locked() + else: + fanout_packet = None + fanout_packet = self._broadcast_tp_payload_packet(fanout_packet) + if fanout_packet is None: + newly_finished = set() + chunk_finished = set() + request_metadata = {} + else: + if not self.is_data_transfer_rank(): + self._apply_async_chunk_fanout_packet(fanout_packet) + newly_finished = set(fanout_packet["newly_finished"]) + chunk_finished = set(fanout_packet["chunk_finished"]) + request_metadata = dict(fanout_packet["request_metadata"]) + else: + with self._lock: + newly_finished = set(self._finished_load_reqs) + self._finished_load_reqs.clear() + chunk_finished = set(self._chunk_finished_req_ids) + self._chunk_finished_req_ids.clear() + request_metadata = dict(self._local_request_metadata) + self._local_request_metadata.clear() + # _send_side_request_payload is the async accumulation buffer for + # future recv chunks. Clearing it on every consumable wake-up drops + # intermediate + # thinker decode spans before the model side can consume them. + # Only terminal chunk_finished requests may release that buffer. + for req_id in chunk_finished: + if req_id not in self._local_stage_payload_cache: + continue + ext_req_id = self._request_ids_mapping.get(req_id, req_id) + self._send_side_request_payload.pop(ext_req_id, None) + if ext_req_id != req_id: + self._send_side_request_payload.pop(req_id, None) + self._chunk_ready_req_ids.update(newly_finished) + + output = OmniConnectorOutput( + chunk_ready_req_ids=set(self._chunk_ready_req_ids), + chunk_finished_req_ids=chunk_finished, + request_metadata=request_metadata, + kv_sent_req_ids=list(self._kv_sent_req_ids), + stage_recv_req_ids=set(self._stage_recv_req_ids), + has_pending_kv_work=self.has_pending_kv_work(), + ) + if output.stage_recv_req_ids or chunk_finished or newly_finished: + logger.debug( + "[Stage-%s] get_omni_connector_output: stage_recv=%s, chunk_finished=%s, chunk_ready=%s", + self._stage_id, + output.stage_recv_req_ids, + chunk_finished, + output.chunk_ready_req_ids, + ) + self._chunk_ready_req_ids.clear() + self._kv_sent_req_ids.clear() + self._stage_recv_req_ids.clear() + return output + + @staticmethod + def _connector_output_has_signals(output: OmniConnectorOutput) -> bool: + return bool( + output.chunk_ready_req_ids + or output.chunk_finished_req_ids + or output.request_metadata + or output.kv_sent_req_ids + or output.stage_recv_req_ids + or output.has_pending_kv_work + ) + + def attach_omni_connector_output(self, result: Any | None) -> Any: + omni_output = self.get_omni_connector_output() + if not self._connector_output_has_signals(omni_output): + return result + + from copy import copy + + from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT + + wrapped = copy(result if result is not None else EMPTY_MODEL_RUNNER_OUTPUT) + wrapped.omni_connector_output = omni_output + return wrapped + + # ------------------------------------------------------------------ # + # Properties for compatibility with custom_process funcs that access + # transfer_manager.put_req_chunk / request_payload / code_prompt_token_ids + # ------------------------------------------------------------------ # + + @property + def put_req_chunk(self) -> dict[str, int]: + return self._put_req_chunk + + @property + def ramp_chunk_count(self) -> dict[str, int]: + return self._ramp_chunk_count + + @property + def request_payload(self) -> dict[str, dict[str, Any]]: + return self._send_side_request_payload + + @request_payload.setter + def request_payload(self, value: dict[str, dict[str, Any]]) -> None: + self._send_side_request_payload = value + + @property + def code_prompt_token_ids(self) -> dict[str, list[list[int]]]: + return self._code_prompt_token_ids + + @property + def connector(self) -> Any | None: + return self._omni_connector + + # ------------------------------------------------------------------ # + # full_payload_mode (recv_full_payload_inputs / send_full_payload_outputs) + # ------------------------------------------------------------------ # + + def recv_full_payload_inputs(self, scheduler_output: Any) -> dict[str, Any] | None: + """Check for incoming full_payload_mode stage inputs (non-blocking). + + Returns a dict mapping ``request_id -> engine_inputs`` for data + that has arrived, or ``None`` if nothing is ready. Stores full + payloads in the local cache and extracts scheduling metadata. + """ + # Fast path: when TP is trivial (no peer ranks waiting on a broadcast) + # and the bg recv thread has not staged anything, skip the lock + TP + # broadcast cycle entirely. _broadcast_tp_payload_packet already + # returns its input unchanged under the same world_size<=1 condition, + # so the original code path was a no-op here on every empty step. + tp_group = self._get_local_tp_group() + if ( + tp_group is None or getattr(tp_group, "world_size", 1) <= 1 + ) and not self._full_payload_pending_broadcast_req_ids: + return None + with self._lock: + results = self._collect_full_payload_results_locked() if self.is_data_transfer_rank() else None + results = self._broadcast_tp_payload_packet(results) + if not results: + return None + with self._lock: + self._stage_recv_req_ids.update(results.keys()) + for req_id in results: + self._pending_load_reqs.pop(req_id, None) + self._apply_staged_payloads_locked(results) + for req_id, payload in results.items(): + self._local_request_metadata[req_id] = self._extract_scheduling_metadata(payload) + logger.debug( + "[Stage-%s] recv_full_payload_inputs: consumed %s reqs: %s, stage_recv_req_ids now=%s", + self._stage_id, + len(results), + list(results.keys()), + self._stage_recv_req_ids, + ) + return results + + def _get_model_config(self) -> Any: + model_config = getattr(self, "model_config", None) + if model_config is not None: + return model_config + return getattr(getattr(self, "vllm_config", None), "model_config", None) + + def _should_accumulate_full_payload_output(self) -> bool: + """Gate send-side full-payload output accumulation only. + + Cached per instance: the result depends only on model_config / + _custom_process_func, both of which are set at init time. Avoid + the per-step dynamic import inside the model decode loop. + """ + if getattr(self, "_omni_connector", None) is None: + # No connector at all: send_full_payload_outputs would no-op. + # Skip the per-step accumulator+build that would otherwise be + # silently discarded. Defends against a terminal stage whose + # custom_process_input_func has a *_full_payload derivative in + # the same module (e.g. dynin stage 2 token2image_to_token2audio + # in pipelines that don't configure any connector at all). + # + # Known limitation: a *terminal-consumer* stage that has a + # connector configured for receiving upstream input is NOT + # caught here -- ``_omni_connector`` is non-None for it, and + # ``_load_custom_func`` may still resolve a ``*_full_payload`` + # derivative from this stage's ``custom_process_input_func``. + # In that case the accumulator builds payloads that + # ``send_full_payload_outputs`` later drops via its own + # connector-side checks (wasted CPU, not a functional bug). + # A topology-aware gate (explicit producer field or pipeline + # is_terminal info) would close the gap; that change is out + # of scope for this PR. + self._should_accumulate_full_payload_output_cached = False + return False + cached = getattr(self, "_should_accumulate_full_payload_output_cached", None) + if cached is not None: + return cached + model_config = self._get_model_config() + if model_config is None: + self._should_accumulate_full_payload_output_cached = False + return False + result = should_accumulate_full_payload_output( + model_config, + getattr(self, "_custom_process_func", None), + ) + self._should_accumulate_full_payload_output_cached = result + return result + + @staticmethod + def _new_full_payload_accumulator(output: dict[str, Any]): + chunks: dict[str, list[torch.Tensor]] = {} + latest: dict[str, Any] = {} + rows: dict[str, int] = {} + for k, v in output.items(): + if isinstance(v, torch.Tensor) and v.dim() >= 2: + chunks[k] = [v] + rows[k] = int(v.shape[0]) + else: + latest[k] = v + return chunks, latest, rows + + @staticmethod + def _materialize_full_payload_entry(entry): + if len(entry) == 2: + return entry + chunks, latest, _rows, request = entry + output = dict(latest) + for k, tensors in chunks.items(): + if tensors: + output[k] = tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=0) + return output, request + + def _resolve_full_payload_replace_keys(self) -> frozenset: + """Per-model REPLACE-key set for the full-payload accumulator. + + Looked up from the stage-input-processor module that ships the model's sync builder + (`model_config.custom_process_input_func.__module__`). The module + declares ``_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str]``; if absent, + returns the empty set. + + Cached per instance. Keys in this set use REPLACE semantics in the + accumulator (subsequent emissions discard prior chunks) instead of + the default CONCAT semantics. Use for tensors that carry the full + result so far rather than per-step deltas (e.g. ``model_outputs``). + """ + cached = getattr(self, "_full_payload_replace_keys_cached", None) + if cached is not None: + return cached + proc = getattr(self, "_custom_process_func", None) + if proc is None: + self._full_payload_replace_keys_cached = frozenset() + return self._full_payload_replace_keys_cached + module_name = getattr(proc, "__module__", None) + if module_name is None: + self._full_payload_replace_keys_cached = frozenset() + return self._full_payload_replace_keys_cached + try: + import sys as _sys + + mod = _sys.modules.get(module_name) or importlib.import_module(module_name) + keys = getattr(mod, "_FULL_PAYLOAD_REPLACE_KEYS", frozenset()) + except ImportError: + logger.debug( + "Could not import stage input processor module %s while resolving " + "_FULL_PAYLOAD_REPLACE_KEYS; using CONCAT semantics for all keys.", + module_name, + exc_info=True, + ) + keys = frozenset() + if not isinstance(keys, (frozenset, set)): + logger.debug( + "Ignoring non-set _FULL_PAYLOAD_REPLACE_KEYS from %s: %s", + module_name, + type(keys).__name__, + ) + keys = frozenset() + self._full_payload_replace_keys_cached = frozenset(keys) + logger.debug( + "Resolved _FULL_PAYLOAD_REPLACE_KEYS for %s: %s", + module_name, + sorted(self._full_payload_replace_keys_cached), + ) + return self._full_payload_replace_keys_cached + + def accumulate_full_payload_output( + self, + req_id: str, + pooler_output: Any, + request: Any, + ) -> None: + """Accumulate pooler_output for a request across steps (full_payload_mode). + + Per-token tensors (2-D+, matching trailing dims) are concatenated + along dim-0. Scalar / global tensors (1-D or 0-D) are replaced + with the latest value. + + Note: codec rows are NOT filtered for zero placeholders here. The + downstream consumer ``_extract_qwen3_full_payload_codec_rows`` crops + codec rows using ``output_token_ids`` as the authoritative source, + which makes any sender-side zero filtering redundant. Skipping the + sender-side ``t.any()`` scan also avoids a per-tensor GPU->CPU device + sync that stalled the decode pipeline. + + The data is actually sent when ``flush_full_payload_outputs`` is called + with the finished request IDs from the next scheduler cycle. + """ + replace_keys = self._resolve_full_payload_replace_keys() + existing = self._pending_full_payload_send.get(req_id) + + if existing is None: + chunks, latest, rows = self._new_full_payload_accumulator(pooler_output) + self._pending_full_payload_send[req_id] = (chunks, latest, rows, request) + return + + if len(existing) == 2: + chunks, latest, rows = self._new_full_payload_accumulator(existing[0]) + else: + chunks, latest, rows, _ = existing + + for k, v in pooler_output.items(): + if v is None: + continue + if k in replace_keys: + # Explicit REPLACE semantics: the new value supersedes any + # prior chunks (e.g. `model_outputs` carries the full result + # so far, not an appendable per-step delta). + latest.pop(k, None) + if isinstance(v, torch.Tensor) and v.dim() >= 2: + chunks[k] = [v] + rows[k] = int(v.shape[0]) + else: + chunks.pop(k, None) + rows.pop(k, None) + latest[k] = v + continue + if isinstance(v, torch.Tensor) and v.dim() >= 2: + if k in chunks and chunks[k] and v.shape[1:] == chunks[k][0].shape[1:]: + chunks[k].append(v) + rows[k] += int(v.shape[0]) + else: + latest.pop(k, None) + chunks[k] = [v] + rows[k] = int(v.shape[0]) + else: + chunks.pop(k, None) + rows.pop(k, None) + latest[k] = v + + self._pending_full_payload_send[req_id] = (chunks, latest, rows, request) + + def flush_full_payload_outputs(self, finished_req_ids: set[str]) -> None: + """Send accumulated full_payload outputs for requests that just finished.""" + pending_req_ids = set(self._pending_full_payload_send.keys()) + if not (finished_req_ids & pending_req_ids): + return + + logger.debug( + "[Stage-%s] flush_full_payload_outputs: finished_req_ids=%s, pending=%s", + self._stage_id, + finished_req_ids, + list(self._pending_full_payload_send.keys()), + ) + to_send: dict[str, tuple[Any, Any]] = {} + for req_id in finished_req_ids: + entry = self._pending_full_payload_send.pop(req_id, None) + if entry is not None: + to_send[req_id] = self._materialize_full_payload_entry(entry) + logger.debug("[Stage-%s] flush_full_payload_outputs: to_send=%s", self._stage_id, list(to_send.keys())) + if to_send: + self.send_full_payload_outputs(scheduler_output=None, outputs=to_send) + + def send_full_payload_outputs( + self, + scheduler_output: Any, + outputs: dict[str, tuple[Any, Any] | Any], + ) -> list[str]: + """Send full_payload stage outputs to the next stage via connector. + + Args: + outputs: Mapping of ``req_id`` to either a + ``(pooling_output, request)`` tuple (preferred) or a raw + payload dict. When a tuple is supplied the request object + is forwarded to ``custom_process_stage_input_func``. + + Returns list of request IDs successfully enqueued. + """ + if self._omni_connector is None: + logger.debug("[Stage-%s] send_full_payload_outputs: connector is None, skip", self._stage_id) + return [] + if not self.is_data_transfer_rank(): + logger.debug( + "[Stage-%s] send_full_payload_outputs: not data_transfer_rank (rank=%s), skip", + self._stage_id, + self._local_rank, + ) + return list(outputs.keys()) + sent_ids: list[str] = [] + next_stage_id = self._next_stage_id + for req_id, value in outputs.items(): + if isinstance(value, tuple) and len(value) == 2: + raw_output, request = value + else: + raw_output, request = value, None + + payload = raw_output + if self._custom_process_func is not None: + payload = self._build_custom_process_payload( + request_id=req_id, + request=request, + pooling_output=raw_output, + ) + if payload is None: + continue + if payload is None: + logger.debug("[Stage-%s] send_full_payload_outputs: payload is None for %s", self._stage_id, req_id) + continue + if isinstance(payload, dict): + audio_codes = self._payload_audio_codes(payload) + if isinstance(audio_codes, torch.Tensor): + code_len = int(audio_codes.numel()) + elif hasattr(audio_codes, "__len__"): + code_len = len(audio_codes) + else: + code_len = None + meta = payload.get("meta") if isinstance(payload.get("meta"), dict) else {} + logger.debug( + "[Stage-%s] send_full_payload_outputs: req=%s payload_keys=%s code_len=%s left_context_size=%s", + self._stage_id, + req_id, + sorted(payload.keys()), + code_len, + meta.get("left_context_size"), + ) + + external_req_id = self._resolve_external_req_id(request, req_id) + chunk_id = self._put_req_chunk[req_id] + self._put_req_chunk[req_id] += 1 + self._ramp_chunk_count[req_id] += 1 + connector_put_key = f"{external_req_id}_{self._stage_id}_{chunk_id}" + + logger.debug( + "[Stage-%s] send_full_payload_outputs: enqueue req=%s put_key=%s next_stage=%s", + self._stage_id, + req_id, + connector_put_key, + next_stage_id, + ) + task = { + "stage_id": self._stage_id, + "next_stage_id": next_stage_id, + "put_key": connector_put_key, + "data": payload, + "request_id": req_id, + } + with self._lock: + self._pending_save_reqs.setdefault(req_id, deque()).append(task) + self._pending_save_counts[req_id] += 1 + sent_ids.append(req_id) + if sent_ids: + self._work_available.set() + return sent_ids + + # ------------------------------------------------------------------ # + # Streaming chunk mode (recv_chunk / send_chunk) + # ------------------------------------------------------------------ # + + def register_chunk_recv(self, request: Any) -> None: + """Register a request for async chunk retrieval by the bg thread. + + Stage-0 has no upstream producer so this is a no-op there. + Skips requests whose batch data has already been received to + prevent the bg thread from polling for non-existent chunks. + """ + if self._stage_id == 0: + return + request_id = request.request_id + # Explicit external_req_id=None must fall back to request_id; + # otherwise recv keys become `None__` and collide + # across requests. + ext = getattr(request, "external_req_id", None) + self._request_ids_mapping[request_id] = ext if ext is not None else request_id + with self._lock: + if request_id in self._stage_recv_req_ids: + return + # Don't re-register if the finish sentinel was already received + if request_id in self._chunk_stream_completed: + return + self._pending_load_reqs[request_id] = request + self._work_available.set() + + def recv_chunk(self) -> dict[str, Any]: + """Collect chunks received by the bg thread since last call. + + Returns a dict ``{request_id: chunk_payload}`` for newly arrived + chunks. Empty dict when nothing is ready. + + This method reads from ``_finished_load_reqs`` without clearing + it -- ``get_omni_connector_output()`` is the sole consumer that + drains and resets ``_finished_load_reqs`` at the end of each + ``execute_model`` cycle. + + Returns **shallow copies** of the cached payloads so that the + caller can read them without racing against the background recv + thread, which may concurrently mutate the live cache entries via + ``dict.update()``. + """ + with self._lock: + finished = set(self._finished_load_reqs) + if not finished: + return {} + # Snapshot the payloads under the lock to avoid racing with + # _poll_single_request which does existing.update(payload_data) + # on the same dict objects. + result = {} + for rid in finished: + payload = self._local_stage_payload_cache.get(rid) + result[rid] = dict(payload) if isinstance(payload, dict) else payload + + self._chunk_ready_req_ids.update(finished) + return result + + def send_chunk( + self, + request: Any, + pooling_output: Any | None = None, + ) -> bool: + """Derive and enqueue one chunk for async sending. + + Payload extraction runs in the caller thread (via + ``custom_process_stage_input_func``); the actual + ``connector.put()`` is done by the background save thread. + Non-KV data is identical across TP ranks; only rank 0 sends. + """ + if self._omni_connector is None: + logger.warning("[Stage-%s] send_chunk: connector is None", self._stage_id) + return False + if not self.is_data_transfer_rank(): + return True + raw_req_id = getattr(request, "request_id", None) or getattr(request, "req_id", None) + request_id = self._resolve_external_req_id(request, raw_req_id) + # Cache the internal→external mapping so that finish sentinels can + # resolve the external ID even after the request is freed. + if raw_req_id and raw_req_id != request_id: + self._request_ids_mapping.setdefault(raw_req_id, request_id) + chunk_id = self._put_req_chunk[request_id] + + payload_data = self._build_custom_process_payload( + request_id=request_id, + request=request, + pooling_output=pooling_output, + ) + if payload_data is None: + if chunk_id == 0: + logger.warning( + "[Stage-%s] send_chunk: payload is None for req=%s chunk=%s (process_func=%s)", + self._stage_id, + request_id, + chunk_id, + self._custom_process_func, + ) + return False + + self._put_req_chunk[request_id] += 1 + self._ramp_chunk_count[request_id] += 1 + next_stage_id = self._next_stage_id + connector_put_key = f"{request_id}_{self._stage_id}_{chunk_id}" + + if chunk_id == 0: + logger.debug( + "[Stage-%s] send_chunk: first chunk enqueued, req=%s key=%s", + self._stage_id, + request_id, + connector_put_key, + ) + + task = { + "stage_id": self._stage_id, + "next_stage_id": next_stage_id, + "put_key": connector_put_key, + "data": payload_data, + "request_id": request_id, + } + with self._lock: + self._pending_save_reqs.setdefault(request_id, deque()).append(task) + self._pending_save_counts[request_id] += 1 + self._work_available.set() + return True + + # ------------------------------------------------------------------ # + # Background I/O threads + # ------------------------------------------------------------------ # + + def _recv_loop(self) -> None: + """Background thread: poll connector for incoming data.""" + _recv_poll_count = 0 + while not self._stop_event.is_set(): + with self._lock: + pending_ids = list(self._pending_load_reqs.keys()) + + if not pending_ids: + self._work_available.wait(timeout=0.01) + self._work_available.clear() + continue + + _recv_poll_count += 1 + if _recv_poll_count % 5000 == 1: + logger.debug( + "[Stage-%s] _recv_loop: polling %s pending reqs: %s (poll#%s)", + self._stage_id, + len(pending_ids), + pending_ids[:5], + _recv_poll_count, + ) + + made_progress = False + for req_id in pending_ids: + if self._stop_event.is_set(): + break + try: + made_progress = self._poll_single_request(req_id) or made_progress + except Exception: + logger.warning("Error receiving data for %s", req_id, exc_info=True) + + if not made_progress and not self._stop_event.is_set(): + self._work_available.wait(timeout=0.005) + self._work_available.clear() + + _MAX_SEND_RETRIES = 3 + + def _save_loop(self) -> None: + """Background thread: send outgoing data via connector.""" + while not self._stop_event.is_set(): + task = None + with self._lock: + for req_id in list(self._pending_save_reqs.keys()): + dq = self._pending_save_reqs[req_id] + if dq: + task = dq.popleft() + if not dq: + del self._pending_save_reqs[req_id] + break + del self._pending_save_reqs[req_id] + + if task is not None: + success = False + try: + success = self._send_single_request(task) + except Exception: + logger.error( + "Error saving data for %s", + task.get("request_id"), + exc_info=True, + ) + if not success: + self._requeue_or_drop_failed_send(task) + continue + + self._work_available.wait(timeout=0.01) + self._work_available.clear() + + def _requeue_or_drop_failed_send(self, task: dict) -> None: + """Re-enqueue a failed send task or drop it after max retries.""" + retry_count = task.get("_retry_count", 0) + 1 + req_id = task.get("request_id") + if retry_count <= self._MAX_SEND_RETRIES: + task["_retry_count"] = retry_count + logger.warning( + "[Stage-%s] Re-enqueuing failed send for %s (retry %d/%d)", + getattr(self, "_stage_id", "?"), + req_id, + retry_count, + self._MAX_SEND_RETRIES, + ) + with self._lock: + dq = self._pending_save_reqs.setdefault(req_id, deque()) + dq.appendleft(task) + else: + logger.error( + "[Stage-%s] Giving up on send for %s after %d retries", + getattr(self, "_stage_id", "?"), + req_id, + self._MAX_SEND_RETRIES, + ) + self._decrement_pending_save_count(req_id) + + # ------------------------------------------------------------------ # + # Chunk-level poll / send (ported from OmniChunkTransferAdapter) + # ------------------------------------------------------------------ # + + def _poll_single_request(self, req_id: str) -> bool: + """Poll connector for one chunk of a request (non-blocking).""" + connector = self._omni_connector + if connector is None: + return False + + if self._async_chunk and self._model_mode != "ar": + with self._lock: + staged_payload = self._local_stage_payload_cache.get(req_id) + metadata_in_flight = req_id in self._local_request_metadata + scheduler_wakeup_pending = req_id in self._finished_load_reqs + if self._payload_is_consumable(staged_payload) or metadata_in_flight or scheduler_wakeup_pending: + logger.debug( + "[Stage-%s] delaying recv for req=%s until staged async payload is handed to scheduler", + self._stage_id, + req_id, + ) + return False + + target_stage_id = self._stage_id - 1 + chunk_id = self._get_req_chunk[req_id] + external_req_id = self._request_ids_mapping.get(req_id, req_id) + connector_get_key = f"{external_req_id}_{target_stage_id}_{chunk_id}" + + if self._async_chunk: + result = self._recv_async_chunk_result( + connector, + str(target_stage_id), + str(self._stage_id), + connector_get_key, + ) + else: + result = self._recv_full_payload_result( + connector, + str(target_stage_id), + str(self._stage_id), + connector_get_key, + ) + + if result is None: + return False + + payload_data, _size = result + if not payload_data: + return False + if isinstance(payload_data, dict): + logger.debug( + "[Stage-%s] recv_chunk_result: req=%s ext=%s key=%s keys=%s finished=%s", + self._stage_id, + req_id, + external_req_id, + connector_get_key, + sorted(payload_data.keys()), + self._payload_finished(payload_data), + ) + + self._get_req_chunk[req_id] += 1 + + if self._async_chunk: + is_finished = self._payload_finished(payload_data) + incoming_payload_consumable = self._payload_is_consumable(payload_data) + + if self._model_mode == "ar": + payload_data = self._accumulate_payload(external_req_id, payload_data) + payload_consumable = incoming_payload_consumable + else: + new_ids = self._payload_audio_codes(payload_data) or [] + if not new_ids and not is_finished: + return False + payload_consumable = self._payload_is_consumable(payload_data) + + with self._lock: + if is_finished: + self._chunk_finished_req_ids.add(req_id) + self._chunk_stream_completed.add(req_id) + # Local cache (RFC §2.4) — merge, don't replace, so that + # earlier chunk keys (e.g. thinker_prefill_embeddings from + # chunk 0) are not overwritten by later chunks. + existing = self._local_stage_payload_cache.get(req_id) + if existing is not None and isinstance(existing, dict) and isinstance(payload_data, dict): + existing.update(payload_data) + else: + self._local_stage_payload_cache[req_id] = payload_data + staged_payload = self._local_stage_payload_cache[req_id] + self._async_chunk_updated_req_ids.add(req_id) + self.put_local_request_metadata(req_id, self._extract_scheduling_metadata(staged_payload)) + # A finish-only sentinel still needs one terminal wake-up so + # the downstream stage can sync the merged local payload and + # flush/finish even when the last recv carries no new + # consumable chunk bytes. + if payload_consumable or is_finished: + self._finished_load_reqs.add(req_id) + if is_finished and not payload_consumable: + logger.debug( + "[Stage-%s] finish sentinel arrived for req=%s without new consumable payload", + self._stage_id, + req_id, + ) + elif not payload_consumable: + logger.debug( + "[Stage-%s] req=%s received metadata-only / non-consumable async payload; delaying wake-up", + self._stage_id, + req_id, + ) + if is_finished: + self._pending_load_reqs.pop(req_id, None) + else: + # full_payload_mode: the complete payload arrives in a single get(), + # so always unregister immediately. + if isinstance(payload_data, dict): + engine_inputs = payload_data.get("engine_inputs", payload_data) + else: + engine_inputs = payload_data + with self._lock: + self._local_stage_payload_cache[req_id] = self._snapshot_payload(engine_inputs) + # Publish full-payload readiness only after the aligned TP broadcast + # path in recv_full_payload_inputs() has materialized the payload on all + # local ranks. Publishing metadata / stage_recv from the background recv + # thread can let the scheduler observe a request before the payload is + # actually visible to the model thread. + self._full_payload_pending_broadcast_req_ids.add(req_id) + self._pending_load_reqs.pop(req_id, None) + logger.debug( + "[Stage-%s] full_payload recv complete: req=%s key=%s payload_type=%s", + self._stage_id, + req_id, + connector_get_key, + type(engine_inputs).__name__, + ) + + logger.debug("[Stage-%s] Received data for key %s", self._stage_id, connector_get_key) + return True + + def _build_custom_process_payload( + self, + request_id: str | None, + request: Any | None, + pooling_output: Any | None, + ) -> Any | None: + """Run the custom process hook with a best-effort finished kwarg.""" + if self._custom_process_func is None: + return None + + kwargs = { + "transfer_manager": self, + "pooling_output": pooling_output, + "request": request, + } + supports_is_finished = getattr( + self, + "_custom_process_supports_is_finished", + self._custom_process_supports_is_finished_kwarg(), + ) + is_finished_fn = getattr(request, "is_finished", None) + if callable(is_finished_fn): + try: + if supports_is_finished is not False: + kwargs["is_finished"] = bool(is_finished_fn()) + except Exception: + logger.debug("request.is_finished() failed for %s", request_id, exc_info=True) + + try: + return self._custom_process_func(**kwargs) + except TypeError as exc: + if "is_finished" not in kwargs or not self._is_unexpected_is_finished_kwarg_error(exc): + logger.exception("custom_process_stage_input_func failed for chunk %s", request_id) + return None + kwargs.pop("is_finished", None) + try: + return self._custom_process_func(**kwargs) + except Exception: + logger.exception("custom_process_stage_input_func failed for chunk %s", request_id) + return None + except Exception: + logger.exception("custom_process_stage_input_func failed for chunk %s", request_id) + return None + + def _custom_process_supports_is_finished_kwarg(self) -> bool | None: + """Return whether the custom process hook accepts `is_finished`.""" + if self._custom_process_func is None: + return None + try: + signature = inspect.signature(self._custom_process_func) + except (TypeError, ValueError): + return None + + for param in signature.parameters.values(): + if param.kind == inspect.Parameter.VAR_KEYWORD: + return True + + is_finished_param = signature.parameters.get("is_finished") + if is_finished_param is None: + return False + return is_finished_param.kind in ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ) + + @staticmethod + def _is_unexpected_is_finished_kwarg_error(exc: TypeError) -> bool: + message = str(exc) + return ( + "unexpected keyword argument 'is_finished'" in message + or 'unexpected keyword argument "is_finished"' in message + or "positional-only arguments passed as keyword arguments: 'is_finished'" in message + ) + + def _send_single_request(self, task: dict) -> bool: + """Send one queued task via connector.put(). + + Returns True on success. On failure (put() raises or returns + ``success=False``), returns False **without** decrementing + ``_pending_save_counts`` so the caller can retry or clean up. + """ + connector = self._omni_connector + if connector is None: + return True + + request_id = task.get("request_id") + payload_data = task.get("data") + if payload_data is None and task.get("request") is not None: + payload_data = self._build_custom_process_payload( + request_id=request_id, + request=task.get("request"), + pooling_output=task.get("pooling_output"), + ) + put_key = task.get("put_key") + + success, _size, _metadata = connector.put( + from_stage=str(task["stage_id"]), + to_stage=str(task["next_stage_id"]), + put_key=put_key, + data=payload_data, + ) + logger.debug( + "[Stage-%s] _send_single_request: put_key=%s success=%s size=%s", + task["stage_id"], + put_key, + success, + _size, + ) + + if not success: + return False + + self._decrement_pending_save_count(request_id) + return True + + def _decrement_pending_save_count(self, request_id: str) -> None: + """Decrement pending save count and run deferred cleanup if zero.""" + cleanup_req_id = None + with self._lock: + remaining = self._pending_save_counts.get(request_id, 0) + if remaining > 1: + self._pending_save_counts[request_id] = remaining - 1 + elif remaining == 1: + self._pending_save_counts.pop(request_id, None) + if request_id in self._deferred_send_cleanup: + self._deferred_send_cleanup.remove(request_id) + cleanup_req_id = request_id + if cleanup_req_id is not None: + self._put_req_chunk.pop(cleanup_req_id, None) + self._send_side_request_payload.pop(cleanup_req_id, None) + self._code_prompt_token_ids.pop(cleanup_req_id, None) + self._cached_ic.pop(cleanup_req_id, None) + self._ramp_chunk_count.pop(cleanup_req_id, None) + self._adaptive_states.pop(cleanup_req_id, None) + + # ------------------------------------------------------------------ # + # Payload accumulation (ported from OmniChunkTransferAdapter) + # ------------------------------------------------------------------ # + + def _accumulate_payload(self, req_id: str, payload_data: OmniPayload) -> OmniPayload: + """Accumulate chunk payloads (concat tensors, extend lists).""" + if req_id not in self._send_side_request_payload: + self._send_side_request_payload[req_id] = dict(payload_data) + return dict(self._send_side_request_payload[req_id]) + + origin = self._send_side_request_payload[req_id] + merged = dict(origin) + raw_ok = payload_data.get("meta", {}).get("override_keys", []) if isinstance(payload_data, dict) else [] + override_keys = {tuple(k) if isinstance(k, list) else k for k in raw_ok} + + for key, value in payload_data.items(): + if isinstance(value, dict): + origin_sub = origin.get(key) + merged_sub = dict(origin_sub) if isinstance(origin_sub, dict) else {} + for qual, qval in value.items(): + if key == "meta" and qual == "finished": + merged_sub[qual] = qval + continue + if (key, qual) in override_keys: + merged_sub[qual] = qval + continue + osv = merged_sub.get(qual) + if isinstance(qval, torch.Tensor) and isinstance(osv, torch.Tensor): + merged_sub[qual] = torch.cat([osv, qval], dim=0) + elif isinstance(qval, list) and isinstance(osv, list): + merged_sub[qual] = osv + qval + else: + merged_sub[qual] = qval + merged[key] = merged_sub + else: + if key in override_keys: + merged[key] = value + continue + ov = origin.get(key) + if isinstance(value, torch.Tensor) and isinstance(ov, torch.Tensor): + merged[key] = torch.cat([ov, value], dim=0) + elif isinstance(value, list) and isinstance(ov, list): + merged[key] = ov + value + else: + merged[key] = value + + self._send_side_request_payload[req_id] = merged + return dict(merged) diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py new file mode 100644 index 00000000000..275a2c0afe6 --- /dev/null +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py @@ -0,0 +1,1061 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Omni connector lifecycle and shared model-runner transport state.""" + +from __future__ import annotations + +import importlib +import inspect +import os +import threading +from collections import defaultdict, deque +from typing import TYPE_CHECKING, Any + +import torch +from vllm.logger import init_logger + +from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory +from vllm_omni.distributed.omni_connectors.utils.config import ( + ConnectorSpec, + get_stage_connector_role, +) + +logger = init_logger("vllm_omni.worker.omni_connector_model_runner_mixin") + +if TYPE_CHECKING: + from vllm_omni.config.model import OmniModelConfig + from vllm_omni.distributed.omni_connectors.connectors.base import ( + OmniConnectorBase, + ) + from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( + OmniKVTransferManager, + ) + + +def needs_omni_connector(model_config: Any) -> bool: + """Whether a runner owns an input, output, or explicitly routed connector.""" + return ( + bool(getattr(model_config, "requires_full_payload_input", False)) + or bool(getattr(model_config, "custom_process_next_stage_input_func", None)) + or get_stage_connector_role(model_config) is not None + ) + + +def _should_create_payload_connector(model_config: Any) -> bool: + """Whether this stage owns runner payload transport for its edge. + + Sender edges may instead be owned solely by KV transfer. Receivers still + need a connector even though they do not declare a downstream payload hook. + """ + if get_stage_connector_role(model_config) != "sender": + return True + + next_stage_func = getattr(model_config, "custom_process_next_stage_input_func", None) + return isinstance(next_stage_func, str) and bool(next_stage_func) + + +def should_accumulate_full_payload_output(model_config, custom_process_func) -> bool: + """Producer-side structural gate. + + Fires iff the stage explicitly declares a downstream full-payload + producer hook via ``custom_process_next_stage_input_func``. Consumer + stages may have ``custom_process_input_func`` values that can be + mechanically derived to ``*_full_payload`` helper names in the same + module; those are intentionally not enough to make the stage a producer. + """ + if custom_process_func is None: + return False + if getattr(model_config, "async_chunk", False): + return False + if getattr(model_config, "final_output", False): + return False + next_stage_func = getattr(model_config, "custom_process_next_stage_input_func", None) + if not isinstance(next_stage_func, str) or not next_stage_func: + return False + return getattr(model_config, "model_stage", None) is not None + + +class _OmniConnectorRuntimeMixin: + """Own connector lifecycle, shared state, and KV transfer delegation.""" + + _omni_connector: Any + _kv_transfer_manager: Any + _async_chunk: bool + _model_mode: str + _stage_id: int + _next_stage_id: int + _from_tp: int + _to_tp: int + _local_rank: int + _custom_process_func_path: str | None + _custom_process_func: Any + _custom_process_supports_is_finished: bool | None + _put_req_chunk: dict[str, int] + _get_req_chunk: dict[str, int] + _ramp_chunk_count: dict[str, int] + _adaptive_states: dict[str, Any] + _send_side_request_payload: dict[str, dict[str, Any]] + _code_prompt_token_ids: dict[str, list[list[int]]] + _cached_ic: dict[str, int] + _request_ids_mapping: dict[str, str] + _pending_load_reqs: dict[str, Any] + _finished_load_reqs: set[str] + _pending_save_reqs: dict[str, deque[Any]] + _pending_save_counts: dict[str, int] + _deferred_send_cleanup: set[str] + _chunk_ready_req_ids: set[str] + _chunk_finished_req_ids: set[str] + _stage_recv_req_ids: set[str] + _full_payload_pending_broadcast_req_ids: set[str] + _async_chunk_updated_req_ids: set[str] + _local_stage_payload_cache: dict[str, dict[str, Any]] + _local_request_metadata: dict[str, dict[str, Any]] + _chunk_stream_completed: set[str] + _pending_full_payload_send: dict[str, tuple[Any, ...]] + _kv_sent_req_ids: list[str] + _kv_pending_transfers: dict[str, dict[str, Any]] + _kv_active_transfers: set[str] + _kv_completed_transfers: set[str] + _kv_triggered_requests: set[str] + _lock: Any + _stop_event: Any + _work_available: Any + _recv_thread: Any + _save_thread: Any + _omni_connector_initialized: bool + _full_payload_replace_keys_cached: frozenset[Any] + _should_accumulate_full_payload_output_cached: bool + + vllm_config: Any + requests: dict[str, Any] + model_intermediate_buffer: dict[str, Any] + + _recv_loop: Any + _save_loop: Any + flush_full_payload_outputs: Any + _custom_process_supports_is_finished_kwarg: Any + _get_local_tp_group: Any + + # ------------------------------------------------------------------ # + # Init / Shutdown + # ------------------------------------------------------------------ # + + def init_omni_connectors( + self, + model_config: OmniModelConfig, + kv_transfer_manager: OmniKVTransferManager | None = None, + ) -> None: + """Initialize connectors and background threads. + + Args: + model_config: Stage-level model config with connector settings. + kv_transfer_manager: Existing KV transfer manager to delegate to. + """ + self._omni_connector: OmniConnectorBase | None = ( + self._create_connector(model_config) if _should_create_payload_connector(model_config) else None + ) + self._kv_transfer_manager = kv_transfer_manager + + self._async_chunk: bool = getattr(model_config, "async_chunk", False) + self._model_mode: str = getattr(model_config, "worker_type", "ar") + stage_id = getattr(model_config, "stage_id", 0) + if isinstance(stage_id, str): + stage_id = int(stage_id) + self._stage_id: int = stage_id if isinstance(stage_id, int) else 0 + + self._custom_process_func_path, self._custom_process_func = self._load_custom_func(model_config) + self._custom_process_supports_is_finished = self._custom_process_supports_is_finished_kwarg() + logger.debug( + "[Stage-%s] init_omni_connectors: async_chunk=%s, custom_process_func=%s, connector=%s, func_path=%s", + self._stage_id, + self._async_chunk, + self._custom_process_func, + type(self._omni_connector).__name__ if self._omni_connector else None, + self._custom_process_func_path, + ) + + # -- next stage ID (from connector config or default stage_id + 1) -- + self._next_stage_id: int = self._resolve_next_stage_id(model_config) + + # -- heterogeneous TP rank support -- + rank_cfg = self._parse_rank_mapping(model_config) + if self._kv_transfer_manager is not None: + topology = getattr(self._kv_transfer_manager, "tp_topology", None) + effective_mapping = ( + getattr(topology, "source_tp_size", None), + getattr(topology, "target_tp_size", None), + getattr(topology, "local_rank", None), + ) + if all(isinstance(value, int) for value in effective_mapping): + rank_cfg = { + "from_tp": effective_mapping[0], + "to_tp": effective_mapping[1], + "local_rank": effective_mapping[2], + } + self._from_tp: int = rank_cfg["from_tp"] + self._to_tp: int = rank_cfg["to_tp"] + self._local_rank: int = rank_cfg["local_rank"] + if self._kv_transfer_manager is not None: + self._kv_transfer_manager.kv_send_key_builder = self.get_rank_aware_kv_send_keys + self._kv_transfer_manager.kv_recv_key_builder = self.get_rank_aware_kv_keys + self._kv_transfer_manager.kv_payload_merger = self._merge_rank_sharded_kv_payloads + self._kv_transfer_manager.kv_payload_slicer = self._slice_rank_sharded_kv_payload + + # -- chunk index tracking (ported from OmniChunkTransferAdapter) -- + self._put_req_chunk: dict[str, int] = defaultdict(int) + self._get_req_chunk: dict[str, int] = defaultdict(int) + # Segment-local chunk counter: incremented alongside _put_req_chunk + # and popped at request cleanup. Note: the mixin path (uniproc mode) + # does not have segment boundary infrastructure; multi-segment support + # is only available via chunk_transfer_adapter (distributed path). + self._ramp_chunk_count: dict[str, int] = defaultdict(int) + self._adaptive_states: dict[str, Any] = {} + # Send-side async accumulation / staging buffer. Receive-side payload + # ownership lives in ``_local_stage_payload_cache``. + self._send_side_request_payload: dict[str, dict[str, Any]] = {} + self._code_prompt_token_ids: dict[str, list[list[int]]] = defaultdict(list) + self._cached_ic: dict[str, int] = {} + self._request_ids_mapping: dict[str, str] = {} + + # -- async I/O state (shared by chunk + full_payload_mode) -- + self._pending_load_reqs: dict[str, Any] = {} + self._finished_load_reqs: set[str] = set() + self._pending_save_reqs: dict[str, deque] = {} + self._pending_save_counts: dict[str, int] = defaultdict(int) + self._deferred_send_cleanup: set[str] = set() + # -- per-cycle output accumulator -- + self._chunk_ready_req_ids: set[str] = set() + self._chunk_finished_req_ids: set[str] = set() + self._stage_recv_req_ids: set[str] = set() + self._full_payload_pending_broadcast_req_ids: set[str] = set() + self._async_chunk_updated_req_ids: set[str] = set() + + # -- Model Runner local payload cache (RFC §2.4) -- + # Full stage payloads land here first on the recv side. We + # intentionally do not write connector recv results straight into + # `model_intermediate_buffer`: runner-owned runtime state is + # materialized later by `_sync_local_stage_payloads()` on the + # model thread. This keeps recv timing separate from execute-step + # visibility and avoids mixing connector I/O with model runtime + # ownership. + self._local_stage_payload_cache: dict[str, dict[str, Any]] = {} + # Lightweight scheduling metadata pending delivery to the Scheduler. + self._local_request_metadata: dict[str, dict[str, Any]] = {} + + # -- persistent set of request IDs whose chunk stream is complete -- + # Prevents re-registration after the finish sentinel has been received. + self._chunk_stream_completed: set[str] = set() + + # -- full_payload_mode: accumulate latest pooler_output per request, + # send only when the request finishes (next-cycle flush) -- + self._pending_full_payload_send: dict[str, tuple[Any, ...]] = {} + + # -- KV sent accumulator -- + self._kv_sent_req_ids: list[str] = [] + + # -- KV transfer lifecycle (absorbed from scheduler) -- + # Requests marked for KV transfer: {req_id: {seq_len, block_ids}} + self._kv_pending_transfers: dict[str, dict[str, Any]] = {} + # Requests whose KV transfer has been submitted but not yet acked + self._kv_active_transfers: set[str] = set() + # Requests whose KV transfer is complete (acked by kv_extracted_req_ids) + self._kv_completed_transfers: set[str] = set() + # Dedup guard: requests that have already triggered KV transfer + self._kv_triggered_requests: set[str] = set() + + self._lock = threading.Lock() + self._stop_event = threading.Event() + self._work_available = threading.Event() + + # Start background threads only when there's a connector + self._recv_thread: threading.Thread | None = None + self._save_thread: threading.Thread | None = None + if self._omni_connector is not None: + self._recv_thread = threading.Thread( + target=self._recv_loop, + daemon=True, + name="omni-mixin-recv", + ) + self._recv_thread.start() + self._save_thread = threading.Thread( + target=self._save_loop, + daemon=True, + name="omni-mixin-save", + ) + self._save_thread.start() + + # Explicit "fully initialised" marker so other parts of the runner + # (e.g. _update_states cleanup) can branch on a stable contract + # instead of probing for private mixin attribute names. Must be set + # only after every field above has been bound, so a partially + # constructed mixin is never observable as initialised. + self._omni_connector_initialized = True + + def shutdown_omni_connectors(self) -> None: + """Stop background threads and release connector resources.""" + self._stop_event.set() + if self._recv_thread is not None: + self._recv_thread.join(timeout=5) + if self._save_thread is not None: + self._save_thread.join(timeout=5) + if self._omni_connector is not None: + try: + self._omni_connector.close() + except Exception: + pass + + def cleanup_finished_request(self, req_id: str) -> None: + """Clean up per-request state after a request is fully finished. + + Call this when a request is freed from the model runner to prevent + memory leaks in the mixin's tracking dicts/sets. + + Two senders use different keys: ``send_chunk`` keys per-request + state under the EXTERNAL id (after mapping resolution), while + ``send_full_payload_outputs`` keys under the INTERNAL id. To cover + both modes (and forward compat with id-rename scenarios) we attempt + cleanup against both keys; the entry that doesn't exist for the + active mode is a no-op pop. Only the key that actually has pending + saves is added to ``_deferred_send_cleanup`` so the bg save's + decrement path drains it without leaving orphans. + """ + # Force-flush any pending full-payload accumulator entry before + # cleanup proceeds. Without this, finished requests with no + # downstream consumer (e.g. text-only on multi-modal arch) leave + # the entry orphaned in _pending_full_payload_send across requests, + # which empirically destabilises subsequent thinker forwards by + # making prefix-cache reuse observe stale accumulator state. The + # flush is idempotent when the entry has already been flushed by the + # scheduler-driven path, but this cleanup path runs for every request, + # so skip it entirely when the request never accumulated a payload. + if req_id in self._pending_full_payload_send: + try: + self.flush_full_payload_outputs({req_id}) + except Exception: + # Cleanup must still proceed regardless of flush errors here -- + # we already gated on ``_omni_connector_initialized`` upstream, + # so any exception here reflects a real connector-side issue + # (shared memory corruption, background thread crash) worth + # surfacing rather than silently swallowing. + logger.warning( + "flush_full_payload_outputs(%s) raised during cleanup; continuing tear-down.", + req_id, + exc_info=True, + ) + + ext_id = self._request_ids_mapping.pop(req_id, None) + keys_to_clean: list[str] = [req_id] + if ext_id is not None and ext_id != req_id: + keys_to_clean.append(ext_id) + + with self._lock: + keys_pending = [k for k in keys_to_clean if self._pending_save_counts.get(k, 0)] + for k in keys_pending: + self._deferred_send_cleanup.add(k) + for k in keys_to_clean: + if k in keys_pending: + continue + self._put_req_chunk.pop(k, None) + self._send_side_request_payload.pop(k, None) + self._code_prompt_token_ids.pop(k, None) + self._cached_ic.pop(k, None) + self._ramp_chunk_count.pop(k, None) + self._adaptive_states.pop(k, None) + self._kv_pending_transfers.pop(req_id, None) + self._kv_active_transfers.discard(req_id) + self._kv_completed_transfers.discard(req_id) + self._kv_triggered_requests.discard(req_id) + self._cleanup_recv_delivery_state(req_id) + + def drop_inactive_request_delivery_state(self, req_id: str) -> None: + """Clear recv-side state for inactive requests.""" + ext_id = self._request_ids_mapping.pop(req_id, None) + if hasattr(self, "_lock"): + with self._lock: + self._drop_send_side_payload_state(req_id, ext_id) + else: + self._drop_send_side_payload_state(req_id, ext_id) + self._cleanup_recv_delivery_state(req_id) + + def _drop_send_side_payload_state(self, req_id: str, ext_id: str | None) -> None: + if ext_id is not None: + self._send_side_request_payload.pop(ext_id, None) + self._cached_ic.pop(ext_id, None) + self._send_side_request_payload.pop(req_id, None) + self._cached_ic.pop(req_id, None) + + def _cleanup_recv_delivery_state(self, req_id: str) -> None: + """Clear recv-side delivery-cycle state.""" + if hasattr(self, "_lock"): + with self._lock: + self._clear_recv_delivery_state(req_id) + else: + self._clear_recv_delivery_state(req_id) + + def _clear_recv_delivery_state(self, req_id: str) -> None: + self._get_req_chunk.pop(req_id, None) + self._pending_load_reqs.pop(req_id, None) + self._finished_load_reqs.discard(req_id) + self._chunk_ready_req_ids.discard(req_id) + self._chunk_finished_req_ids.discard(req_id) + self._chunk_stream_completed.discard(req_id) + self._stage_recv_req_ids.discard(req_id) + self._full_payload_pending_broadcast_req_ids.discard(req_id) + self._async_chunk_updated_req_ids.discard(req_id) + self._local_stage_payload_cache.pop(req_id, None) + self._local_request_metadata.pop(req_id, None) + + def prune_inactive_requests(self, active_req_ids: Any) -> set[str]: + """Drop connector state for requests that no longer exist locally. + + Preempted / unscheduled requests are expected to stay in + ``self.requests`` and therefore remain untouched. This only prunes + stale request IDs that have already fallen out of the active request + map, preventing background recv/send bookkeeping from outliving the + request lifecycle. + """ + if active_req_ids is None: + return set() + + active_req_ids = set(active_req_ids) + pending_req_ids = set(getattr(self, "_pending_load_reqs", {}).keys()) + received_req_ids = set(getattr(self, "_stage_recv_req_ids", set())) + received_req_ids.update(getattr(self, "_full_payload_pending_broadcast_req_ids", set())) + received_req_ids.update(getattr(self, "_local_request_metadata", {}).keys()) + # Pending recv requests may not yet be in the caller's active set + # (e.g. WAITING_FOR_CHUNK requests live in the coordinator's internal + # queues, not in model runner self.requests). Protect them so that + # legitimate waiting requests are not pruned. + # + # Likewise, a full payload can arrive on the background recv thread + # after the scheduler_output snapshot for the current execute_model() + # cycle was already materialized. Those requests may briefly live only + # in recv-side buffers/local cache until the next scheduler cycle wakes + # them up; pruning them here drops the payload before stage_recv can be + # published. + active_req_ids.update(pending_req_ids) + active_req_ids.update(received_req_ids) + stale_req_ids: set[str] = set() + + # NOTE: _pending_load_reqs is excluded from the scan list because + # all its entries are unconditionally protected above. The mixin + # cannot distinguish a legitimately-waiting pending recv from an + # orphaned one (only the coordinator/scheduler knows). + # + # Requests with freshly received full payloads / local stage payloads + # are also protected above. Their scheduler wake-up may lag the recv + # thread by one execute_model() cycle, especially when the request was + # added after the current scheduler_output snapshot. + # + # Orphaned pending recv entries (e.g. from upstream stage crash) are + # handled by collect_timed_out_request_ids() -- on + # OmniSchedulingCoordinator for full-payload requests, and on + # OmniChunkTransferAdapter for async-chunk ones -- which detect + # wait-time violations. The scheduler then removes the request from + # its queues, sets FINISHED_ERROR, and calls _free_request() which + # ultimately triggers cleanup_finished_request() here. + for attr_name in ( + "_request_ids_mapping", + "_get_req_chunk", + "_finished_load_reqs", + "_chunk_ready_req_ids", + "_chunk_finished_req_ids", + "_chunk_stream_completed", + "_stage_recv_req_ids", + "_full_payload_pending_broadcast_req_ids", + "_async_chunk_updated_req_ids", + "_local_stage_payload_cache", + "_local_request_metadata", + "_kv_pending_transfers", + "_kv_active_transfers", + "_kv_completed_transfers", + "_kv_triggered_requests", + ): + state = getattr(self, attr_name, None) + if isinstance(state, dict): + stale_req_ids.update(req_id for req_id in state if req_id not in active_req_ids) + elif isinstance(state, set): + stale_req_ids.update(req_id for req_id in state if req_id not in active_req_ids) + + for req_id in stale_req_ids: + self.cleanup_finished_request(req_id) + + return stale_req_ids + + def drop_inactive_request_runtime_state(self, req_id: str) -> None: + """Clear inactive request state used by both the runner and mixin. + + This centralizes the model-runner-side cleanup pattern so + ``OmniGPUModelRunner`` can reuse it instead of open-coding the same + inactive-request state mutations. + """ + if hasattr(self, "model_intermediate_buffer"): + self.model_intermediate_buffer.pop(req_id, None) + self.drop_inactive_request_delivery_state(req_id) + + # ------------------------------------------------------------------ # + # Helpers + # ------------------------------------------------------------------ # + + @staticmethod + def _create_connector(model_config: Any) -> OmniConnectorBase | None: + """Create a connector from model_config, or None if unconfigured.""" + connector_config = getattr(model_config, "stage_connector_config", None) + if connector_config is None: + return None + + if not isinstance(connector_config, dict): + connector_config = { + "name": getattr(connector_config, "name", None), + "extra": getattr(connector_config, "extra", None), + } + + name = connector_config.get("name") + if not isinstance(name, str) or not name.strip(): + raise RuntimeError("Invalid stage connector config: missing connector name") + name = name.strip() + + extra = connector_config.get("extra") + if extra is None: + extra = {} + elif not isinstance(extra, dict): + raise RuntimeError(f"Invalid extra config for connector {name}: expected dict, got {type(extra).__name__}") + + spec = ConnectorSpec(name=name, extra=extra) + try: + return OmniConnectorFactory.create_connector(spec) + except Exception as exc: + raise RuntimeError(f"Failed to create connector {name}") from exc + + @classmethod + def _load_custom_func(cls, model_config: Any) -> tuple[str | None, Any | None]: + """Load the connector payload builder for the downstream stage. + + Preferred source is ``custom_process_next_stage_input_func``. Some + full_payload_mode configs (async_chunk=false) only expose the next-stage prompt builder via + ``custom_process_input_func`` (for example ``thinker2talker``), while the + connector payload builder lives beside it as ``thinker2talker_full_payload``. + In that case, derive the full_payload_mode builder path automatically. + """ + candidates: list[str] = [] + + next_stage_func = getattr(model_config, "custom_process_next_stage_input_func", None) + if isinstance(next_stage_func, str) and next_stage_func: + candidates.append(next_stage_func) + + if not getattr(model_config, "async_chunk", False): + input_func = getattr(model_config, "custom_process_input_func", None) + if isinstance(input_func, str) and input_func: + try: + module_path, func_name = input_func.rsplit(".", 1) + if func_name.endswith("_full_payload") or func_name.endswith("_batch"): + candidates.append(f"{module_path}.{func_name}") + else: + candidates.append(f"{module_path}.{func_name}_full_payload") + candidates.append(f"{module_path}.{func_name}_batch") + candidates.append(input_func) + except ValueError: + candidates.append(input_func) + + tried: set[str] = set() + for func_path in candidates: + if func_path in tried: + continue + tried.add(func_path) + try: + module_path, func_name = func_path.rsplit(".", 1) + module = importlib.import_module(module_path) + func = getattr(module, func_name, None) + if callable(func): + if not cls._is_connector_payload_builder(func): + logger.debug( + "Skipping incompatible connector payload hook %s; signature=%s", + func_path, + inspect.signature(func), + ) + continue + return func_path, func + except Exception: + logger.warning("Failed to load custom func: %s", func_path, exc_info=True) + + return None, None + + @staticmethod + def _is_connector_payload_builder(func: Any) -> bool: + """Whether *func* matches the mixin payload-builder contract.""" + try: + signature = inspect.signature(func) + except (TypeError, ValueError): + return False + + params = signature.parameters + if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in params.values()): + return True + + required = {"transfer_manager", "pooling_output", "request"} + supported = { + name + for name, param in params.items() + if param.kind + in ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ) + } + return required.issubset(supported) + + def _resolve_external_req_id(self, request: Any, fallback_req_id: str) -> str: + """Resolve the external request ID consistently. + + Checks ``_request_ids_mapping`` first (populated by + ``register_chunk_recv``), then falls back to the request's + ``external_req_id`` attribute, and finally to the given + ``fallback_req_id``. + """ + mapped = self._request_ids_mapping.get(fallback_req_id) + if mapped is not None: + return mapped + if request is not None: + # external_req_id may be explicitly None; fall back. + ext = getattr(request, "external_req_id", None) + if ext is not None: + return ext + return fallback_req_id + + def _resolve_next_stage_id(self, model_config: Any) -> int: + """Determine the downstream stage ID from connector config. + + Falls back to ``stage_id + 1`` when the config does not specify + a ``to_stage`` explicitly. + """ + connector_config = getattr(model_config, "stage_connector_config", None) + if connector_config is not None: + if isinstance(connector_config, dict): + to_stage = connector_config.get("to_stage") + else: + to_stage = getattr(connector_config, "to_stage", None) + if isinstance(to_stage, int): + return to_stage + if isinstance(to_stage, str) and to_stage.strip(): + return int(to_stage) + return self._stage_id + 1 + + @staticmethod + def _parse_rank_mapping(model_config: Any) -> dict[str, int]: + """Parse rank_mapping from connector config (optional). + + Returns ``{"from_tp": int, "to_tp": int, "local_rank": int}``. + When ``rank_mapping`` is absent, assumes 1:1 homogeneous mapping. + """ + connector_config = getattr(model_config, "stage_connector_config", None) + if connector_config is not None and not isinstance(connector_config, dict): + connector_config = getattr(connector_config, "__dict__", {}) + + rank_mapping: dict = {} + if isinstance(connector_config, dict): + rank_mapping = connector_config.get("rank_mapping", {}) + + from_tp = int(rank_mapping.get("from_tp", 1)) + to_tp = int(rank_mapping.get("to_tp", 1)) + + local_rank = 0 + try: + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + except (ValueError, TypeError): + pass + + return {"from_tp": from_tp, "to_tp": to_tp, "local_rank": local_rank} + + # ------------------------------------------------------------------ # + # KV cache (delegates to OmniKVTransferManager) + # ------------------------------------------------------------------ # + + def send_kv_cache( + self, + finished_reqs: dict[str, dict[str, Any]], + kv_caches: list[torch.Tensor], + block_size: int, + cache_dtype: str, + request_id_resolver: Any | None = None, + ) -> list[str]: + """Send KV cache for finished requests. + + Delegates to the existing ``OmniKVTransferManager``. + """ + if self._kv_transfer_manager is None: + return list(finished_reqs.keys()) if finished_reqs else [] + result = self._kv_transfer_manager.handle_finished_requests_kv_transfer( + finished_reqs=finished_reqs, + kv_caches=kv_caches, + block_size=block_size, + cache_dtype=cache_dtype, + request_id_resolver=request_id_resolver, + ) + if result: + self._kv_sent_req_ids.extend(result) + return result + + def recv_kv_cache( + self, + request_id: str, + target_device: torch.device | None = None, + ) -> tuple[dict[str, Any] | None, int]: + """Receive KV cache for a request. + + Delegates to the existing ``OmniKVTransferManager``. + """ + if self._kv_transfer_manager is None: + return None, 0 + return self._kv_transfer_manager.receive_kv_cache_for_request( + request_id=request_id, + target_device=target_device, + ) + + def receive_cfg_companion_kv_payloads( + self, + cfg_request_ids: dict[str, str], + target_device: torch.device | None = None, + ) -> dict[str, tuple[dict[str, Any] | None, int]]: + """Receive raw CFG companion KV payloads keyed by role.""" + return { + role: self.recv_kv_cache(companion_rid, target_device=target_device) + for role, companion_rid in cfg_request_ids.items() + } + + def receive_multi_kv_cache( + self, + req: Any, + cfg_kv_collect_func: Any | None = None, + target_device: torch.device | None = None, + ) -> bool: + """Receive primary and optional companion KV caches for a request. + + The mixin owns the runner-facing orchestration: primary KV receive, + companion payload fetch, and applying any model-specific CFG fields back + onto ``req.sampling_params``. + """ + if self._kv_transfer_manager is None: + return False + + request_id = getattr(req, "request_id", None) + if not request_id: + logger.warning("Request has no ID, cannot receive KV cache") + return False + + active_requests = getattr(self, "requests", None) + if active_requests is not None and request_id not in active_requests: + logger.debug("Skip receiving KV cache for inactive request %s", request_id) + return False + + primary_ok = False + data, _size = self.recv_kv_cache(request_id, target_device=target_device) + if data: + self._kv_transfer_manager.apply_kv_cache_to_request(req, data) + primary_ok = True + + cfg_ids = getattr(getattr(req, "sampling_params", None), "cfg_kv_request_ids", None) + if cfg_ids and cfg_kv_collect_func: + try: + cfg_role_payloads = self.receive_cfg_companion_kv_payloads( + cfg_ids, + target_device=target_device, + ) + cfg_kvs = cfg_kv_collect_func(request_id, cfg_role_payloads) + if cfg_kvs and hasattr(req, "sampling_params") and req.sampling_params is not None: + for key, value in cfg_kvs.items(): + setattr(req.sampling_params, key, value) + logger.debug("Applied CFG KV caches: %s", list(cfg_kvs.keys())) + except Exception: + logger.exception("Failed to collect CFG KV caches for %s", request_id) + + return primary_ok + + # ------------------------------------------------------------------ # + # Rank-aware KV transfer routing + # ------------------------------------------------------------------ # + + def get_rank_aware_kv_keys( + self, + req_id: str, + from_stage: int, + to_stage: int | None = None, + chunk_id: int = 0, + ) -> list[str]: + """Build recv-side connector keys for all remote ranks this rank needs. + + For heterogeneous TP receive, the local rank is the target rank and must + fetch one or more source-rank shards keyed as ``from_rank -> to_rank``. + """ + if self._from_tp <= 1 and self._to_tp <= 1: + resolved_to_stage = self._next_stage_id if to_stage is None else to_stage + return [f"omni_{from_stage}_to_{resolved_to_stage}_kv_cache_{req_id}"] + + remote_ranks = self.get_kv_remote_ranks() + return [ + self.get_kv_connector_key( + req_id=req_id, + from_stage=from_stage, + chunk_id=chunk_id, + from_rank=remote_rank, + to_rank=self._local_rank, + ) + for remote_rank in remote_ranks + ] + + def get_kv_target_ranks_for_send(self) -> list[int]: + """Determine which target ranks this local rank should send KV shards to.""" + self._validate_kv_tp_topology() + if self._from_tp == self._to_tp: + return [self._local_rank] + if self._from_tp > self._to_tp: + tp_ratio = self._from_tp // self._to_tp + return [self._local_rank // tp_ratio] + tp_ratio = self._to_tp // self._from_tp + base_rank = self._local_rank * tp_ratio + return [base_rank + i for i in range(tp_ratio)] + + def get_rank_aware_kv_send_keys( + self, + req_id: str, + from_stage: int, + to_stage: int | None = None, + chunk_id: int = 0, + ) -> list[str]: + """Build send-side connector keys for this rank's KV shard(s).""" + if self._from_tp <= 1 and self._to_tp <= 1: + resolved_to_stage = self._next_stage_id if to_stage is None else to_stage + return [f"omni_{from_stage}_to_{resolved_to_stage}_kv_cache_{req_id}"] + + target_ranks = self.get_kv_target_ranks_for_send() + return [ + self.get_kv_connector_key( + req_id=req_id, + from_stage=from_stage, + chunk_id=chunk_id, + from_rank=self._local_rank, + to_rank=target_rank, + ) + for target_rank in target_ranks + ] + + @staticmethod + def _merge_rank_sharded_kv_payloads(payloads: list[dict[str, Any]]) -> dict[str, Any] | None: + """Merge multiple source-rank KV shards for one target rank.""" + payloads = [payload for payload in payloads if isinstance(payload, dict)] + if not payloads: + return None + if len(payloads) == 1: + return payloads[0] + + merged = dict(payloads[0]) + layer_blocks = merged.get("layer_blocks") + if not isinstance(layer_blocks, dict): + return merged + + def _merge_tensor_lists(name: str) -> list[torch.Tensor | None]: + merged_list: list[torch.Tensor | None] = [] + cache_lists = [payload.get("layer_blocks", {}).get(name, []) for payload in payloads] + max_len = max((len(cache_list) for cache_list in cache_lists), default=0) + for idx in range(max_len): + tensors = [cache_list[idx] for cache_list in cache_lists if idx < len(cache_list)] + tensors = [tensor for tensor in tensors if isinstance(tensor, torch.Tensor)] + if not tensors: + merged_list.append(None) + elif len(tensors) == 1: + merged_list.append(tensors[0]) + else: + merged_list.append(torch.cat(tensors, dim=-2).contiguous()) + return merged_list + + merged["layer_blocks"] = { + "key_cache": _merge_tensor_lists("key_cache"), + "value_cache": _merge_tensor_lists("value_cache"), + } + metadata = dict(merged.get("metadata", {})) + metadata["merged_remote_rank_count"] = len(payloads) + merged["metadata"] = metadata + return merged + + def _slice_rank_sharded_kv_payload(self, payload: dict[str, Any] | None) -> dict[str, Any] | None: + """Slice a duplicated source-rank KV shard for ``from_tp < to_tp`` cases.""" + if payload is None or self._from_tp >= self._to_tp: + return payload + + tp_ratio = self._to_tp // self._from_tp + shard_index = self._local_rank % tp_ratio + layer_blocks = payload.get("layer_blocks") if isinstance(payload, dict) else None + if not isinstance(layer_blocks, dict): + return payload + + def _slice_tensor_list(name: str) -> list[torch.Tensor | None]: + sliced: list[torch.Tensor | None] = [] + for tensor in layer_blocks.get(name, []): + if not isinstance(tensor, torch.Tensor) or tensor.ndim < 2: + sliced.append(tensor) + continue + head_dim = tensor.shape[-2] + if head_dim % tp_ratio != 0: + sliced.append(tensor) + continue + per_rank = head_dim // tp_ratio + start = shard_index * per_rank + sliced.append(tensor.narrow(-2, start, per_rank).contiguous()) + return sliced + + payload = dict(payload) + payload["layer_blocks"] = { + "key_cache": _slice_tensor_list("key_cache"), + "value_cache": _slice_tensor_list("value_cache"), + } + metadata = dict(payload.get("metadata", {})) + metadata["sliced_for_local_rank"] = self._local_rank + payload["metadata"] = metadata + return payload + + def should_replicate_payload(self) -> bool: + """Whether non-KV payloads should be replicated across ranks. + + Data payloads (stage inputs, chunks) are identical after all-gather, + so only rank 0 transfers them. KV payloads are rank-specific and + all ranks participate. + """ + return self._local_rank != 0 + + def get_kv_rank_mapping(self) -> dict[str, Any]: + """Return the current rank mapping configuration. + + Useful for debugging and for downstream code that needs to know + the TP topology without re-parsing model config. + """ + return { + "from_tp": self._from_tp, + "to_tp": self._to_tp, + "local_rank": self._local_rank, + "remote_ranks": self.get_kv_remote_ranks(), + "is_data_transfer_rank": self.is_data_transfer_rank(), + } + + # ------------------------------------------------------------------ # + # KV transfer lifecycle (RFC – mixin-owned) + # ------------------------------------------------------------------ # + + def mark_kv_transfer( + self, + req_id: str, + seq_len: int, + block_ids: list[int], + custom_metadata: dict[str, Any] | None = None, + ) -> None: + """Mark a request as needing KV cache transfer. + + Called by the scheduler when a transfer trigger fires. The mixin + owns the lifecycle from this point: pending → active → completed. + """ + if req_id in self._kv_pending_transfers: + return + self._kv_triggered_requests.add(req_id) + transfer = { + "seq_len": seq_len, + "block_ids": block_ids, + } + if custom_metadata is not None: + transfer["custom_metadata"] = custom_metadata + self._kv_pending_transfers[req_id] = transfer + + def drain_pending_kv_transfers(self) -> dict[str, dict[str, Any]]: + """Drain pending KV transfers and move them to active. + + Returns ``{req_id: {seq_len, block_ids}}`` for the model runner + to submit to ``send_kv_cache``. + """ + if not self._kv_pending_transfers: + return {} + pending = dict(self._kv_pending_transfers) + self._kv_active_transfers.update(pending.keys()) + self._kv_pending_transfers.clear() + return pending + + def ack_kv_transfers(self, req_ids: list[str] | set[str]) -> None: + """Acknowledge completed KV transfers (from kv_extracted_req_ids). + + Moves requests from active to completed so the scheduler can + safely free their blocks. + """ + for req_id in req_ids: + self._kv_active_transfers.discard(req_id) + self._kv_completed_transfers.add(req_id) + + def drain_completed_kv_transfers(self) -> set[str]: + """Drain and return completed KV transfer request IDs. + + The scheduler calls this to know which requests' blocks can be freed. + """ + completed = set(self._kv_completed_transfers) + self._kv_completed_transfers.clear() + return completed + + def is_kv_transfer_triggered(self, req_id: str) -> bool: + """Check if a request has already triggered KV transfer.""" + return req_id in self._kv_triggered_requests + + def has_pending_kv_work(self) -> bool: + """True if any KV transfers are pending, active, or awaiting ack.""" + return bool(self._kv_pending_transfers or self._kv_active_transfers or self._kv_completed_transfers) + + # ------------------------------------------------------------------ # + # Heterogeneous TP rank support + # ------------------------------------------------------------------ # + + def _validate_kv_tp_topology(self) -> None: + """Reject heterogeneous TP mappings that cannot be routed losslessly.""" + if self._from_tp <= 0 or self._to_tp <= 0: + raise ValueError(f"Invalid KV TP mapping: from_tp={self._from_tp}, to_tp={self._to_tp}") + larger = max(self._from_tp, self._to_tp) + smaller = min(self._from_tp, self._to_tp) + if larger % smaller != 0: + raise ValueError( + f"KV TP mapping must be divisible for rank-aware routing: from_tp={self._from_tp}, to_tp={self._to_tp}" + ) + + def get_kv_remote_ranks(self) -> list[int]: + """Determine which remote ranks this local rank exchanges KV with. + + Follows vLLM's ``TpKVTopology.get_target_remote_ranks()`` pattern: + - ``from_tp > to_tp``: each to-rank reads from multiple from-ranks + - ``from_tp < to_tp``: multiple to-ranks read from the same from-rank + - ``from_tp == to_tp``: 1:1 mapping + """ + self._validate_kv_tp_topology() + if self._from_tp == self._to_tp: + return [self._local_rank] + + if self._from_tp > self._to_tp: + tp_ratio = self._from_tp // self._to_tp + return [self._local_rank * tp_ratio + i for i in range(tp_ratio)] + else: + tp_ratio = self._to_tp // self._from_tp + return [self._local_rank // tp_ratio] + + def is_data_transfer_rank(self) -> bool: + """Whether this rank should participate in data (non-KV) transfer. + + Ordinary stage payloads are TP-identical, so exactly one TP rank + should talk to the connector. When TP is initialized, use TP rank 0 + so the connector leader matches TP-local broadcast source rank. + Otherwise fall back to LOCAL_RANK==0 for the single-rank case. + """ + tp_group = self._get_local_tp_group() + if tp_group is not None and getattr(tp_group, "world_size", 1) > 1: + return getattr(tp_group, "rank_in_group", 0) == 0 + return self._local_rank == 0 + + def get_kv_connector_key( + self, + req_id: str, + from_stage: int, + chunk_id: int, + from_rank: int, + to_rank: int, + ) -> str: + """Build connector key that includes rank info for KV transfers.""" + return f"{req_id}_{from_stage}_{chunk_id}_{from_rank}_{to_rank}" diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py index 1c8fe5b9f57..3f94c16a725 100644 --- a/vllm_omni/worker/omni_connector_model_runner_mixin.py +++ b/vllm_omni/worker/omni_connector_model_runner_mixin.py @@ -1,2341 +1,27 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unified data-plane communication mixin for Model Runners. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Public model-runner interface for Omni connector transport.""" -All connector.put()/get() calls are consolidated here. Background I/O -threads handle async_chunk and full_payload_mode transfers; KV cache is delegated to -the existing OmniKVTransferManager (to be absorbed later). - -The mixin reports transfer results via OmniConnectorOutput so that the -Scheduler can make scheduling decisions without ever touching a connector. -""" - -from __future__ import annotations - -import importlib -import inspect -import os -import threading -from collections import defaultdict, deque -from typing import TYPE_CHECKING, Any - -import torch -from vllm.distributed.parallel_state import get_tp_group -from vllm.logger import init_logger - -from vllm_omni.data_entry_keys import OmniPayload -from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory -from vllm_omni.distributed.omni_connectors.utils.config import ( - ConnectorSpec, - get_stage_connector_role, +from vllm_omni.distributed.omni_connectors.model_runner.omni_connector_payload_transport import ( + _OmniConnectorPayloadTransportMixin, ) -from vllm_omni.outputs import OmniConnectorOutput -from vllm_omni.worker.payload_span import ( - get_tensor_span, - merge_tensor_spans, +from vllm_omni.distributed.omni_connectors.model_runner.omni_connector_runtime import ( + _should_create_payload_connector, + needs_omni_connector, + should_accumulate_full_payload_output, ) -_EMBED_SPAN_GROUPS: tuple[tuple[str, str, str], ...] = (("decode", "decode_token_start", "decode_token_end"),) - -if TYPE_CHECKING: - from vllm_omni.config.model import OmniModelConfig - from vllm_omni.distributed.omni_connectors.connectors.base import ( - OmniConnectorBase, - ) - from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( - OmniKVTransferManager, - ) - -logger = init_logger(__name__) - - -def needs_omni_connector(model_config: Any) -> bool: - """Whether a runner owns an input, output, or explicitly routed connector.""" - return ( - bool(getattr(model_config, "requires_full_payload_input", False)) - or bool(getattr(model_config, "custom_process_next_stage_input_func", None)) - or get_stage_connector_role(model_config) is not None - ) - - -def _should_create_payload_connector(model_config: Any) -> bool: - """Whether this stage owns runner payload transport for its edge. - - Sender edges may instead be owned solely by KV transfer. Receivers still - need a connector even though they do not declare a downstream payload hook. - """ - if get_stage_connector_role(model_config) != "sender": - return True - - next_stage_func = getattr(model_config, "custom_process_next_stage_input_func", None) - return isinstance(next_stage_func, str) and bool(next_stage_func) - - -def should_accumulate_full_payload_output(model_config, custom_process_func) -> bool: - """Producer-side structural gate. - - Fires iff the stage explicitly declares a downstream full-payload - producer hook via ``custom_process_next_stage_input_func``. Consumer - stages may have ``custom_process_input_func`` values that can be - mechanically derived to ``*_full_payload`` helper names in the same - module; those are intentionally not enough to make the stage a producer. - """ - if custom_process_func is None: - return False - if getattr(model_config, "async_chunk", False): - return False - if getattr(model_config, "final_output", False): - return False - next_stage_func = getattr(model_config, "custom_process_next_stage_input_func", None) - if not isinstance(next_stage_func, str) or not next_stage_func: - return False - return getattr(model_config, "model_stage", None) is not None - +__all__ = [ + "OmniConnectorModelRunnerMixin", + "_should_create_payload_connector", + "needs_omni_connector", + "should_accumulate_full_payload_output", +] -class OmniConnectorModelRunnerMixin: - """Unified data-plane communication mixin for Model Runners. - Provides three transfer modes through a single pair of bg I/O threads: - - **full_payload_mode**: ``recv_full_payload_inputs`` / ``send_full_payload_outputs`` - - **Streaming (async_chunk)**: ``recv_chunk`` / ``send_chunk`` - - **KV cache**: ``send_kv_cache`` / ``recv_kv_cache`` (delegates to - the existing ``OmniKVTransferManager``) +class OmniConnectorModelRunnerMixin(_OmniConnectorPayloadTransportMixin): + """Unified data-plane communication interface for model runners. - The mixin owns connector instances and background threads. It never - touches scheduling queues -- readiness is communicated to the Scheduler - via ``OmniConnectorOutput``. + Runtime ownership and payload transport live under + ``distributed.omni_connectors``. """ - - # ------------------------------------------------------------------ # - # Init / Shutdown - # ------------------------------------------------------------------ # - - def init_omni_connectors( - self, - model_config: OmniModelConfig, - kv_transfer_manager: OmniKVTransferManager | None = None, - ) -> None: - """Initialize connectors and background threads. - - Args: - model_config: Stage-level model config with connector settings. - kv_transfer_manager: Existing KV transfer manager to delegate to. - """ - self._omni_connector: OmniConnectorBase | None = ( - self._create_connector(model_config) if _should_create_payload_connector(model_config) else None - ) - self._kv_transfer_manager = kv_transfer_manager - - self._async_chunk: bool = getattr(model_config, "async_chunk", False) - self._model_mode: str = getattr(model_config, "worker_type", "ar") - stage_id = getattr(model_config, "stage_id", 0) - if isinstance(stage_id, str): - stage_id = int(stage_id) - self._stage_id: int = stage_id if isinstance(stage_id, int) else 0 - - self._custom_process_func_path, self._custom_process_func = self._load_custom_func(model_config) - self._custom_process_supports_is_finished = self._custom_process_supports_is_finished_kwarg() - logger.debug( - "[Stage-%s] init_omni_connectors: async_chunk=%s, custom_process_func=%s, connector=%s, func_path=%s", - self._stage_id, - self._async_chunk, - self._custom_process_func, - type(self._omni_connector).__name__ if self._omni_connector else None, - self._custom_process_func_path, - ) - - # -- next stage ID (from connector config or default stage_id + 1) -- - self._next_stage_id: int = self._resolve_next_stage_id(model_config) - - # -- heterogeneous TP rank support -- - rank_cfg = self._parse_rank_mapping(model_config) - if self._kv_transfer_manager is not None: - topology = getattr(self._kv_transfer_manager, "tp_topology", None) - effective_mapping = ( - getattr(topology, "source_tp_size", None), - getattr(topology, "target_tp_size", None), - getattr(topology, "local_rank", None), - ) - if all(isinstance(value, int) for value in effective_mapping): - rank_cfg = { - "from_tp": effective_mapping[0], - "to_tp": effective_mapping[1], - "local_rank": effective_mapping[2], - } - self._from_tp: int = rank_cfg["from_tp"] - self._to_tp: int = rank_cfg["to_tp"] - self._local_rank: int = rank_cfg["local_rank"] - if self._kv_transfer_manager is not None: - self._kv_transfer_manager.kv_send_key_builder = self.get_rank_aware_kv_send_keys - self._kv_transfer_manager.kv_recv_key_builder = self.get_rank_aware_kv_keys - self._kv_transfer_manager.kv_payload_merger = self._merge_rank_sharded_kv_payloads - self._kv_transfer_manager.kv_payload_slicer = self._slice_rank_sharded_kv_payload - - # -- chunk index tracking (ported from OmniChunkTransferAdapter) -- - self._put_req_chunk: dict[str, int] = defaultdict(int) - self._get_req_chunk: dict[str, int] = defaultdict(int) - # Segment-local chunk counter: incremented alongside _put_req_chunk - # and popped at request cleanup. Note: the mixin path (uniproc mode) - # does not have segment boundary infrastructure; multi-segment support - # is only available via chunk_transfer_adapter (distributed path). - self._ramp_chunk_count: dict[str, int] = defaultdict(int) - self._adaptive_states: dict[str, Any] = {} - # Send-side async accumulation / staging buffer. Receive-side payload - # ownership lives in ``_local_stage_payload_cache``. - self._send_side_request_payload: dict[str, dict[str, Any]] = {} - self._code_prompt_token_ids: dict[str, list[list[int]]] = defaultdict(list) - self._cached_ic: dict[str, int] = {} - self._request_ids_mapping: dict[str, str] = {} - - # -- async I/O state (shared by chunk + full_payload_mode) -- - self._pending_load_reqs: dict[str, Any] = {} - self._finished_load_reqs: set[str] = set() - self._pending_save_reqs: dict[str, deque] = {} - self._pending_save_counts: dict[str, int] = defaultdict(int) - self._deferred_send_cleanup: set[str] = set() - # -- per-cycle output accumulator -- - self._chunk_ready_req_ids: set[str] = set() - self._chunk_finished_req_ids: set[str] = set() - self._stage_recv_req_ids: set[str] = set() - self._full_payload_pending_broadcast_req_ids: set[str] = set() - self._async_chunk_updated_req_ids: set[str] = set() - - # -- Model Runner local payload cache (RFC §2.4) -- - # Full stage payloads land here first on the recv side. We - # intentionally do not write connector recv results straight into - # `model_intermediate_buffer`: runner-owned runtime state is - # materialized later by `_sync_local_stage_payloads()` on the - # model thread. This keeps recv timing separate from execute-step - # visibility and avoids mixing connector I/O with model runtime - # ownership. - self._local_stage_payload_cache: dict[str, dict[str, Any]] = {} - # Lightweight scheduling metadata pending delivery to the Scheduler. - self._local_request_metadata: dict[str, dict[str, Any]] = {} - - # -- persistent set of request IDs whose chunk stream is complete -- - # Prevents re-registration after the finish sentinel has been received. - self._chunk_stream_completed: set[str] = set() - - # -- full_payload_mode: accumulate latest pooler_output per request, - # send only when the request finishes (next-cycle flush) -- - self._pending_full_payload_send: dict[str, tuple[Any, ...]] = {} - - # -- KV sent accumulator -- - self._kv_sent_req_ids: list[str] = [] - - # -- KV transfer lifecycle (absorbed from scheduler) -- - # Requests marked for KV transfer: {req_id: {seq_len, block_ids}} - self._kv_pending_transfers: dict[str, dict[str, Any]] = {} - # Requests whose KV transfer has been submitted but not yet acked - self._kv_active_transfers: set[str] = set() - # Requests whose KV transfer is complete (acked by kv_extracted_req_ids) - self._kv_completed_transfers: set[str] = set() - # Dedup guard: requests that have already triggered KV transfer - self._kv_triggered_requests: set[str] = set() - - self._lock = threading.Lock() - self._stop_event = threading.Event() - self._work_available = threading.Event() - - # Start background threads only when there's a connector - self._recv_thread: threading.Thread | None = None - self._save_thread: threading.Thread | None = None - if self._omni_connector is not None: - self._recv_thread = threading.Thread( - target=self._recv_loop, - daemon=True, - name="omni-mixin-recv", - ) - self._recv_thread.start() - self._save_thread = threading.Thread( - target=self._save_loop, - daemon=True, - name="omni-mixin-save", - ) - self._save_thread.start() - - # Explicit "fully initialised" marker so other parts of the runner - # (e.g. _update_states cleanup) can branch on a stable contract - # instead of probing for private mixin attribute names. Must be set - # only after every field above has been bound, so a partially - # constructed mixin is never observable as initialised. - self._omni_connector_initialized = True - - def shutdown_omni_connectors(self) -> None: - """Stop background threads and release connector resources.""" - self._stop_event.set() - if self._recv_thread is not None: - self._recv_thread.join(timeout=5) - if self._save_thread is not None: - self._save_thread.join(timeout=5) - if self._omni_connector is not None: - try: - self._omni_connector.close() - except Exception: - pass - - def cleanup_finished_request(self, req_id: str) -> None: - """Clean up per-request state after a request is fully finished. - - Call this when a request is freed from the model runner to prevent - memory leaks in the mixin's tracking dicts/sets. - - Two senders use different keys: ``send_chunk`` keys per-request - state under the EXTERNAL id (after mapping resolution), while - ``send_full_payload_outputs`` keys under the INTERNAL id. To cover - both modes (and forward compat with id-rename scenarios) we attempt - cleanup against both keys; the entry that doesn't exist for the - active mode is a no-op pop. Only the key that actually has pending - saves is added to ``_deferred_send_cleanup`` so the bg save's - decrement path drains it without leaving orphans. - """ - # Force-flush any pending full-payload accumulator entry before - # cleanup proceeds. Without this, finished requests with no - # downstream consumer (e.g. text-only on multi-modal arch) leave - # the entry orphaned in _pending_full_payload_send across requests, - # which empirically destabilises subsequent thinker forwards by - # making prefix-cache reuse observe stale accumulator state. The - # flush is idempotent when the entry has already been flushed by the - # scheduler-driven path, but this cleanup path runs for every request, - # so skip it entirely when the request never accumulated a payload. - if req_id in self._pending_full_payload_send: - try: - self.flush_full_payload_outputs({req_id}) - except Exception: - # Cleanup must still proceed regardless of flush errors here -- - # we already gated on ``_omni_connector_initialized`` upstream, - # so any exception here reflects a real connector-side issue - # (shared memory corruption, background thread crash) worth - # surfacing rather than silently swallowing. - logger.warning( - "flush_full_payload_outputs(%s) raised during cleanup; continuing tear-down.", - req_id, - exc_info=True, - ) - - ext_id = self._request_ids_mapping.pop(req_id, None) - keys_to_clean: list[str] = [req_id] - if ext_id is not None and ext_id != req_id: - keys_to_clean.append(ext_id) - - with self._lock: - keys_pending = [k for k in keys_to_clean if self._pending_save_counts.get(k, 0)] - for k in keys_pending: - self._deferred_send_cleanup.add(k) - for k in keys_to_clean: - if k in keys_pending: - continue - self._put_req_chunk.pop(k, None) - self._send_side_request_payload.pop(k, None) - self._code_prompt_token_ids.pop(k, None) - self._cached_ic.pop(k, None) - self._ramp_chunk_count.pop(k, None) - self._adaptive_states.pop(k, None) - self._kv_pending_transfers.pop(req_id, None) - self._kv_active_transfers.discard(req_id) - self._kv_completed_transfers.discard(req_id) - self._kv_triggered_requests.discard(req_id) - self._cleanup_recv_delivery_state(req_id) - - def drop_inactive_request_delivery_state(self, req_id: str) -> None: - """Clear recv-side state for inactive requests.""" - ext_id = self._request_ids_mapping.pop(req_id, None) - if hasattr(self, "_lock"): - with self._lock: - self._drop_send_side_payload_state(req_id, ext_id) - else: - self._drop_send_side_payload_state(req_id, ext_id) - self._cleanup_recv_delivery_state(req_id) - - def _drop_send_side_payload_state(self, req_id: str, ext_id: str | None) -> None: - if ext_id is not None: - self._send_side_request_payload.pop(ext_id, None) - self._cached_ic.pop(ext_id, None) - self._send_side_request_payload.pop(req_id, None) - self._cached_ic.pop(req_id, None) - - def _cleanup_recv_delivery_state(self, req_id: str) -> None: - """Clear recv-side delivery-cycle state.""" - if hasattr(self, "_lock"): - with self._lock: - self._clear_recv_delivery_state(req_id) - else: - self._clear_recv_delivery_state(req_id) - - def _clear_recv_delivery_state(self, req_id: str) -> None: - self._get_req_chunk.pop(req_id, None) - self._pending_load_reqs.pop(req_id, None) - self._finished_load_reqs.discard(req_id) - self._chunk_ready_req_ids.discard(req_id) - self._chunk_finished_req_ids.discard(req_id) - self._chunk_stream_completed.discard(req_id) - self._stage_recv_req_ids.discard(req_id) - self._full_payload_pending_broadcast_req_ids.discard(req_id) - self._async_chunk_updated_req_ids.discard(req_id) - self._local_stage_payload_cache.pop(req_id, None) - self._local_request_metadata.pop(req_id, None) - - def prune_inactive_requests(self, active_req_ids: Any) -> set[str]: - """Drop connector state for requests that no longer exist locally. - - Preempted / unscheduled requests are expected to stay in - ``self.requests`` and therefore remain untouched. This only prunes - stale request IDs that have already fallen out of the active request - map, preventing background recv/send bookkeeping from outliving the - request lifecycle. - """ - if active_req_ids is None: - return set() - - active_req_ids = set(active_req_ids) - pending_req_ids = set(getattr(self, "_pending_load_reqs", {}).keys()) - received_req_ids = set(getattr(self, "_stage_recv_req_ids", set())) - received_req_ids.update(getattr(self, "_full_payload_pending_broadcast_req_ids", set())) - received_req_ids.update(getattr(self, "_local_request_metadata", {}).keys()) - # Pending recv requests may not yet be in the caller's active set - # (e.g. WAITING_FOR_CHUNK requests live in the coordinator's internal - # queues, not in model runner self.requests). Protect them so that - # legitimate waiting requests are not pruned. - # - # Likewise, a full payload can arrive on the background recv thread - # after the scheduler_output snapshot for the current execute_model() - # cycle was already materialized. Those requests may briefly live only - # in recv-side buffers/local cache until the next scheduler cycle wakes - # them up; pruning them here drops the payload before stage_recv can be - # published. - active_req_ids.update(pending_req_ids) - active_req_ids.update(received_req_ids) - stale_req_ids: set[str] = set() - - # NOTE: _pending_load_reqs is excluded from the scan list because - # all its entries are unconditionally protected above. The mixin - # cannot distinguish a legitimately-waiting pending recv from an - # orphaned one (only the coordinator/scheduler knows). - # - # Requests with freshly received full payloads / local stage payloads - # are also protected above. Their scheduler wake-up may lag the recv - # thread by one execute_model() cycle, especially when the request was - # added after the current scheduler_output snapshot. - # - # Orphaned pending recv entries (e.g. from upstream stage crash) are - # handled by collect_timed_out_request_ids() -- on - # OmniSchedulingCoordinator for full-payload requests, and on - # OmniChunkTransferAdapter for async-chunk ones -- which detect - # wait-time violations. The scheduler then removes the request from - # its queues, sets FINISHED_ERROR, and calls _free_request() which - # ultimately triggers cleanup_finished_request() here. - for attr_name in ( - "_request_ids_mapping", - "_get_req_chunk", - "_finished_load_reqs", - "_chunk_ready_req_ids", - "_chunk_finished_req_ids", - "_chunk_stream_completed", - "_stage_recv_req_ids", - "_full_payload_pending_broadcast_req_ids", - "_async_chunk_updated_req_ids", - "_local_stage_payload_cache", - "_local_request_metadata", - "_kv_pending_transfers", - "_kv_active_transfers", - "_kv_completed_transfers", - "_kv_triggered_requests", - ): - state = getattr(self, attr_name, None) - if isinstance(state, dict): - stale_req_ids.update(req_id for req_id in state if req_id not in active_req_ids) - elif isinstance(state, set): - stale_req_ids.update(req_id for req_id in state if req_id not in active_req_ids) - - for req_id in stale_req_ids: - self.cleanup_finished_request(req_id) - - return stale_req_ids - - # ------------------------------------------------------------------ # - # Local payload cache (RFC §2.4 – Model Runner ownership) - # ------------------------------------------------------------------ # - - def put_local_stage_payload(self, req_id: str, payload: OmniPayload) -> None: - """Store a full stage payload in the local cache.""" - self._local_stage_payload_cache[req_id] = payload - - def get_local_stage_payload(self, req_id: str) -> OmniPayload | None: - """Read a stage payload without removing it.""" - return self._local_stage_payload_cache.get(req_id) - - def pop_local_stage_payload(self, req_id: str) -> OmniPayload | None: - """Remove and return a stage payload (consume after use).""" - return self._local_stage_payload_cache.pop(req_id, None) - - def put_local_request_metadata(self, req_id: str, metadata: dict[str, Any]) -> None: - """Store lightweight scheduling metadata for a request.""" - self._local_request_metadata[req_id] = metadata - - def get_local_request_metadata(self, req_id: str) -> dict[str, Any] | None: - """Retrieve scheduling metadata for a request.""" - return self._local_request_metadata.get(req_id) - - # ------------------------------------------------------------------ # - # Scheduling metadata extraction - # ------------------------------------------------------------------ # - - @classmethod - def _extract_scheduling_metadata(cls, payload: OmniPayload) -> dict[str, Any]: - """Extract only the fields the scheduler needs from a full payload.""" - extracted: dict[str, Any] = {} - meta = payload.get("meta") if isinstance(payload, dict) else None - meta = meta if isinstance(meta, dict) else {} - - if "next_stage_prompt_len" in meta: - extracted["next_stage_prompt_len"] = meta["next_stage_prompt_len"] - elif "next_stage_prompt_len" in payload: - logger.warning_once( - "legacy flat 'next_stage_prompt_len' key in payload; expected 'meta.next_stage_prompt_len'" - ) - extracted["next_stage_prompt_len"] = payload["next_stage_prompt_len"] - - audio_codes = cls._payload_audio_codes(payload) - if audio_codes is not None: - extracted["code_predictor_codes"] = audio_codes - - if "left_context_size" in meta: - extracted["left_context_size"] = meta["left_context_size"] - elif "left_context_size" in payload: - logger.warning_once("legacy flat 'left_context_size' key in payload; expected 'meta.left_context_size'") - - return extracted - - _NON_CONSUMABLE_PAYLOAD_KEYS: set[tuple[str, str]] = { - ("meta", "finished"), - ("meta", "override_keys"), - ("meta", "next_stage_prompt_len"), - ("meta", "left_context_size"), - ("ids", "output"), - ("embed", "decode_token_start"), - ("embed", "decode_token_end"), - } - - @staticmethod - def _payload_value_has_content(value: Any) -> bool: - if value is None: - return False - if isinstance(value, torch.Tensor): - return value.numel() > 0 - if isinstance(value, (list, tuple, dict, set)): - return len(value) > 0 - return True - - @staticmethod - def _payload_finished(payload: Any) -> bool: - if not isinstance(payload, dict): - return False - if "finished" in payload: - logger.warning_once("legacy flat 'finished' key in payload; expected 'meta.finished'") - meta = payload.get("meta") - if not isinstance(meta, dict) or "finished" not in meta: - return False - flag = meta["finished"] - if isinstance(flag, torch.Tensor): - return flag.numel() == 1 and bool(flag.item()) - return bool(flag) - - @staticmethod - def _payload_audio_codes(payload: Any) -> Any: - if not isinstance(payload, dict): - return None - if "code_predictor_codes" in payload: - logger.warning_once("legacy flat 'code_predictor_codes' key in payload; expected 'codes.audio'") - codes = payload.get("codes") - if isinstance(codes, dict): - return codes.get("audio") - return None - - @classmethod - def _payload_is_consumable(cls, payload: OmniPayload | None) -> bool: - """Return True when an async payload can drive a real forward step. - - Metadata-only wake-ups should not transition WAITING_FOR_CHUNK requests - back to schedulable state. In particular, a widened token horizon without - any newly visible thinker decode embeds should not force a placeholder-only - talker decode step. - """ - if not isinstance(payload, dict) or not payload: - return False - - embed = payload.get("embed") - if isinstance(embed, dict): - decode_embeddings = embed.get("decode") - if isinstance(decode_embeddings, torch.Tensor): - if decode_embeddings.ndim == 0: - return True - return decode_embeddings.numel() > 0 and decode_embeddings.shape[0] > 0 - - audio_codes = cls._payload_audio_codes(payload) - if audio_codes is not None: - if isinstance(audio_codes, torch.Tensor): - return audio_codes.numel() > 0 - if hasattr(audio_codes, "__len__"): - return len(audio_codes) > 0 - return True - - for key, value in payload.items(): - if isinstance(value, dict): - for sk, sv in value.items(): - if (key, sk) in cls._NON_CONSUMABLE_PAYLOAD_KEYS: - continue - if cls._payload_value_has_content(sv): - return True - continue - if cls._payload_value_has_content(value): - return True - return False - - @staticmethod - def _get_local_tp_group() -> Any | None: - """Return the local TP group when tensor parallelism is initialized.""" - try: - return get_tp_group() - except Exception: - return None - - def _recv_ordinary_stage_result( - self, - connector: OmniConnectorBase, - from_stage: str, - to_stage: str, - connector_get_key: str, - ) -> Any: - """Receive one ordinary non-KV stage payload on the local leader rank only.""" - tp_group = self._get_local_tp_group() - if tp_group is None or getattr(tp_group, "world_size", 1) <= 1: - return connector.get(from_stage, to_stage, connector_get_key) - if not self.is_data_transfer_rank(): - return None - return connector.get(from_stage, to_stage, connector_get_key) - - def _recv_full_payload_result( - self, - connector: OmniConnectorBase, - from_stage: str, - to_stage: str, - connector_get_key: str, - ) -> Any: - """Receive one full-payload transfer on the local leader rank only.""" - return self._recv_ordinary_stage_result( - connector, - from_stage, - to_stage, - connector_get_key, - ) - - def _recv_async_chunk_result( - self, - connector: OmniConnectorBase, - from_stage: str, - to_stage: str, - connector_get_key: str, - ) -> Any: - """Receive one ordinary async chunk on the local leader rank only.""" - return self._recv_ordinary_stage_result( - connector, - from_stage, - to_stage, - connector_get_key, - ) - - @staticmethod - def _snapshot_payload(payload: Any) -> Any: - if isinstance(payload, dict): - return dict(payload) - return payload - - def _broadcast_tp_payload_packet(self, packet: Any) -> Any: - """Broadcast one ordinary payload packet from TP rank 0 when TP is active.""" - tp_group = self._get_local_tp_group() - if tp_group is None or getattr(tp_group, "world_size", 1) <= 1: - return packet - leader_packet = packet if self.is_data_transfer_rank() else None - return tp_group.broadcast_object(leader_packet, src=0) - - def _apply_staged_payloads_locked(self, staged_payloads: dict[str, Any]) -> None: - for req_id, payload in staged_payloads.items(): - self._local_stage_payload_cache[req_id] = self._snapshot_payload(payload) - - def _collect_full_payload_results_locked(self) -> dict[str, Any] | None: - if not self._full_payload_pending_broadcast_req_ids: - return None - results: dict[str, Any] = {} - missing_req_ids: list[str] = [] - for req_id in tuple(self._full_payload_pending_broadcast_req_ids): - payload = self._local_stage_payload_cache.get(req_id) - if payload is None: - missing_req_ids.append(req_id) - continue - results[req_id] = self._snapshot_payload(payload) - self._full_payload_pending_broadcast_req_ids.discard(req_id) - if missing_req_ids: - logger.warning( - "[Stage-%s] _collect_full_payload_results_locked: " - "pending full-payload reqs missing from local cache: %s", - self._stage_id, - missing_req_ids, - ) - return results or None - - def _collect_async_chunk_fanout_packet_locked(self) -> dict[str, Any] | None: - payload_req_ids = set(self._async_chunk_updated_req_ids) - payload_req_ids.update(self._finished_load_reqs) - payload_req_ids.update(self._chunk_finished_req_ids) - payload_req_ids.update(self._local_request_metadata) - if not ( - payload_req_ids or self._finished_load_reqs or self._chunk_finished_req_ids or self._local_request_metadata - ): - return None - - staged_payloads = { - req_id: self._snapshot_payload(self._local_stage_payload_cache[req_id]) - for req_id in payload_req_ids - if req_id in self._local_stage_payload_cache - } - packet = { - "staged_payloads": staged_payloads, - "request_metadata": dict(self._local_request_metadata), - "newly_finished": set(self._finished_load_reqs), - "chunk_finished": set(self._chunk_finished_req_ids), - } - - self._async_chunk_updated_req_ids.clear() - self._finished_load_reqs.clear() - self._chunk_finished_req_ids.clear() - self._local_request_metadata.clear() - - for req_id in packet["chunk_finished"]: - if req_id not in self._local_stage_payload_cache: - continue - ext_req_id = self._request_ids_mapping.get(req_id, req_id) - self._send_side_request_payload.pop(ext_req_id, None) - if ext_req_id != req_id: - self._send_side_request_payload.pop(req_id, None) - - return packet - - def _apply_async_chunk_fanout_packet(self, packet: dict[str, Any]) -> None: - staged_payloads = packet.get("staged_payloads", {}) - chunk_finished = set(packet.get("chunk_finished", ())) - with self._lock: - self._apply_staged_payloads_locked(staged_payloads) - for req_id in chunk_finished: - self._pending_load_reqs.pop(req_id, None) - self._chunk_stream_completed.add(req_id) - - # ------------------------------------------------------------------ # - # full_payload_mode (recv_full_payload_inputs / send_full_payload_outputs) - # ------------------------------------------------------------------ # - - def recv_full_payload_inputs(self, scheduler_output: Any) -> dict[str, Any] | None: - """Check for incoming full_payload_mode stage inputs (non-blocking). - - Returns a dict mapping ``request_id -> engine_inputs`` for data - that has arrived, or ``None`` if nothing is ready. Stores full - payloads in the local cache and extracts scheduling metadata. - """ - # Fast path: when TP is trivial (no peer ranks waiting on a broadcast) - # and the bg recv thread has not staged anything, skip the lock + TP - # broadcast cycle entirely. _broadcast_tp_payload_packet already - # returns its input unchanged under the same world_size<=1 condition, - # so the original code path was a no-op here on every empty step. - tp_group = self._get_local_tp_group() - if ( - tp_group is None or getattr(tp_group, "world_size", 1) <= 1 - ) and not self._full_payload_pending_broadcast_req_ids: - return None - with self._lock: - results = self._collect_full_payload_results_locked() if self.is_data_transfer_rank() else None - results = self._broadcast_tp_payload_packet(results) - if not results: - return None - with self._lock: - self._stage_recv_req_ids.update(results.keys()) - for req_id in results: - self._pending_load_reqs.pop(req_id, None) - self._apply_staged_payloads_locked(results) - for req_id, payload in results.items(): - self._local_request_metadata[req_id] = self._extract_scheduling_metadata(payload) - logger.debug( - "[Stage-%s] recv_full_payload_inputs: consumed %s reqs: %s, stage_recv_req_ids now=%s", - self._stage_id, - len(results), - list(results.keys()), - self._stage_recv_req_ids, - ) - return results - - def _get_model_config(self) -> Any: - model_config = getattr(self, "model_config", None) - if model_config is not None: - return model_config - return getattr(getattr(self, "vllm_config", None), "model_config", None) - - def _should_accumulate_full_payload_output(self) -> bool: - """Gate send-side full-payload output accumulation only. - - Cached per instance: the result depends only on model_config / - _custom_process_func, both of which are set at init time. Avoid - the per-step dynamic import inside the model decode loop. - """ - if getattr(self, "_omni_connector", None) is None: - # No connector at all: send_full_payload_outputs would no-op. - # Skip the per-step accumulator+build that would otherwise be - # silently discarded. Defends against a terminal stage whose - # custom_process_input_func has a *_full_payload derivative in - # the same module (e.g. dynin stage 2 token2image_to_token2audio - # in pipelines that don't configure any connector at all). - # - # Known limitation: a *terminal-consumer* stage that has a - # connector configured for receiving upstream input is NOT - # caught here -- ``_omni_connector`` is non-None for it, and - # ``_load_custom_func`` may still resolve a ``*_full_payload`` - # derivative from this stage's ``custom_process_input_func``. - # In that case the accumulator builds payloads that - # ``send_full_payload_outputs`` later drops via its own - # connector-side checks (wasted CPU, not a functional bug). - # A topology-aware gate (explicit producer field or pipeline - # is_terminal info) would close the gap; that change is out - # of scope for this PR. - self._should_accumulate_full_payload_output_cached = False - return False - cached = getattr(self, "_should_accumulate_full_payload_output_cached", None) - if cached is not None: - return cached - model_config = self._get_model_config() - if model_config is None: - self._should_accumulate_full_payload_output_cached = False - return False - result = should_accumulate_full_payload_output( - model_config, - getattr(self, "_custom_process_func", None), - ) - self._should_accumulate_full_payload_output_cached = result - return result - - @staticmethod - def _new_full_payload_accumulator(output: dict[str, Any]): - chunks: dict[str, list[torch.Tensor]] = {} - latest: dict[str, Any] = {} - rows: dict[str, int] = {} - for k, v in output.items(): - if isinstance(v, torch.Tensor) and v.dim() >= 2: - chunks[k] = [v] - rows[k] = int(v.shape[0]) - else: - latest[k] = v - return chunks, latest, rows - - @staticmethod - def _materialize_full_payload_entry(entry): - if len(entry) == 2: - return entry - chunks, latest, _rows, request = entry - output = dict(latest) - for k, tensors in chunks.items(): - if tensors: - output[k] = tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=0) - return output, request - - def _resolve_full_payload_replace_keys(self) -> frozenset: - """Per-model REPLACE-key set for the full-payload accumulator. - - Looked up from the stage-input-processor module that ships the model's sync builder - (`model_config.custom_process_input_func.__module__`). The module - declares ``_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str]``; if absent, - returns the empty set. - - Cached per instance. Keys in this set use REPLACE semantics in the - accumulator (subsequent emissions discard prior chunks) instead of - the default CONCAT semantics. Use for tensors that carry the full - result so far rather than per-step deltas (e.g. ``model_outputs``). - """ - cached = getattr(self, "_full_payload_replace_keys_cached", None) - if cached is not None: - return cached - proc = getattr(self, "_custom_process_func", None) - if proc is None: - self._full_payload_replace_keys_cached = frozenset() - return self._full_payload_replace_keys_cached - module_name = getattr(proc, "__module__", None) - if module_name is None: - self._full_payload_replace_keys_cached = frozenset() - return self._full_payload_replace_keys_cached - try: - import sys as _sys - - mod = _sys.modules.get(module_name) or importlib.import_module(module_name) - keys = getattr(mod, "_FULL_PAYLOAD_REPLACE_KEYS", frozenset()) - except ImportError: - logger.debug( - "Could not import stage input processor module %s while resolving " - "_FULL_PAYLOAD_REPLACE_KEYS; using CONCAT semantics for all keys.", - module_name, - exc_info=True, - ) - keys = frozenset() - if not isinstance(keys, (frozenset, set)): - logger.debug( - "Ignoring non-set _FULL_PAYLOAD_REPLACE_KEYS from %s: %s", - module_name, - type(keys).__name__, - ) - keys = frozenset() - self._full_payload_replace_keys_cached = frozenset(keys) - logger.debug( - "Resolved _FULL_PAYLOAD_REPLACE_KEYS for %s: %s", - module_name, - sorted(self._full_payload_replace_keys_cached), - ) - return self._full_payload_replace_keys_cached - - def accumulate_full_payload_output( - self, - req_id: str, - pooler_output: Any, - request: Any, - ) -> None: - """Accumulate pooler_output for a request across steps (full_payload_mode). - - Per-token tensors (2-D+, matching trailing dims) are concatenated - along dim-0. Scalar / global tensors (1-D or 0-D) are replaced - with the latest value. - - Note: codec rows are NOT filtered for zero placeholders here. The - downstream consumer ``_extract_qwen3_full_payload_codec_rows`` crops - codec rows using ``output_token_ids`` as the authoritative source, - which makes any sender-side zero filtering redundant. Skipping the - sender-side ``t.any()`` scan also avoids a per-tensor GPU->CPU device - sync that stalled the decode pipeline. - - The data is actually sent when ``flush_full_payload_outputs`` is called - with the finished request IDs from the next scheduler cycle. - """ - replace_keys = self._resolve_full_payload_replace_keys() - existing = self._pending_full_payload_send.get(req_id) - - if existing is None: - chunks, latest, rows = self._new_full_payload_accumulator(pooler_output) - self._pending_full_payload_send[req_id] = (chunks, latest, rows, request) - return - - if len(existing) == 2: - chunks, latest, rows = self._new_full_payload_accumulator(existing[0]) - else: - chunks, latest, rows, _ = existing - - for k, v in pooler_output.items(): - if v is None: - continue - if k in replace_keys: - # Explicit REPLACE semantics: the new value supersedes any - # prior chunks (e.g. `model_outputs` carries the full result - # so far, not an appendable per-step delta). - latest.pop(k, None) - if isinstance(v, torch.Tensor) and v.dim() >= 2: - chunks[k] = [v] - rows[k] = int(v.shape[0]) - else: - chunks.pop(k, None) - rows.pop(k, None) - latest[k] = v - continue - if isinstance(v, torch.Tensor) and v.dim() >= 2: - if k in chunks and chunks[k] and v.shape[1:] == chunks[k][0].shape[1:]: - chunks[k].append(v) - rows[k] += int(v.shape[0]) - else: - latest.pop(k, None) - chunks[k] = [v] - rows[k] = int(v.shape[0]) - else: - chunks.pop(k, None) - rows.pop(k, None) - latest[k] = v - - self._pending_full_payload_send[req_id] = (chunks, latest, rows, request) - - def flush_full_payload_outputs(self, finished_req_ids: set[str]) -> None: - """Send accumulated full_payload outputs for requests that just finished.""" - pending_req_ids = set(self._pending_full_payload_send.keys()) - if not (finished_req_ids & pending_req_ids): - return - - logger.debug( - "[Stage-%s] flush_full_payload_outputs: finished_req_ids=%s, pending=%s", - self._stage_id, - finished_req_ids, - list(self._pending_full_payload_send.keys()), - ) - to_send: dict[str, tuple[Any, Any]] = {} - for req_id in finished_req_ids: - entry = self._pending_full_payload_send.pop(req_id, None) - if entry is not None: - to_send[req_id] = self._materialize_full_payload_entry(entry) - logger.debug("[Stage-%s] flush_full_payload_outputs: to_send=%s", self._stage_id, list(to_send.keys())) - if to_send: - self.send_full_payload_outputs(scheduler_output=None, outputs=to_send) - - def send_full_payload_outputs( - self, - scheduler_output: Any, - outputs: dict[str, tuple[Any, Any] | Any], - ) -> list[str]: - """Send full_payload stage outputs to the next stage via connector. - - Args: - outputs: Mapping of ``req_id`` to either a - ``(pooling_output, request)`` tuple (preferred) or a raw - payload dict. When a tuple is supplied the request object - is forwarded to ``custom_process_stage_input_func``. - - Returns list of request IDs successfully enqueued. - """ - if self._omni_connector is None: - logger.debug("[Stage-%s] send_full_payload_outputs: connector is None, skip", self._stage_id) - return [] - if not self.is_data_transfer_rank(): - logger.debug( - "[Stage-%s] send_full_payload_outputs: not data_transfer_rank (rank=%s), skip", - self._stage_id, - self._local_rank, - ) - return list(outputs.keys()) - sent_ids: list[str] = [] - next_stage_id = self._next_stage_id - for req_id, value in outputs.items(): - if isinstance(value, tuple) and len(value) == 2: - raw_output, request = value - else: - raw_output, request = value, None - - payload = raw_output - if self._custom_process_func is not None: - payload = self._build_custom_process_payload( - request_id=req_id, - request=request, - pooling_output=raw_output, - ) - if payload is None: - continue - if payload is None: - logger.debug("[Stage-%s] send_full_payload_outputs: payload is None for %s", self._stage_id, req_id) - continue - if isinstance(payload, dict): - audio_codes = self._payload_audio_codes(payload) - if isinstance(audio_codes, torch.Tensor): - code_len = int(audio_codes.numel()) - elif hasattr(audio_codes, "__len__"): - code_len = len(audio_codes) - else: - code_len = None - meta = payload.get("meta") if isinstance(payload.get("meta"), dict) else {} - logger.debug( - "[Stage-%s] send_full_payload_outputs: req=%s payload_keys=%s code_len=%s left_context_size=%s", - self._stage_id, - req_id, - sorted(payload.keys()), - code_len, - meta.get("left_context_size"), - ) - - external_req_id = self._resolve_external_req_id(request, req_id) - chunk_id = self._put_req_chunk[req_id] - self._put_req_chunk[req_id] += 1 - self._ramp_chunk_count[req_id] += 1 - connector_put_key = f"{external_req_id}_{self._stage_id}_{chunk_id}" - - logger.debug( - "[Stage-%s] send_full_payload_outputs: enqueue req=%s put_key=%s next_stage=%s", - self._stage_id, - req_id, - connector_put_key, - next_stage_id, - ) - task = { - "stage_id": self._stage_id, - "next_stage_id": next_stage_id, - "put_key": connector_put_key, - "data": payload, - "request_id": req_id, - } - with self._lock: - self._pending_save_reqs.setdefault(req_id, deque()).append(task) - self._pending_save_counts[req_id] += 1 - sent_ids.append(req_id) - if sent_ids: - self._work_available.set() - return sent_ids - - # ------------------------------------------------------------------ # - # Streaming chunk mode (recv_chunk / send_chunk) - # ------------------------------------------------------------------ # - - def register_chunk_recv(self, request: Any) -> None: - """Register a request for async chunk retrieval by the bg thread. - - Stage-0 has no upstream producer so this is a no-op there. - Skips requests whose batch data has already been received to - prevent the bg thread from polling for non-existent chunks. - """ - if self._stage_id == 0: - return - request_id = request.request_id - # Explicit external_req_id=None must fall back to request_id; - # otherwise recv keys become `None__` and collide - # across requests. - ext = getattr(request, "external_req_id", None) - self._request_ids_mapping[request_id] = ext if ext is not None else request_id - with self._lock: - if request_id in self._stage_recv_req_ids: - return - # Don't re-register if the finish sentinel was already received - if request_id in self._chunk_stream_completed: - return - self._pending_load_reqs[request_id] = request - self._work_available.set() - - def recv_chunk(self) -> dict[str, Any]: - """Collect chunks received by the bg thread since last call. - - Returns a dict ``{request_id: chunk_payload}`` for newly arrived - chunks. Empty dict when nothing is ready. - - This method reads from ``_finished_load_reqs`` without clearing - it -- ``get_omni_connector_output()`` is the sole consumer that - drains and resets ``_finished_load_reqs`` at the end of each - ``execute_model`` cycle. - - Returns **shallow copies** of the cached payloads so that the - caller can read them without racing against the background recv - thread, which may concurrently mutate the live cache entries via - ``dict.update()``. - """ - with self._lock: - finished = set(self._finished_load_reqs) - if not finished: - return {} - # Snapshot the payloads under the lock to avoid racing with - # _poll_single_request which does existing.update(payload_data) - # on the same dict objects. - result = {} - for rid in finished: - payload = self._local_stage_payload_cache.get(rid) - result[rid] = dict(payload) if isinstance(payload, dict) else payload - - self._chunk_ready_req_ids.update(finished) - return result - - def send_chunk( - self, - request: Any, - pooling_output: Any | None = None, - ) -> bool: - """Derive and enqueue one chunk for async sending. - - Payload extraction runs in the caller thread (via - ``custom_process_stage_input_func``); the actual - ``connector.put()`` is done by the background save thread. - Non-KV data is identical across TP ranks; only rank 0 sends. - """ - if self._omni_connector is None: - logger.warning("[Stage-%s] send_chunk: connector is None", self._stage_id) - return False - if not self.is_data_transfer_rank(): - return True - raw_req_id = getattr(request, "request_id", None) or getattr(request, "req_id", None) - request_id = self._resolve_external_req_id(request, raw_req_id) - # Cache the internal→external mapping so that finish sentinels can - # resolve the external ID even after the request is freed. - if raw_req_id and raw_req_id != request_id: - self._request_ids_mapping.setdefault(raw_req_id, request_id) - chunk_id = self._put_req_chunk[request_id] - - payload_data = self._build_custom_process_payload( - request_id=request_id, - request=request, - pooling_output=pooling_output, - ) - if payload_data is None: - if chunk_id == 0: - logger.warning( - "[Stage-%s] send_chunk: payload is None for req=%s chunk=%s (process_func=%s)", - self._stage_id, - request_id, - chunk_id, - self._custom_process_func, - ) - return False - - self._put_req_chunk[request_id] += 1 - self._ramp_chunk_count[request_id] += 1 - next_stage_id = self._next_stage_id - connector_put_key = f"{request_id}_{self._stage_id}_{chunk_id}" - - if chunk_id == 0: - logger.debug( - "[Stage-%s] send_chunk: first chunk enqueued, req=%s key=%s", - self._stage_id, - request_id, - connector_put_key, - ) - - task = { - "stage_id": self._stage_id, - "next_stage_id": next_stage_id, - "put_key": connector_put_key, - "data": payload_data, - "request_id": request_id, - } - with self._lock: - self._pending_save_reqs.setdefault(request_id, deque()).append(task) - self._pending_save_counts[request_id] += 1 - self._work_available.set() - return True - - # ------------------------------------------------------------------ # - # KV cache (delegates to OmniKVTransferManager) - # ------------------------------------------------------------------ # - - def send_kv_cache( - self, - finished_reqs: dict[str, dict[str, Any]], - kv_caches: list[torch.Tensor], - block_size: int, - cache_dtype: str, - request_id_resolver: Any | None = None, - ) -> list[str]: - """Send KV cache for finished requests. - - Delegates to the existing ``OmniKVTransferManager``. - """ - if self._kv_transfer_manager is None: - return list(finished_reqs.keys()) if finished_reqs else [] - result = self._kv_transfer_manager.handle_finished_requests_kv_transfer( - finished_reqs=finished_reqs, - kv_caches=kv_caches, - block_size=block_size, - cache_dtype=cache_dtype, - request_id_resolver=request_id_resolver, - ) - if result: - self._kv_sent_req_ids.extend(result) - return result - - def recv_kv_cache( - self, - request_id: str, - target_device: torch.device | None = None, - ) -> tuple[dict[str, Any] | None, int]: - """Receive KV cache for a request. - - Delegates to the existing ``OmniKVTransferManager``. - """ - if self._kv_transfer_manager is None: - return None, 0 - return self._kv_transfer_manager.receive_kv_cache_for_request( - request_id=request_id, - target_device=target_device, - ) - - def receive_cfg_companion_kv_payloads( - self, - cfg_request_ids: dict[str, str], - target_device: torch.device | None = None, - ) -> dict[str, tuple[dict[str, Any] | None, int]]: - """Receive raw CFG companion KV payloads keyed by role.""" - return { - role: self.recv_kv_cache(companion_rid, target_device=target_device) - for role, companion_rid in cfg_request_ids.items() - } - - def receive_multi_kv_cache( - self, - req: Any, - cfg_kv_collect_func: Any | None = None, - target_device: torch.device | None = None, - ) -> bool: - """Receive primary and optional companion KV caches for a request. - - The mixin owns the runner-facing orchestration: primary KV receive, - companion payload fetch, and applying any model-specific CFG fields back - onto ``req.sampling_params``. - """ - if self._kv_transfer_manager is None: - return False - - request_id = getattr(req, "request_id", None) - if not request_id: - logger.warning("Request has no ID, cannot receive KV cache") - return False - - active_requests = getattr(self, "requests", None) - if active_requests is not None and request_id not in active_requests: - logger.debug("Skip receiving KV cache for inactive request %s", request_id) - return False - - primary_ok = False - data, _size = self.recv_kv_cache(request_id, target_device=target_device) - if data: - self._kv_transfer_manager.apply_kv_cache_to_request(req, data) - primary_ok = True - - cfg_ids = getattr(getattr(req, "sampling_params", None), "cfg_kv_request_ids", None) - if cfg_ids and cfg_kv_collect_func: - try: - cfg_role_payloads = self.receive_cfg_companion_kv_payloads( - cfg_ids, - target_device=target_device, - ) - cfg_kvs = cfg_kv_collect_func(request_id, cfg_role_payloads) - if cfg_kvs and hasattr(req, "sampling_params") and req.sampling_params is not None: - for key, value in cfg_kvs.items(): - setattr(req.sampling_params, key, value) - logger.debug("Applied CFG KV caches: %s", list(cfg_kvs.keys())) - except Exception: - logger.exception("Failed to collect CFG KV caches for %s", request_id) - - return primary_ok - - # ------------------------------------------------------------------ # - # Rank-aware KV transfer routing - # ------------------------------------------------------------------ # - - def get_rank_aware_kv_keys( - self, - req_id: str, - from_stage: int, - to_stage: int | None = None, - chunk_id: int = 0, - ) -> list[str]: - """Build recv-side connector keys for all remote ranks this rank needs. - - For heterogeneous TP receive, the local rank is the target rank and must - fetch one or more source-rank shards keyed as ``from_rank -> to_rank``. - """ - if self._from_tp <= 1 and self._to_tp <= 1: - resolved_to_stage = self._next_stage_id if to_stage is None else to_stage - return [f"omni_{from_stage}_to_{resolved_to_stage}_kv_cache_{req_id}"] - - remote_ranks = self.get_kv_remote_ranks() - return [ - self.get_kv_connector_key( - req_id=req_id, - from_stage=from_stage, - chunk_id=chunk_id, - from_rank=remote_rank, - to_rank=self._local_rank, - ) - for remote_rank in remote_ranks - ] - - def get_kv_target_ranks_for_send(self) -> list[int]: - """Determine which target ranks this local rank should send KV shards to.""" - self._validate_kv_tp_topology() - if self._from_tp == self._to_tp: - return [self._local_rank] - if self._from_tp > self._to_tp: - tp_ratio = self._from_tp // self._to_tp - return [self._local_rank // tp_ratio] - tp_ratio = self._to_tp // self._from_tp - base_rank = self._local_rank * tp_ratio - return [base_rank + i for i in range(tp_ratio)] - - def get_rank_aware_kv_send_keys( - self, - req_id: str, - from_stage: int, - to_stage: int | None = None, - chunk_id: int = 0, - ) -> list[str]: - """Build send-side connector keys for this rank's KV shard(s).""" - if self._from_tp <= 1 and self._to_tp <= 1: - resolved_to_stage = self._next_stage_id if to_stage is None else to_stage - return [f"omni_{from_stage}_to_{resolved_to_stage}_kv_cache_{req_id}"] - - target_ranks = self.get_kv_target_ranks_for_send() - return [ - self.get_kv_connector_key( - req_id=req_id, - from_stage=from_stage, - chunk_id=chunk_id, - from_rank=self._local_rank, - to_rank=target_rank, - ) - for target_rank in target_ranks - ] - - @staticmethod - def _merge_rank_sharded_kv_payloads(payloads: list[dict[str, Any]]) -> dict[str, Any] | None: - """Merge multiple source-rank KV shards for one target rank.""" - payloads = [payload for payload in payloads if isinstance(payload, dict)] - if not payloads: - return None - if len(payloads) == 1: - return payloads[0] - - merged = dict(payloads[0]) - layer_blocks = merged.get("layer_blocks") - if not isinstance(layer_blocks, dict): - return merged - - def _merge_tensor_lists(name: str) -> list[torch.Tensor | None]: - merged_list: list[torch.Tensor | None] = [] - cache_lists = [payload.get("layer_blocks", {}).get(name, []) for payload in payloads] - max_len = max((len(cache_list) for cache_list in cache_lists), default=0) - for idx in range(max_len): - tensors = [cache_list[idx] for cache_list in cache_lists if idx < len(cache_list)] - tensors = [tensor for tensor in tensors if isinstance(tensor, torch.Tensor)] - if not tensors: - merged_list.append(None) - elif len(tensors) == 1: - merged_list.append(tensors[0]) - else: - merged_list.append(torch.cat(tensors, dim=-2).contiguous()) - return merged_list - - merged["layer_blocks"] = { - "key_cache": _merge_tensor_lists("key_cache"), - "value_cache": _merge_tensor_lists("value_cache"), - } - metadata = dict(merged.get("metadata", {})) - metadata["merged_remote_rank_count"] = len(payloads) - merged["metadata"] = metadata - return merged - - def _slice_rank_sharded_kv_payload(self, payload: dict[str, Any] | None) -> dict[str, Any] | None: - """Slice a duplicated source-rank KV shard for ``from_tp < to_tp`` cases.""" - if payload is None or self._from_tp >= self._to_tp: - return payload - - tp_ratio = self._to_tp // self._from_tp - shard_index = self._local_rank % tp_ratio - layer_blocks = payload.get("layer_blocks") if isinstance(payload, dict) else None - if not isinstance(layer_blocks, dict): - return payload - - def _slice_tensor_list(name: str) -> list[torch.Tensor | None]: - sliced: list[torch.Tensor | None] = [] - for tensor in layer_blocks.get(name, []): - if not isinstance(tensor, torch.Tensor) or tensor.ndim < 2: - sliced.append(tensor) - continue - head_dim = tensor.shape[-2] - if head_dim % tp_ratio != 0: - sliced.append(tensor) - continue - per_rank = head_dim // tp_ratio - start = shard_index * per_rank - sliced.append(tensor.narrow(-2, start, per_rank).contiguous()) - return sliced - - payload = dict(payload) - payload["layer_blocks"] = { - "key_cache": _slice_tensor_list("key_cache"), - "value_cache": _slice_tensor_list("value_cache"), - } - metadata = dict(payload.get("metadata", {})) - metadata["sliced_for_local_rank"] = self._local_rank - payload["metadata"] = metadata - return payload - - def should_replicate_payload(self) -> bool: - """Whether non-KV payloads should be replicated across ranks. - - Data payloads (stage inputs, chunks) are identical after all-gather, - so only rank 0 transfers them. KV payloads are rank-specific and - all ranks participate. - """ - return self._local_rank != 0 - - def get_kv_rank_mapping(self) -> dict[str, Any]: - """Return the current rank mapping configuration. - - Useful for debugging and for downstream code that needs to know - the TP topology without re-parsing model config. - """ - return { - "from_tp": self._from_tp, - "to_tp": self._to_tp, - "local_rank": self._local_rank, - "remote_ranks": self.get_kv_remote_ranks(), - "is_data_transfer_rank": self.is_data_transfer_rank(), - } - - # ------------------------------------------------------------------ # - # KV transfer lifecycle (RFC – mixin-owned) - # ------------------------------------------------------------------ # - - def mark_kv_transfer( - self, - req_id: str, - seq_len: int, - block_ids: list[int], - custom_metadata: dict[str, Any] | None = None, - ) -> None: - """Mark a request as needing KV cache transfer. - - Called by the scheduler when a transfer trigger fires. The mixin - owns the lifecycle from this point: pending → active → completed. - """ - if req_id in self._kv_pending_transfers: - return - self._kv_triggered_requests.add(req_id) - transfer = { - "seq_len": seq_len, - "block_ids": block_ids, - } - if custom_metadata is not None: - transfer["custom_metadata"] = custom_metadata - self._kv_pending_transfers[req_id] = transfer - - def drain_pending_kv_transfers(self) -> dict[str, dict[str, Any]]: - """Drain pending KV transfers and move them to active. - - Returns ``{req_id: {seq_len, block_ids}}`` for the model runner - to submit to ``send_kv_cache``. - """ - if not self._kv_pending_transfers: - return {} - pending = dict(self._kv_pending_transfers) - self._kv_active_transfers.update(pending.keys()) - self._kv_pending_transfers.clear() - return pending - - def ack_kv_transfers(self, req_ids: list[str] | set[str]) -> None: - """Acknowledge completed KV transfers (from kv_extracted_req_ids). - - Moves requests from active to completed so the scheduler can - safely free their blocks. - """ - for req_id in req_ids: - self._kv_active_transfers.discard(req_id) - self._kv_completed_transfers.add(req_id) - - def drain_completed_kv_transfers(self) -> set[str]: - """Drain and return completed KV transfer request IDs. - - The scheduler calls this to know which requests' blocks can be freed. - """ - completed = set(self._kv_completed_transfers) - self._kv_completed_transfers.clear() - return completed - - def is_kv_transfer_triggered(self, req_id: str) -> bool: - """Check if a request has already triggered KV transfer.""" - return req_id in self._kv_triggered_requests - - def has_pending_kv_work(self) -> bool: - """True if any KV transfers are pending, active, or awaiting ack.""" - return bool(self._kv_pending_transfers or self._kv_active_transfers or self._kv_completed_transfers) - - # Output aggregation - # ------------------------------------------------------------------ # - - def get_omni_connector_output(self) -> OmniConnectorOutput: - """Collect and reset transfer results for this execute_model cycle. - - ``request_metadata`` carries only lightweight scheduling metadata. - Full payloads remain owned by the Model Runner local cache for all - paths. - """ - if not hasattr(self, "_lock"): - return OmniConnectorOutput() - - tp_group = self._get_local_tp_group() - if self._async_chunk and tp_group is not None and getattr(tp_group, "world_size", 1) > 1: - if self.is_data_transfer_rank(): - with self._lock: - fanout_packet = self._collect_async_chunk_fanout_packet_locked() - else: - fanout_packet = None - fanout_packet = self._broadcast_tp_payload_packet(fanout_packet) - if fanout_packet is None: - newly_finished = set() - chunk_finished = set() - request_metadata = {} - else: - if not self.is_data_transfer_rank(): - self._apply_async_chunk_fanout_packet(fanout_packet) - newly_finished = set(fanout_packet["newly_finished"]) - chunk_finished = set(fanout_packet["chunk_finished"]) - request_metadata = dict(fanout_packet["request_metadata"]) - else: - with self._lock: - newly_finished = set(self._finished_load_reqs) - self._finished_load_reqs.clear() - chunk_finished = set(self._chunk_finished_req_ids) - self._chunk_finished_req_ids.clear() - request_metadata = dict(self._local_request_metadata) - self._local_request_metadata.clear() - # _send_side_request_payload is the async accumulation buffer for - # future recv chunks. Clearing it on every consumable wake-up drops - # intermediate - # thinker decode spans before the model side can consume them. - # Only terminal chunk_finished requests may release that buffer. - for req_id in chunk_finished: - if req_id not in self._local_stage_payload_cache: - continue - ext_req_id = self._request_ids_mapping.get(req_id, req_id) - self._send_side_request_payload.pop(ext_req_id, None) - if ext_req_id != req_id: - self._send_side_request_payload.pop(req_id, None) - self._chunk_ready_req_ids.update(newly_finished) - - output = OmniConnectorOutput( - chunk_ready_req_ids=set(self._chunk_ready_req_ids), - chunk_finished_req_ids=chunk_finished, - request_metadata=request_metadata, - kv_sent_req_ids=list(self._kv_sent_req_ids), - stage_recv_req_ids=set(self._stage_recv_req_ids), - has_pending_kv_work=self.has_pending_kv_work(), - ) - if output.stage_recv_req_ids or chunk_finished or newly_finished: - logger.debug( - "[Stage-%s] get_omni_connector_output: stage_recv=%s, chunk_finished=%s, chunk_ready=%s", - self._stage_id, - output.stage_recv_req_ids, - chunk_finished, - output.chunk_ready_req_ids, - ) - self._chunk_ready_req_ids.clear() - self._kv_sent_req_ids.clear() - self._stage_recv_req_ids.clear() - return output - - @staticmethod - def _connector_output_has_signals(output: OmniConnectorOutput) -> bool: - return bool( - output.chunk_ready_req_ids - or output.chunk_finished_req_ids - or output.request_metadata - or output.kv_sent_req_ids - or output.stage_recv_req_ids - or output.has_pending_kv_work - ) - - def attach_omni_connector_output(self, result: Any | None) -> Any: - omni_output = self.get_omni_connector_output() - if not self._connector_output_has_signals(omni_output): - return result - - from copy import copy - - from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT - - wrapped = copy(result if result is not None else EMPTY_MODEL_RUNNER_OUTPUT) - wrapped.omni_connector_output = omni_output - return wrapped - - # ------------------------------------------------------------------ # - # Properties for compatibility with custom_process funcs that access - # transfer_manager.put_req_chunk / request_payload / code_prompt_token_ids - # ------------------------------------------------------------------ # - - @property - def put_req_chunk(self) -> dict[str, int]: - return self._put_req_chunk - - @property - def ramp_chunk_count(self) -> dict[str, int]: - return self._ramp_chunk_count - - @property - def request_payload(self) -> dict[str, dict[str, Any]]: - return self._send_side_request_payload - - @request_payload.setter - def request_payload(self, value: dict[str, dict[str, Any]]) -> None: - self._send_side_request_payload = value - - @property - def code_prompt_token_ids(self) -> dict[str, list[list[int]]]: - return self._code_prompt_token_ids - - @property - def connector(self) -> Any | None: - return self._omni_connector - - # ------------------------------------------------------------------ # - # Background I/O threads - # ------------------------------------------------------------------ # - - def _recv_loop(self) -> None: - """Background thread: poll connector for incoming data.""" - _recv_poll_count = 0 - while not self._stop_event.is_set(): - with self._lock: - pending_ids = list(self._pending_load_reqs.keys()) - - if not pending_ids: - self._work_available.wait(timeout=0.01) - self._work_available.clear() - continue - - _recv_poll_count += 1 - if _recv_poll_count % 5000 == 1: - logger.debug( - "[Stage-%s] _recv_loop: polling %s pending reqs: %s (poll#%s)", - self._stage_id, - len(pending_ids), - pending_ids[:5], - _recv_poll_count, - ) - - made_progress = False - for req_id in pending_ids: - if self._stop_event.is_set(): - break - try: - made_progress = self._poll_single_request(req_id) or made_progress - except Exception: - logger.warning("Error receiving data for %s", req_id, exc_info=True) - - if not made_progress and not self._stop_event.is_set(): - self._work_available.wait(timeout=0.005) - self._work_available.clear() - - _MAX_SEND_RETRIES = 3 - - def _save_loop(self) -> None: - """Background thread: send outgoing data via connector.""" - while not self._stop_event.is_set(): - task = None - with self._lock: - for req_id in list(self._pending_save_reqs.keys()): - dq = self._pending_save_reqs[req_id] - if dq: - task = dq.popleft() - if not dq: - del self._pending_save_reqs[req_id] - break - del self._pending_save_reqs[req_id] - - if task is not None: - success = False - try: - success = self._send_single_request(task) - except Exception: - logger.error( - "Error saving data for %s", - task.get("request_id"), - exc_info=True, - ) - if not success: - self._requeue_or_drop_failed_send(task) - continue - - self._work_available.wait(timeout=0.01) - self._work_available.clear() - - def _requeue_or_drop_failed_send(self, task: dict) -> None: - """Re-enqueue a failed send task or drop it after max retries.""" - retry_count = task.get("_retry_count", 0) + 1 - req_id = task.get("request_id") - if retry_count <= self._MAX_SEND_RETRIES: - task["_retry_count"] = retry_count - logger.warning( - "[Stage-%s] Re-enqueuing failed send for %s (retry %d/%d)", - getattr(self, "_stage_id", "?"), - req_id, - retry_count, - self._MAX_SEND_RETRIES, - ) - with self._lock: - dq = self._pending_save_reqs.setdefault(req_id, deque()) - dq.appendleft(task) - else: - logger.error( - "[Stage-%s] Giving up on send for %s after %d retries", - getattr(self, "_stage_id", "?"), - req_id, - self._MAX_SEND_RETRIES, - ) - self._decrement_pending_save_count(req_id) - - # ------------------------------------------------------------------ # - # Chunk-level poll / send (ported from OmniChunkTransferAdapter) - # ------------------------------------------------------------------ # - - def _poll_single_request(self, req_id: str) -> bool: - """Poll connector for one chunk of a request (non-blocking).""" - connector = self._omni_connector - if connector is None: - return False - - if self._async_chunk and self._model_mode != "ar": - with self._lock: - staged_payload = self._local_stage_payload_cache.get(req_id) - metadata_in_flight = req_id in self._local_request_metadata - scheduler_wakeup_pending = req_id in self._finished_load_reqs - if self._payload_is_consumable(staged_payload) or metadata_in_flight or scheduler_wakeup_pending: - logger.debug( - "[Stage-%s] delaying recv for req=%s until staged async payload is handed to scheduler", - self._stage_id, - req_id, - ) - return False - - target_stage_id = self._stage_id - 1 - chunk_id = self._get_req_chunk[req_id] - external_req_id = self._request_ids_mapping.get(req_id, req_id) - connector_get_key = f"{external_req_id}_{target_stage_id}_{chunk_id}" - - if self._async_chunk: - result = self._recv_async_chunk_result( - connector, - str(target_stage_id), - str(self._stage_id), - connector_get_key, - ) - else: - result = self._recv_full_payload_result( - connector, - str(target_stage_id), - str(self._stage_id), - connector_get_key, - ) - - if result is None: - return False - - payload_data, _size = result - if not payload_data: - return False - if isinstance(payload_data, dict): - logger.debug( - "[Stage-%s] recv_chunk_result: req=%s ext=%s key=%s keys=%s finished=%s", - self._stage_id, - req_id, - external_req_id, - connector_get_key, - sorted(payload_data.keys()), - self._payload_finished(payload_data), - ) - - self._get_req_chunk[req_id] += 1 - - if self._async_chunk: - is_finished = self._payload_finished(payload_data) - incoming_payload_consumable = self._payload_is_consumable(payload_data) - - if self._model_mode == "ar": - payload_data = self._accumulate_payload(external_req_id, payload_data) - payload_consumable = incoming_payload_consumable - else: - new_ids = self._payload_audio_codes(payload_data) or [] - if not new_ids and not is_finished: - return False - payload_consumable = self._payload_is_consumable(payload_data) - - with self._lock: - if is_finished: - self._chunk_finished_req_ids.add(req_id) - self._chunk_stream_completed.add(req_id) - # Local cache (RFC §2.4) — merge, don't replace, so that - # earlier chunk keys (e.g. thinker_prefill_embeddings from - # chunk 0) are not overwritten by later chunks. - existing = self._local_stage_payload_cache.get(req_id) - if existing is not None and isinstance(existing, dict) and isinstance(payload_data, dict): - existing.update(payload_data) - else: - self._local_stage_payload_cache[req_id] = payload_data - staged_payload = self._local_stage_payload_cache[req_id] - self._async_chunk_updated_req_ids.add(req_id) - self.put_local_request_metadata(req_id, self._extract_scheduling_metadata(staged_payload)) - # A finish-only sentinel still needs one terminal wake-up so - # the downstream stage can sync the merged local payload and - # flush/finish even when the last recv carries no new - # consumable chunk bytes. - if payload_consumable or is_finished: - self._finished_load_reqs.add(req_id) - if is_finished and not payload_consumable: - logger.debug( - "[Stage-%s] finish sentinel arrived for req=%s without new consumable payload", - self._stage_id, - req_id, - ) - elif not payload_consumable: - logger.debug( - "[Stage-%s] req=%s received metadata-only / non-consumable async payload; delaying wake-up", - self._stage_id, - req_id, - ) - if is_finished: - self._pending_load_reqs.pop(req_id, None) - else: - # full_payload_mode: the complete payload arrives in a single get(), - # so always unregister immediately. - if isinstance(payload_data, dict): - engine_inputs = payload_data.get("engine_inputs", payload_data) - else: - engine_inputs = payload_data - with self._lock: - self._local_stage_payload_cache[req_id] = self._snapshot_payload(engine_inputs) - # Publish full-payload readiness only after the aligned TP broadcast - # path in recv_full_payload_inputs() has materialized the payload on all - # local ranks. Publishing metadata / stage_recv from the background recv - # thread can let the scheduler observe a request before the payload is - # actually visible to the model thread. - self._full_payload_pending_broadcast_req_ids.add(req_id) - self._pending_load_reqs.pop(req_id, None) - logger.debug( - "[Stage-%s] full_payload recv complete: req=%s key=%s payload_type=%s", - self._stage_id, - req_id, - connector_get_key, - type(engine_inputs).__name__, - ) - - logger.debug("[Stage-%s] Received data for key %s", self._stage_id, connector_get_key) - return True - - def _build_custom_process_payload( - self, - request_id: str | None, - request: Any | None, - pooling_output: Any | None, - ) -> Any | None: - """Run the custom process hook with a best-effort finished kwarg.""" - if self._custom_process_func is None: - return None - - kwargs = { - "transfer_manager": self, - "pooling_output": pooling_output, - "request": request, - } - supports_is_finished = getattr( - self, - "_custom_process_supports_is_finished", - self._custom_process_supports_is_finished_kwarg(), - ) - is_finished_fn = getattr(request, "is_finished", None) - if callable(is_finished_fn): - try: - if supports_is_finished is not False: - kwargs["is_finished"] = bool(is_finished_fn()) - except Exception: - logger.debug("request.is_finished() failed for %s", request_id, exc_info=True) - - try: - return self._custom_process_func(**kwargs) - except TypeError as exc: - if "is_finished" not in kwargs or not self._is_unexpected_is_finished_kwarg_error(exc): - logger.exception("custom_process_stage_input_func failed for chunk %s", request_id) - return None - kwargs.pop("is_finished", None) - try: - return self._custom_process_func(**kwargs) - except Exception: - logger.exception("custom_process_stage_input_func failed for chunk %s", request_id) - return None - except Exception: - logger.exception("custom_process_stage_input_func failed for chunk %s", request_id) - return None - - def _custom_process_supports_is_finished_kwarg(self) -> bool | None: - """Return whether the custom process hook accepts `is_finished`.""" - if self._custom_process_func is None: - return None - try: - signature = inspect.signature(self._custom_process_func) - except (TypeError, ValueError): - return None - - for param in signature.parameters.values(): - if param.kind == inspect.Parameter.VAR_KEYWORD: - return True - - is_finished_param = signature.parameters.get("is_finished") - if is_finished_param is None: - return False - return is_finished_param.kind in ( - inspect.Parameter.POSITIONAL_OR_KEYWORD, - inspect.Parameter.KEYWORD_ONLY, - ) - - @staticmethod - def _is_unexpected_is_finished_kwarg_error(exc: TypeError) -> bool: - message = str(exc) - return ( - "unexpected keyword argument 'is_finished'" in message - or 'unexpected keyword argument "is_finished"' in message - or "positional-only arguments passed as keyword arguments: 'is_finished'" in message - ) - - def _send_single_request(self, task: dict) -> bool: - """Send one queued task via connector.put(). - - Returns True on success. On failure (put() raises or returns - ``success=False``), returns False **without** decrementing - ``_pending_save_counts`` so the caller can retry or clean up. - """ - connector = self._omni_connector - if connector is None: - return True - - request_id = task.get("request_id") - payload_data = task.get("data") - if payload_data is None and task.get("request") is not None: - payload_data = self._build_custom_process_payload( - request_id=request_id, - request=task.get("request"), - pooling_output=task.get("pooling_output"), - ) - put_key = task.get("put_key") - - success, _size, _metadata = connector.put( - from_stage=str(task["stage_id"]), - to_stage=str(task["next_stage_id"]), - put_key=put_key, - data=payload_data, - ) - logger.debug( - "[Stage-%s] _send_single_request: put_key=%s success=%s size=%s", - task["stage_id"], - put_key, - success, - _size, - ) - - if not success: - return False - - self._decrement_pending_save_count(request_id) - return True - - def _decrement_pending_save_count(self, request_id: str) -> None: - """Decrement pending save count and run deferred cleanup if zero.""" - cleanup_req_id = None - with self._lock: - remaining = self._pending_save_counts.get(request_id, 0) - if remaining > 1: - self._pending_save_counts[request_id] = remaining - 1 - elif remaining == 1: - self._pending_save_counts.pop(request_id, None) - if request_id in self._deferred_send_cleanup: - self._deferred_send_cleanup.remove(request_id) - cleanup_req_id = request_id - if cleanup_req_id is not None: - self._put_req_chunk.pop(cleanup_req_id, None) - self._send_side_request_payload.pop(cleanup_req_id, None) - self._code_prompt_token_ids.pop(cleanup_req_id, None) - self._cached_ic.pop(cleanup_req_id, None) - self._ramp_chunk_count.pop(cleanup_req_id, None) - self._adaptive_states.pop(cleanup_req_id, None) - - # ------------------------------------------------------------------ # - # Payload accumulation (ported from OmniChunkTransferAdapter) - # ------------------------------------------------------------------ # - - def _accumulate_payload(self, req_id: str, payload_data: OmniPayload) -> OmniPayload: - """Accumulate chunk payloads (concat tensors, extend lists).""" - if req_id not in self._send_side_request_payload: - self._send_side_request_payload[req_id] = dict(payload_data) - return dict(self._send_side_request_payload[req_id]) - - origin = self._send_side_request_payload[req_id] - merged = dict(origin) - raw_ok = payload_data.get("meta", {}).get("override_keys", []) if isinstance(payload_data, dict) else [] - override_keys = {tuple(k) if isinstance(k, list) else k for k in raw_ok} - - for key, value in payload_data.items(): - if isinstance(value, dict): - origin_sub = origin.get(key) - merged_sub = dict(origin_sub) if isinstance(origin_sub, dict) else {} - span_handled: set[str] = set() - if key == "embed" and isinstance(origin_sub, dict): - for tk, sk, ek in _EMBED_SPAN_GROUPS: - if tk not in value or (key, tk) in override_keys: - continue - span = merge_tensor_spans( - get_tensor_span(origin_sub, tensor_key=tk, start_key=sk, end_key=ek), - get_tensor_span(value, tensor_key=tk, start_key=sk, end_key=ek), - ) - if span is None: - continue - t, s, e = span - merged_sub[tk] = t - merged_sub[sk] = s - merged_sub[ek] = e - span_handled |= {tk, sk, ek} - for qual, qval in value.items(): - if qual in span_handled: - continue - if key == "meta" and qual == "finished": - merged_sub[qual] = qval - continue - if (key, qual) in override_keys: - merged_sub[qual] = qval - continue - osv = merged_sub.get(qual) - if isinstance(qval, torch.Tensor) and isinstance(osv, torch.Tensor): - merged_sub[qual] = torch.cat([osv, qval], dim=0) - elif isinstance(qval, list) and isinstance(osv, list): - merged_sub[qual] = osv + qval - else: - merged_sub[qual] = qval - merged[key] = merged_sub - else: - if key in override_keys: - merged[key] = value - continue - ov = origin.get(key) - if isinstance(value, torch.Tensor) and isinstance(ov, torch.Tensor): - merged[key] = torch.cat([ov, value], dim=0) - elif isinstance(value, list) and isinstance(ov, list): - merged[key] = ov + value - else: - merged[key] = value - - self._send_side_request_payload[req_id] = merged - return dict(merged) - - def drop_inactive_request_runtime_state(self, req_id: str) -> None: - """Clear inactive request state used by both the runner and mixin. - - This centralizes the model-runner-side cleanup pattern so - ``OmniGPUModelRunner`` can reuse it instead of open-coding the same - inactive-request state mutations. - """ - if hasattr(self, "model_intermediate_buffer"): - self.model_intermediate_buffer.pop(req_id, None) - self.drop_inactive_request_delivery_state(req_id) - - # ------------------------------------------------------------------ # - # Helpers - # ------------------------------------------------------------------ # - - @staticmethod - def _create_connector(model_config: Any) -> OmniConnectorBase | None: - """Create a connector from model_config, or None if unconfigured.""" - connector_config = getattr(model_config, "stage_connector_config", None) - if connector_config is None: - return None - - if not isinstance(connector_config, dict): - connector_config = { - "name": getattr(connector_config, "name", None), - "extra": getattr(connector_config, "extra", None), - } - - name = connector_config.get("name") - if not isinstance(name, str) or not name.strip(): - raise RuntimeError("Invalid stage connector config: missing connector name") - name = name.strip() - - extra = connector_config.get("extra") - if extra is None: - extra = {} - elif not isinstance(extra, dict): - raise RuntimeError(f"Invalid extra config for connector {name}: expected dict, got {type(extra).__name__}") - - spec = ConnectorSpec(name=name, extra=extra) - try: - return OmniConnectorFactory.create_connector(spec) - except Exception as exc: - raise RuntimeError(f"Failed to create connector {name}") from exc - - @staticmethod - def _load_custom_func(model_config: Any) -> tuple[str | None, Any | None]: - """Load the connector payload builder for the downstream stage. - - Preferred source is ``custom_process_next_stage_input_func``. Some - full_payload_mode configs (async_chunk=false) only expose the next-stage prompt builder via - ``custom_process_input_func`` (for example ``thinker2talker``), while the - connector payload builder lives beside it as ``thinker2talker_full_payload``. - In that case, derive the full_payload_mode builder path automatically. - """ - candidates: list[str] = [] - - next_stage_func = getattr(model_config, "custom_process_next_stage_input_func", None) - if isinstance(next_stage_func, str) and next_stage_func: - candidates.append(next_stage_func) - - if not getattr(model_config, "async_chunk", False): - input_func = getattr(model_config, "custom_process_input_func", None) - if isinstance(input_func, str) and input_func: - try: - module_path, func_name = input_func.rsplit(".", 1) - if func_name.endswith("_full_payload") or func_name.endswith("_batch"): - candidates.append(f"{module_path}.{func_name}") - else: - candidates.append(f"{module_path}.{func_name}_full_payload") - candidates.append(f"{module_path}.{func_name}_batch") - candidates.append(input_func) - except ValueError: - candidates.append(input_func) - - tried: set[str] = set() - for func_path in candidates: - if func_path in tried: - continue - tried.add(func_path) - try: - module_path, func_name = func_path.rsplit(".", 1) - module = importlib.import_module(module_path) - func = getattr(module, func_name, None) - if callable(func): - if not OmniConnectorModelRunnerMixin._is_connector_payload_builder(func): - logger.debug( - "Skipping incompatible connector payload hook %s; signature=%s", - func_path, - inspect.signature(func), - ) - continue - return func_path, func - except Exception: - logger.warning("Failed to load custom func: %s", func_path, exc_info=True) - - return None, None - - @staticmethod - def _is_connector_payload_builder(func: Any) -> bool: - """Whether *func* matches the mixin payload-builder contract.""" - try: - signature = inspect.signature(func) - except (TypeError, ValueError): - return False - - params = signature.parameters - if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in params.values()): - return True - - required = {"transfer_manager", "pooling_output", "request"} - supported = { - name - for name, param in params.items() - if param.kind - in ( - inspect.Parameter.POSITIONAL_OR_KEYWORD, - inspect.Parameter.KEYWORD_ONLY, - ) - } - return required.issubset(supported) - - def _resolve_external_req_id(self, request: Any, fallback_req_id: str) -> str: - """Resolve the external request ID consistently. - - Checks ``_request_ids_mapping`` first (populated by - ``register_chunk_recv``), then falls back to the request's - ``external_req_id`` attribute, and finally to the given - ``fallback_req_id``. - """ - mapped = self._request_ids_mapping.get(fallback_req_id) - if mapped is not None: - return mapped - if request is not None: - # external_req_id may be explicitly None; fall back. - ext = getattr(request, "external_req_id", None) - if ext is not None: - return ext - return fallback_req_id - - def _resolve_next_stage_id(self, model_config: Any) -> int: - """Determine the downstream stage ID from connector config. - - Falls back to ``stage_id + 1`` when the config does not specify - a ``to_stage`` explicitly. - """ - connector_config = getattr(model_config, "stage_connector_config", None) - if connector_config is not None: - if isinstance(connector_config, dict): - to_stage = connector_config.get("to_stage") - else: - to_stage = getattr(connector_config, "to_stage", None) - if isinstance(to_stage, int): - return to_stage - if isinstance(to_stage, str) and to_stage.strip(): - return int(to_stage) - return self._stage_id + 1 - - @staticmethod - def _parse_rank_mapping(model_config: Any) -> dict[str, int]: - """Parse rank_mapping from connector config (optional). - - Returns ``{"from_tp": int, "to_tp": int, "local_rank": int}``. - When ``rank_mapping`` is absent, assumes 1:1 homogeneous mapping. - """ - connector_config = getattr(model_config, "stage_connector_config", None) - if connector_config is not None and not isinstance(connector_config, dict): - connector_config = getattr(connector_config, "__dict__", {}) - - rank_mapping: dict = {} - if isinstance(connector_config, dict): - rank_mapping = connector_config.get("rank_mapping", {}) - - from_tp = int(rank_mapping.get("from_tp", 1)) - to_tp = int(rank_mapping.get("to_tp", 1)) - - local_rank = 0 - try: - local_rank = int(os.environ.get("LOCAL_RANK", "0")) - except (ValueError, TypeError): - pass - - return {"from_tp": from_tp, "to_tp": to_tp, "local_rank": local_rank} - - # ------------------------------------------------------------------ # - # Heterogeneous TP rank support - # ------------------------------------------------------------------ # - - def _validate_kv_tp_topology(self) -> None: - """Reject heterogeneous TP mappings that cannot be routed losslessly.""" - if self._from_tp <= 0 or self._to_tp <= 0: - raise ValueError(f"Invalid KV TP mapping: from_tp={self._from_tp}, to_tp={self._to_tp}") - larger = max(self._from_tp, self._to_tp) - smaller = min(self._from_tp, self._to_tp) - if larger % smaller != 0: - raise ValueError( - f"KV TP mapping must be divisible for rank-aware routing: from_tp={self._from_tp}, to_tp={self._to_tp}" - ) - - def get_kv_remote_ranks(self) -> list[int]: - """Determine which remote ranks this local rank exchanges KV with. - - Follows vLLM's ``TpKVTopology.get_target_remote_ranks()`` pattern: - - ``from_tp > to_tp``: each to-rank reads from multiple from-ranks - - ``from_tp < to_tp``: multiple to-ranks read from the same from-rank - - ``from_tp == to_tp``: 1:1 mapping - """ - self._validate_kv_tp_topology() - if self._from_tp == self._to_tp: - return [self._local_rank] - - if self._from_tp > self._to_tp: - tp_ratio = self._from_tp // self._to_tp - return [self._local_rank * tp_ratio + i for i in range(tp_ratio)] - else: - tp_ratio = self._to_tp // self._from_tp - return [self._local_rank // tp_ratio] - - def is_data_transfer_rank(self) -> bool: - """Whether this rank should participate in data (non-KV) transfer. - - Ordinary stage payloads are TP-identical, so exactly one TP rank - should talk to the connector. When TP is initialized, use TP rank 0 - so the connector leader matches TP-local broadcast source rank. - Otherwise fall back to LOCAL_RANK==0 for the single-rank case. - """ - tp_group = self._get_local_tp_group() - if tp_group is not None and getattr(tp_group, "world_size", 1) > 1: - return getattr(tp_group, "rank_in_group", 0) == 0 - return self._local_rank == 0 - - def get_kv_connector_key( - self, - req_id: str, - from_stage: int, - chunk_id: int, - from_rank: int, - to_rank: int, - ) -> str: - """Build connector key that includes rank info for KV transfers.""" - return f"{req_id}_{from_stage}_{chunk_id}_{from_rank}_{to_rank}" diff --git a/vllm_omni/worker/payload_span.py b/vllm_omni/worker/payload_span.py deleted file mode 100644 index 994392343a9..00000000000 --- a/vllm_omni/worker/payload_span.py +++ /dev/null @@ -1,64 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Helpers for explicit thinker decode span metadata.""" - -from collections.abc import Mapping -from typing import Any - -import torch - -THINKER_DECODE_EMBEDDINGS_KEY = "thinker_decode_embeddings" -THINKER_OUTPUT_TOKEN_IDS_KEY = "thinker_output_token_ids" -THINKER_DECODE_TOKEN_START_KEY = "thinker_decode_embeddings_token_start" -THINKER_DECODE_TOKEN_END_KEY = "thinker_decode_embeddings_token_end" - -CACHED_THINKER_DECODE_EMBEDDINGS_KEY = "cached_thinker_decode_embeddings" -CACHED_THINKER_DECODE_TOKEN_START_KEY = "cached_thinker_decode_embeddings_token_start" -CACHED_THINKER_DECODE_TOKEN_END_KEY = "cached_thinker_decode_embeddings_token_end" - -TensorSpan = tuple[torch.Tensor, int, int] - - -def get_tensor_span(payload: Mapping[str, Any], *, tensor_key: str, start_key: str, end_key: str) -> TensorSpan | None: - tensor = payload.get(tensor_key) - start = payload.get(start_key) - end = payload.get(end_key) - if not isinstance(tensor, torch.Tensor): - return None - if not isinstance(start, int) or not isinstance(end, int): - return None - if start < 0 or end < start or (end - start) != int(tensor.shape[0]): - return None - return tensor, start, end - - -def merge_tensor_spans(existing_span: TensorSpan | None, incoming_span: TensorSpan | None) -> TensorSpan | None: - if existing_span is None or incoming_span is None: - return None - - existing_tensor, existing_start, existing_end = existing_span - incoming_tensor, incoming_start, incoming_end = incoming_span - if incoming_tensor.device != existing_tensor.device or incoming_tensor.dtype != existing_tensor.dtype: - incoming_tensor = incoming_tensor.to(device=existing_tensor.device, dtype=existing_tensor.dtype) - if incoming_start == existing_end: - return torch.cat([existing_tensor, incoming_tensor], dim=0), existing_start, incoming_end - if incoming_start < existing_end: - overlap = existing_end - incoming_start - if overlap >= int(incoming_tensor.shape[0]): - return existing_tensor, existing_start, existing_end - trimmed_tensor = incoming_tensor[overlap:] - return ( - torch.cat([existing_tensor, trimmed_tensor], dim=0), - existing_start, - existing_end + int(trimmed_tensor.shape[0]), - ) - return None - - -def get_tensor_span_row(span: TensorSpan | None, index: int) -> torch.Tensor | None: - if span is None: - return None - tensor, start, end = span - if index < start or index >= end: - return None - return tensor[index - start] From ab56148505330cfeb935af4b6f453e87287f20b3 Mon Sep 17 00:00:00 2001 From: Mu GuanLin <1203789601@qq.com> Date: Wed, 9 Sep 2026 14:37:01 +0800 Subject: [PATCH 23/93] [Bugfix] Make LTX vocoder decoding deterministic (#7231) Signed-off-by: mglyn <1203789601@qq.com> --- tests/diffusion/models/ltx2/test_ltx2_vae.py | 23 ++++++++++ tests/e2e/accuracy/ltx/run_ltx25_reference.py | 26 ++++++++++- .../diffusion/models/ltx2/ltx2_runtime.py | 46 ++++++++++++------- 3 files changed, 77 insertions(+), 18 deletions(-) diff --git a/tests/diffusion/models/ltx2/test_ltx2_vae.py b/tests/diffusion/models/ltx2/test_ltx2_vae.py index 964bb7d06cb..b9bc52d9452 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_vae.py +++ b/tests/diffusion/models/ltx2/test_ltx2_vae.py @@ -12,6 +12,29 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@pytest.mark.parametrize("initial_deterministic", [False, True]) +@pytest.mark.parametrize("fail", [False, True]) +def test_ltx_vocoder_deterministic_context(monkeypatch, initial_deterministic, fail): + from vllm_omni.diffusion.models.ltx2.ltx2_runtime import _deterministic_ltx_vocoder + + monkeypatch.setattr(torch.backends.cudnn, "deterministic", initial_deterministic) + if torch.backends.cudnn.is_available(): + monkeypatch.setattr(torch.backends.cudnn, "benchmark_limit", 17) + settings = { + name: getattr(torch.backends.cudnn, name) for name in ("enabled", "benchmark", "benchmark_limit", "allow_tf32") + } + try: + with _deterministic_ltx_vocoder(): + assert torch.backends.cudnn.deterministic + assert {name: getattr(torch.backends.cudnn, name) for name in settings} == settings + if fail: + raise RuntimeError("injected failure") + except RuntimeError as exc: + assert fail and str(exc) == "injected failure" + assert torch.backends.cudnn.deterministic == initial_deterministic + assert {name: getattr(torch.backends.cudnn, name) for name in settings} == settings + + def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch): import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime diff --git a/tests/e2e/accuracy/ltx/run_ltx25_reference.py b/tests/e2e/accuracy/ltx/run_ltx25_reference.py index cbc8917a869..14c862555e9 100644 --- a/tests/e2e/accuracy/ltx/run_ltx25_reference.py +++ b/tests/e2e/accuracy/ltx/run_ltx25_reference.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 """Generate raw LTX video and audio outputs with the official or Omni runtime.""" @@ -114,6 +114,29 @@ def _insert_official_paths(official_root: Path) -> None: sys.path.insert(0, path) +def _configure_official_vocoder_determinism() -> None: + """Make the pinned official LTX-2.5 vocoder a stable reference.""" + from ltx_core.model.audio_vae.vocoder import VocoderWithBWE + + original_forward = VocoderWithBWE.forward + if getattr(original_forward, "_vllm_omni_deterministic", False): + return + + def deterministic_forward(self: Any, mel_spec: torch.Tensor) -> torch.Tensor: + device_type = mel_spec.device.type + if device_type != "cuda": + return original_forward(self, mel_spec) + previous = torch.backends.cudnn.deterministic + try: + torch.backends.cudnn.deterministic = True + return original_forward(self, mel_spec) + finally: + torch.backends.cudnn.deterministic = previous + + setattr(deterministic_forward, "_vllm_omni_deterministic", True) + setattr(VocoderWithBWE, "forward", deterministic_forward) + + def _configure_official_sdpa(pipeline: Any) -> None: """Pin official connector and denoiser attention to cuDNN, with MATH fallback.""" from ltx_core.loader.attention_ops import set_attention_module_op @@ -187,6 +210,7 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None: if args.official_root is None: raise ValueError("Official backend requires --official-root") _insert_official_paths(args.official_root) + _configure_official_vocoder_determinism() # Keep Gemma on the same SDPA path in both subprocesses. DiT attention is # still pinned independently to the explicit cuDNN backend below. torch.backends.cuda.enable_cudnn_sdp(False) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py index 2c51b17d90b..9e682ffef42 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py @@ -65,27 +65,39 @@ logger = init_logger(__name__) +@contextmanager +def _deterministic_ltx_vocoder(): + previous = torch.backends.cudnn.deterministic + try: + torch.backends.cudnn.deterministic = True + yield + finally: + torch.backends.cudnn.deterministic = previous + + def _run_ltx_vocoder(vocoder: nn.Module, generated_mel: torch.Tensor) -> torch.Tensor: """Run the BWE vocoder in FP32, matching the official LTX pipeline.""" - if not hasattr(vocoder, "bwe_generator"): - return vocoder(generated_mel) - - input_dtype = generated_mel.dtype device_type = generated_mel.device.type - module_dtype = next(vocoder.parameters()).dtype - if device_type == "mps": - if module_dtype != torch.float32: - vocoder.float() - try: - return vocoder(generated_mel.float()).to(input_dtype) - finally: + cudnn_context = _deterministic_ltx_vocoder() if device_type == "cuda" else nullcontext() + with cudnn_context: + if not hasattr(vocoder, "bwe_generator"): + return vocoder(generated_mel) + + input_dtype = generated_mel.dtype + module_dtype = next(vocoder.parameters()).dtype + if device_type == "mps": if module_dtype != torch.float32: - vocoder.to(module_dtype) - - # BF16 errors compound through the BWE model's long convolution stack. - # FP32 autocast upcasts each op without materializing a second FP32 model. - with torch.autocast(device_type=device_type, dtype=torch.float32): - return vocoder(generated_mel.float()).to(input_dtype) + vocoder.float() + try: + return vocoder(generated_mel.float()).to(input_dtype) + finally: + if module_dtype != torch.float32: + vocoder.to(module_dtype) + + # BF16 errors compound through the BWE model's long convolution stack. + # FP32 autocast upcasts each op without materializing a second FP32 model. + with torch.autocast(device_type=device_type, dtype=torch.float32): + return vocoder(generated_mel.float()).to(input_dtype) def _prepare_ltx2_video_output( From 7712701578d8359b5af958bb2df9023c90e21a6d Mon Sep 17 00:00:00 2001 From: Sparks <41097544+Sparks-M@users.noreply.github.com> Date: Wed, 9 Sep 2026 15:54:46 +0800 Subject: [PATCH 24/93] [Doc] [Recipe] Add FLUX.1-schnell recipe for RTX 5090 32GB (#7299) Signed-off-by: Sparks-M <41097544+Sparks-M@users.noreply.github.com> --- recipes/README.md | 1 + recipes/black-forest-labs/FLUX.1-schnell.md | 160 ++++++++++++++++++++ 2 files changed, 161 insertions(+) create mode 100644 recipes/black-forest-labs/FLUX.1-schnell.md diff --git a/recipes/README.md b/recipes/README.md index 8ac877567fa..6b80f8a0343 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -42,6 +42,7 @@ recipes/ | -------- | ------ | ---------- | | [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB | | [`Bagel/BAGEL-7B-MoT.md`](./Bagel/BAGEL-7B-MoT.md) | Text-to-image with shared online/offline examples | 1x A100 80GB / 2x CUDA GPUs | +| [`black-forest-labs/FLUX.1-schnell.md`](./black-forest-labs/FLUX.1-schnell.md) | Offline text-to-image generation with CPU offload | 1x NVIDIA GeForce RTX 5090 32GB | | [`Boogu/Boogu-Image.md`](./Boogu/Boogu-Image.md) | Text-to-image online serving (Boogu-Image-0.1-Base) | 1x A100/H100 40GB+ | | [`BosonAI/Higgs-Audio-V3-TTS.md`](./BosonAI/Higgs-Audio-V3-TTS.md) | Online + offline multilingual TTS with voice cloning | 1x H100 80GB | | [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB | diff --git a/recipes/black-forest-labs/FLUX.1-schnell.md b/recipes/black-forest-labs/FLUX.1-schnell.md new file mode 100644 index 00000000000..9ae6f481741 --- /dev/null +++ b/recipes/black-forest-labs/FLUX.1-schnell.md @@ -0,0 +1,160 @@ +# FLUX.1-schnell + +> Offline text-to-image generation on one NVIDIA GeForce RTX 5090 32 GB + +## Summary + +- Vendor: Black Forest Labs +- Model: `black-forest-labs/FLUX.1-schnell` +- Task: Text-to-image generation +- Mode: Offline inference with the shared text-to-image runner +- Maintainer: Community + +## When to use this recipe + +Use this recipe as a personally validated starting point for generating a +1024x1024 image with FLUX.1-schnell on one NVIDIA GeForce RTX 5090. The tested +configuration uses BF16, four denoising steps and model-level CPU offload. + +This recipe covers single-image generation. It does not attempt to optimize +throughput with quantization or cache acceleration. + +## References + +- Model card: +- Canonical offline text-to-image guide: + [`docs/user_guide/examples/offline_inference/text_to_image.md`](../../docs/user_guide/examples/offline_inference/text_to_image.md) +- Shared runnable example: + [`examples/offline_inference/text_to_image`](../../examples/offline_inference/text_to_image) +- CPU offload guide: + [`docs/user_guide/diffusion/cpu_offload.md`](../../docs/user_guide/diffusion/cpu_offload.md) +- Community recipe tracker: + [#2645](https://github.com/vllm-project/vllm-omni/issues/2645) + +## Hardware Support + +## GPU + +### 1x NVIDIA GeForce RTX 5090 32 GB + +This configuration was personally validated end to end on September 9, 2026. + +#### Environment + +- OS: Ubuntu 22.04.3 x86_64 +- Python: 3.12.7 +- GPU: 1x NVIDIA GeForce RTX 5090, 32,607 MiB reported by `nvidia-smi` +- Host RAM: approximately 62 GiB (64 GB class) +- NVIDIA driver: 595.71.05; CUDA 13.2 reported by `nvidia-smi` +- PyTorch: 2.13.0+cu130; PyTorch CUDA runtime: 13.0 +- vLLM: 0.28.0 +- vLLM-Omni: `0.28.1.dev112+ga31690041` +- Diffusers: 0.40.0 +- Transformers: 5.14.1 + +#### Command + +Run the shared example from the vLLM-Omni repository root in an environment +with vLLM-Omni installed: + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model black-forest-labs/FLUX.1-schnell \ + --prompt 'A red fox standing in fresh snow, cinematic photography' \ + --height 1024 --width 1024 \ + --num-inference-steps 4 --guidance-scale 0 --seed 42 \ + --enable-cpu-offload \ + --output flux-schnell-rtx5090.png +``` + +The command uses the public model ID for portability. The measured runs used +already downloaded FLUX.1-schnell weights; initial Hub download and +authentication were not part of the inference validation. + +#### Verification + +Verify that the runner succeeded and inspect the saved image: + +```bash +python - <<'PY' +import hashlib +from pathlib import Path + +from PIL import Image, ImageStat + +path = Path("flux-schnell-rtx5090.png") +with Image.open(path) as image: + image.load() + assert image.format == "PNG" + assert image.size == (1024, 1024) + rgb = image.convert("RGB") + stats = ImageStat.Stat(rgb) + assert max(stats.stddev) > 1, "Image is nearly uniform" + print(f"format={image.format}") + print(f"size={image.width}x{image.height}") + print(f"mode={image.mode}") + print(f"bytes={path.stat().st_size}") + print("rgb_mean=" + ",".join(f"{value:.3f}" for value in stats.mean)) + print(f"rgb_extrema={rgb.getextrema()}") + +print(f"sha256={hashlib.sha256(path.read_bytes()).hexdigest()}") +PY +``` + +The repeated experiment completed 36 image requests (6 additional warmups and +30 measured outputs) with successful process exits. Every PNG passed decoding, +1024x1024 dimension and non-uniform pixel checks. SHA-256 hashes were recorded +for the artifacts; matching those hashes is not required across environments. +These checks establish file integrity, not a quantitative image-quality score. + +#### Performance + +The experiment used three independent processes. Each process initialized the +model once, generated two warmup images, then generated ten measured images +sequentially. All 30 measured images used BF16, CPU offload, 1024x1024, four +steps, guidance scale 0 and TP=1, without concurrent requests. + +| Generation measurement | Result | +| --- | ---: | +| Measured images | 30 | +| Mean latency | 25.8411 seconds | +| P95 latency (nearest-rank) | 26.1446 seconds | + +| Startup / resource measurement | Maximum | +| --- | ---: | +| Model initialization (seconds) | 64.73 | +| Full-lifetime RAM PSS peak (MiB) | 49306 | +| Full-lifetime device GPU memory peak (MiB) | 28121 | +| Measured PyTorch reserved GPU memory peak (MiB) | 27360 | + +Latency measures only `Omni.generate()`, excluding image saving and +verification. The two warmup images per process are excluded. P95 uses the +nearest-rank method over the 30 measured images. Model initialization includes +the engine's built-in dummy warmup; weights were already downloaded and OS file +caches were not cleared. + +RAM and whole-device GPU memory were sampled throughout each process lifetime +at a target interval of 200 ms. RAM is the maximum simultaneous sum of PSS for +the inference process and its workers, read from Linux `smaps_rollup`. +Whole-device GPU memory is the maximum reported by `nvidia-smi`. Both are +sampled peaks and may miss shorter spikes. + +The PyTorch value is the maximum worker-reported `peak_memory_mb` across the +30 measured requests. Despite the field name, its unit is MiB; it measures the +PyTorch allocator's reserved-memory high-water mark rather than whole-device +usage. + +#### Notes + +- `--enable-cpu-offload` swaps the transformer and text encoders between host + and device memory at phase boundaries. It reduces GPU residency but adds + transfer overhead and requires sufficient host RAM. The minimum host RAM + requirement was not measured. +- CPU offload does not imply CPU-only computation. The tested model-level + offload backend keeps the VAE on the GPU. +- Only offline single-image generation on one GPU with BF16, 1024x1024, + four steps and CPU offload was validated. Online serving, FP8, other + resolutions, concurrency, execution without offload and smaller host RAM + configurations are outside this recipe's validated scope. +- The upstream revision of the downloaded weights was not recorded, so the + measurements do not establish exact checkpoint reproducibility. From c66df06c0ffd48cf4f7a0b01434673a3cd7467b5 Mon Sep 17 00:00:00 2001 From: chi030303 <106855944+chi030303@users.noreply.github.com> Date: Wed, 9 Sep 2026 15:56:14 +0800 Subject: [PATCH 25/93] [Doc] Qwen3-TTS: add 0.6B on 1x A100 40GB (#7289) Signed-off-by: chi030303 <106855944+chi030303@users.noreply.github.com> --- recipes/Qwen/Qwen3-TTS.md | 88 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 88 insertions(+) diff --git a/recipes/Qwen/Qwen3-TTS.md b/recipes/Qwen/Qwen3-TTS.md index 0d930d58f7a..d0dbe57dfe7 100644 --- a/recipes/Qwen/Qwen3-TTS.md +++ b/recipes/Qwen/Qwen3-TTS.md @@ -263,6 +263,94 @@ curl -X POST http://localhost:8091/v1/audio/speech \ gpu_memory_utilization: 0.15 ``` +### 1x A100 40GB (0.6B CustomVoice) + +#### Environment + +- OS: Linux 4.18.0-553 (RHEL 8.10), x86_64 +- Python: 3.12.14 +- PyTorch: 2.13.0+cu129 +- Driver / runtime: NVIDIA 570.124.06 / CUDA 12.8 (`nvidia-smi`); PyTorch CUDA 12.9 +- GPU: NVIDIA A100-SXM4-40GB, 40960 MiB +- vLLM version: 0.28.0+cu129 +- vLLM-Omni version or commit: `3204a0b2ade0f05424b7f11e3bcc03cb3542e0c6` + +#### Command + +```bash +vllm serve Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \ + --deploy-config vllm_omni/deploy/qwen3_tts.yaml \ + --omni --port 8091 +``` + +The default deploy config (`qwen3_tts.yaml`) works without modification on this +A100 40GB. Both stages (talker + code2wav) share GPU 0 with +`gpu_memory_utilization: 0.3` each. + +#### Verification + +**English synthesis:** + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, this is Qwen3-TTS running on A100 40GB.", + "voice": "vivian", + "language": "English" + }' --output test_english.wav +``` + +**Chinese synthesis:** + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "你好，这是在A100 40GB上运行的语音合成测试。", + "voice": "vivian", + "language": "Chinese" + }' --output test_chinese.wav +``` + +**With emotion instruction:** + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "I am so excited about this!", + "voice": "vivian", + "language": "English", + "instructions": "Speak with great enthusiasm" + }' --output test_emotion.wav +``` + +#### Notes + +- Memory usage: **16283 MiB / 40960 MiB** peak (`nvidia-smi`) during the three + verification requests with the default deploy config + (`gpu_memory_utilization: 0.3` per stage). Stage 0 used 12404 MiB and stage 1 + used 3862 MiB. The 0.6B weights occupy only ~2.4 GiB (Stage 0: 1.91 GiB, + Stage 1: 0.45 GiB); the rest is KV cache reserved as a fraction of device + memory, so the same `0.3` setting reads higher in GiB on 40GB than the + ~13.5 GiB idle figure on the 24GB 4090 profile above. To shrink the + reservation, use the `gpu_memory_utilization: 0.15` override documented in + that 4090 section. +- vLLM 0.28.0's default wheel targets CUDA 13.0. Driver 570 cannot load that + variant. This qualification used the official CUDA 12.9 wheel (same pin as + the A100 Wan2.2 recipe): + + ```bash + uv pip install \ + 'https://github.com/vllm-project/vllm/releases/download/v0.28.0/vllm-0.28.0%2Bcu129-cp38-abi3-manylinux_2_28_x86_64.whl' \ + --torch-backend=cu129 + ``` + + Confirm `torch.__version__` contains `cu129`. If a CPU build is already + installed, `uv pip install --torch-backend=cu129` is a no-op; use + `--reinstall`. + ### 1x AMD MI300X, 1.7B checkpoints #### Environment From ffa991b24fcf9540126ce574f9c7b8fa174af96d Mon Sep 17 00:00:00 2001 From: Mu GuanLin <1203789601@qq.com> Date: Wed, 9 Sep 2026 16:31:16 +0800 Subject: [PATCH 26/93] [Perf][Model] Add optimized LTX-2.5 DiffVAE operators (#7308) Signed-off-by: mglyn <1203789601@qq.com> --- .../ltx2/test_ltx2_diffvae_attention.py | 89 +++ .../models/ltx2/test_ltx2_diffvae_ops.py | 409 ++++++++++++++ .../models/ltx2/test_ltx2_ops_platform.py | 123 ++++ tests/diffusion/models/ltx2/test_ltx2_vae.py | 189 ++++++- .../diffusion/models/ltx2/ltx2_components.py | 14 +- .../diffusion/models/ltx2/ops/__init__.py | 93 ++++ .../diffusion/models/ltx2/ops/_numerics.py | 86 +++ vllm_omni/diffusion/models/ltx2/ops/fna.py | 525 ++++++++++++++++++ .../diffusion/models/ltx2/ops/qk_rms_norm.py | 377 +++++++++++++ .../models/ltx2/ops/residual_adaln.py | 489 ++++++++++++++++ .../diffusion/models/ltx2/ops/rope_tables.py | 50 ++ vllm_omni/diffusion/models/ltx2/ops/swiglu.py | 206 +++++++ .../diffusion/models/ltx2/vae/__init__.py | 4 + .../decoder.py} | 185 ++++++ .../distributed.py} | 2 +- 15 files changed, 2825 insertions(+), 16 deletions(-) create mode 100644 tests/diffusion/models/ltx2/test_ltx2_diffvae_attention.py create mode 100644 tests/diffusion/models/ltx2/test_ltx2_diffvae_ops.py create mode 100644 tests/diffusion/models/ltx2/test_ltx2_ops_platform.py create mode 100644 vllm_omni/diffusion/models/ltx2/ops/__init__.py create mode 100644 vllm_omni/diffusion/models/ltx2/ops/_numerics.py create mode 100644 vllm_omni/diffusion/models/ltx2/ops/fna.py create mode 100644 vllm_omni/diffusion/models/ltx2/ops/qk_rms_norm.py create mode 100644 vllm_omni/diffusion/models/ltx2/ops/residual_adaln.py create mode 100644 vllm_omni/diffusion/models/ltx2/ops/rope_tables.py create mode 100644 vllm_omni/diffusion/models/ltx2/ops/swiglu.py create mode 100644 vllm_omni/diffusion/models/ltx2/vae/__init__.py rename vllm_omni/diffusion/models/ltx2/{ltx2_diffusion_decoder.py => vae/decoder.py} (57%) rename vllm_omni/diffusion/models/ltx2/{ltx2_diffusion_decoder_distributed.py => vae/distributed.py} (99%) diff --git a/tests/diffusion/models/ltx2/test_ltx2_diffvae_attention.py b/tests/diffusion/models/ltx2/test_ltx2_diffvae_attention.py new file mode 100644 index 00000000000..6cdaf197d7c --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_diffvae_attention.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Native FNA against an independent, test-only explicit neighborhood mask.""" + +import pytest +import torch +from torch.nn.attention import SDPBackend, sdpa_kernel + +from vllm_omni.diffusion.models.ltx2.ops import resolve_ltx2_vae_operators +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion] + + +def _selected_fna(): + if not current_omni_platform.is_available(): + pytest.skip("No accelerator is available") + operators = resolve_ltx2_vae_operators(current_omni_platform.get_torch_device()) + if operators is None or operators.fna is None: + pytest.skip("LTX DiffVAE FNA is unavailable for this device or installation") + return operators.fna + + +def _load_fna_ops(): + return pytest.importorskip("vllm_omni.diffusion.models.ltx2.ops.fna", exc_type=ImportError) + + +def _pack(x, tile): + b, t, h, w, heads, dim = x.shape + tt, th, tw = tile + nt, nh, nw = (t + tt - 1) // tt, (h + th - 1) // th, (w + tw - 1) // tw + padded = x.new_zeros((b, nt * tt, nh * th, nw * tw, heads, dim)) + padded[:, :t, :h, :w] = x + # Outer and inner spatial tiles both flatten in W,H,T order. + return ( + padded.view(b, nt, tt, nh, th, nw, tw, heads, dim) + .permute(0, 5, 3, 1, 6, 4, 2, 7, 8) + .reshape(b, nt * nh * nw * tt * th * tw, heads, dim) + ) + + +def _unpack(x, shape): + t, h, w = shape + b, _, heads, dim = x.shape + nt, nh, nw = ((n + 3) // 4 for n in shape) + return ( + x.view(b, nw, nh, nt, 4, 4, 4, heads, dim) + .permute(0, 3, 6, 2, 5, 1, 4, 7, 8) + .reshape(b, nt * 4, nh * 4, nw * 4, heads, dim)[:, :t, :h, :w] + ) + + +@pytest.mark.parametrize("batch,shape", [(1, (11, 11, 11)), (2, (13, 12, 17))]) +@torch.inference_mode() +def test_fna_matches_explicit_inward_shift_mask(batch, shape): + fna = _selected_fna() + generator = torch.Generator(device="cuda").manual_seed(42) + q, k, v = [ + torch.randn(batch, *shape, 4, 64, device="cuda", dtype=torch.bfloat16, generator=generator) for _ in range(3) + ] + q = q * 0.125 + coordinates = torch.stack( + torch.meshgrid(*(torch.arange(n, device="cuda") for n in shape), indexing="ij"), dim=-1 + ).reshape(-1, 3) + start = (coordinates - 5).clamp_min(0).minimum(torch.tensor(shape, device="cuda") - 11) + mask = ((coordinates[None] >= start[:, None]) & (coordinates[None] < start[:, None] + 11)).all(-1) + with sdpa_kernel(SDPBackend.MATH): + expected = ( + torch.nn.functional.scaled_dot_product_attention( + *(x.flatten(1, 3).transpose(1, 2).float() for x in (q, k, v)), attn_mask=mask, scale=1.0 + ) + .transpose(1, 2) + .reshape(q.shape) + ) + actual = _unpack( + fna(_pack(q, (4, 4, 4)), _pack(k, (4, 4, 8)), _pack(v, (4, 4, 8)), shape=shape, window=(11, 11, 11)), + shape, + ).float() + assert actual.isfinite().all() + assert float((actual - expected).norm() / expected.norm()) < 0.005 + torch.testing.assert_close(actual, expected, atol=0.008, rtol=0.025) + + +def test_fna_rejects_cpu_without_fallback(): + fna_ops = _load_fna_ops() + x = torch.zeros(1, 64, 4, 64, dtype=torch.bfloat16) + with pytest.raises(ValueError, match="CUDA BF16"): + fna_ops.fna3d_tilelang(x, x, x, shape=(11, 11, 11), window=(11, 11, 11)) diff --git a/tests/diffusion/models/ltx2/test_ltx2_diffvae_ops.py b/tests/diffusion/models/ltx2/test_ltx2_diffvae_ops.py new file mode 100644 index 00000000000..a66a7f7c876 --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_diffvae_ops.py @@ -0,0 +1,409 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import subprocess +import sys + +import pytest +import torch + +from vllm_omni.diffusion.models.ltx2.ops import qk_rms_norm as qk_rms_norm_ops +from vllm_omni.diffusion.models.ltx2.ops import residual_adaln as residual_adaln_ops +from vllm_omni.diffusion.models.ltx2.ops import resolve_ltx2_vae_operators +from vllm_omni.diffusion.models.ltx2.ops import swiglu as swiglu_ops +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion] + +_DIM_SPLIT = (16, 24, 24) + + +def _selected_operators(): + if not current_omni_platform.is_available(): + pytest.skip("No accelerator is available") + operators = resolve_ltx2_vae_operators(current_omni_platform.get_torch_device()) + if operators is None: + pytest.skip("No LTX DiffVAE operators are available for this device") + return operators + + +def _selected_fna(): + operators = _selected_operators() + if operators.fna is None: + pytest.skip("LTX DiffVAE FNA is unavailable for this device or installation") + return operators.fna + + +def _load_fna_ops(): + return pytest.importorskip("vllm_omni.diffusion.models.ltx2.ops.fna", exc_type=ImportError) + + +def test_diffvae_model_and_tests_do_not_import_tilelang_backend() -> None: + # Use a fresh interpreter: reloading the model in this process would leave + # already-imported distributed subclasses pointing at the old base class. + result = subprocess.run( + [ + sys.executable, + "-c", + "import sys; " + "sys.modules['vllm_omni.diffusion.models.ltx2.ops.fna'] = None; " + "sys.modules['vllm.tilelang_utils'] = None; " + "from vllm_omni.diffusion.models.ltx2.vae.decoder import LTX2VideoDiffusionDecoderModel; " + "import tests.diffusion.models.ltx2.test_ltx2_diffvae_attention; " + "import tests.diffusion.models.ltx2.test_ltx2_diffvae_ops", + ], + capture_output=True, + text=True, + timeout=60, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +def test_diffvae_fna_tile_map_covers_exact_neighborhood() -> None: + fna_ops = _load_fna_ops() + problem = fna_ops.Problem(13, 14, 17, (11, 11, 11)) + indices, counts = fna_ops.build_kv_tile_map(problem) + q_grid_t, q_grid_h, _ = problem.q_grid + kv_grid_t, kv_grid_h, _ = problem.kv_grid + + for q_block in range(indices.shape[0]): + selected = set(indices[q_block, : counts[q_block]].tolist()) + q_outer_t = q_block % q_grid_t + q_outer_h = (q_block // q_grid_t) % q_grid_h + q_outer_w = q_block // (q_grid_t * q_grid_h) + for q_inner in range(fna_ops.Q_TOKENS): + query = ( + q_outer_t * 4 + q_inner % 4, + q_outer_h * 4 + (q_inner // 4) % 4, + q_outer_w * 4 + q_inner // 16, + ) + if any(index >= length for index, length in zip(query, problem.shape, strict=True)): + continue + starts = tuple( + fna_ops._window_start(index, window, length) + for index, window, length in zip(query, problem.window, problem.shape, strict=True) + ) + for key_t in range(starts[0], starts[0] + problem.window[0]): + for key_h in range(starts[1], starts[1] + problem.window[1]): + for key_w in range(starts[2], starts[2] + problem.window[2]): + key_block = ((key_w // 8) * kv_grid_h + key_h // 4) * kv_grid_t + key_t // 4 + assert key_block in selected + + +def test_diffvae_fna_metadata_uses_static_kernel_shapes() -> None: + fna_ops = _load_fna_ops() + for shape in ((11, 11, 11), (13, 14, 17), (66, 96, 128), (79, 192, 192)): + problem = fna_ops.Problem(*shape, (11, 11, 11)) + _, t_patterns, h_patterns, w_patterns = fna_ops.build_kv_tile_metadata(problem) + assert {table.numel() for table in (t_patterns, h_patterns, w_patterns)} == {fna_ops.PATTERN_TABLE_SIZE} + assert {limit for _, _, limit in fna_ops.build_kv_tile_buckets(problem)} <= set(fna_ops.KV_BUCKET_LIMITS) + for ids, metadata, limit in fna_ops.build_kv_tile_buckets(problem): + assert metadata.shape == (ids.numel(), limit) + + +def test_diffvae_fna_rejects_oversized_packed_tile_grid() -> None: + fna_ops = _load_fna_ops() + with pytest.raises(ValueError, match="15-bit metadata"): + fna_ops.Problem(128, 256, 256, (11, 11, 11)) + + +def test_diffvae_fna_propagates_launch_failures( + monkeypatch: pytest.MonkeyPatch, +) -> None: + fna_ops = _load_fna_ops() + fna = _selected_fna() + problem = fna_ops.Problem(11, 11, 11, (11, 11, 11)) + query = torch.randn(1, problem.q_length, fna_ops.HEADS, 64, device="cuda", dtype=torch.bfloat16) + key = torch.randn(1, problem.kv_length, fna_ops.HEADS, 64, device="cuda", dtype=torch.bfloat16) + value = torch.randn_like(key) + calls = 0 + + def fail_launch(*args, **kwargs): + nonlocal calls + calls += 1 + raise RuntimeError("injected launch failure") + + monkeypatch.setattr(fna_ops, "_fna3d_tilelang", fail_launch) + with torch.inference_mode(): + for _ in range(2): + with pytest.raises(RuntimeError, match="injected launch failure"): + fna( + query, + key, + value, + shape=problem.shape, + window=problem.window, + ) + assert calls == 2 + + +@pytest.fixture +def isolated_runtime_state(): + targets = ( + (qk_rms_norm_ops, "_FAILED_ROPE_KEYS"), + (qk_rms_norm_ops, "_VERIFIED_ROPE_KEYS"), + (swiglu_ops, "_FAILED_KEYS"), + (swiglu_ops, "_VERIFIED_KEYS"), + (residual_adaln_ops, "_FAILED_ADALN_KEYS"), + (residual_adaln_ops, "_VERIFIED_ADALN_KEYS"), + (residual_adaln_ops, "_FAILED_ADD_DEVICES"), + (residual_adaln_ops, "_VERIFIED_ADD_DEVICES"), + ) + saved = [set(getattr(module, name)) for module, name in targets] + for module, name in targets: + getattr(module, name).clear() + try: + yield + finally: + for (module, name), state in zip(targets, saved, strict=True): + target = getattr(module, name) + target.clear() + target.update(state) + + +@pytest.mark.parametrize( + ("dim", "rows"), + [ + (256, 257), + (256, 16385), + (256, 131075), + (512, 65539), + (1024, 32771), + (2048, 16387), + ], +) +def test_diffvae_swiglu_tiled_is_bit_exact(dim: int, rows: int, isolated_runtime_state) -> None: + from diffusers.models.autoencoders.ltx2_diffusion_decoder import LTX2VideoVaeSwiGLU + + operators = _selected_operators() + generator = torch.Generator(device="cuda").manual_seed(20260829) + hidden_states = torch.randn((1, 1, 1, rows, dim), device="cuda", dtype=torch.bfloat16, generator=generator) + gate_weight = torch.randn((4 * dim, dim), device="cuda", dtype=torch.bfloat16, generator=generator) + up_weight = torch.randn_like(gate_weight) + down_weight = torch.randn((dim, 4 * dim), device="cuda", dtype=torch.bfloat16, generator=generator) + with torch.device("meta"): + reference = LTX2VideoVaeSwiGLU(dim, 4 * dim) + reference.w_gate.weight = torch.nn.Parameter(gate_weight, requires_grad=False) + reference.w_up.weight = torch.nn.Parameter(up_weight, requires_grad=False) + reference.w_down.weight = torch.nn.Parameter(down_weight, requires_grad=False) + + with torch.inference_mode(): + expected = reference(hidden_states) + actual = operators.swiglu(hidden_states, gate_weight, up_weight, down_weight) + + assert actual is not None + assert torch.equal(actual, expected) + + +@pytest.mark.parametrize( + ("dim", "expected_rows"), + [(256, 131072), (512, 65536), (1024, 32768), (2048, 16384)], +) +def test_diffvae_swiglu_tile_rows_keep_a_constant_hidden_budget(dim: int, expected_rows: int) -> None: + assert swiglu_ops._tile_rows(1_000_000, 4 * dim) == expected_rows + + +def test_diffvae_swiglu_tiled_permanently_falls_back_after_failure( + monkeypatch: pytest.MonkeyPatch, + isolated_runtime_state, +) -> None: + operators = _selected_operators() + hidden_states = torch.randn(1, 1, 1, 17, 256, device="cuda", dtype=torch.bfloat16) + gate_weight = torch.randn(1024, 256, device="cuda", dtype=torch.bfloat16) + up_weight = torch.randn_like(gate_weight) + down_weight = torch.randn(256, 1024, device="cuda", dtype=torch.bfloat16) + calls = 0 + + def fail_launch(*args, **kwargs): + nonlocal calls + calls += 1 + raise RuntimeError("injected launch failure") + + monkeypatch.setattr(swiglu_ops, "_launch", fail_launch) + with torch.inference_mode(): + assert operators.swiglu(hidden_states, gate_weight, up_weight, down_weight) is None + assert operators.swiglu(hidden_states, gate_weight, up_weight, down_weight) is None + + assert calls == 1 + + +@pytest.mark.parametrize("residual_count", [1, 2]) +def test_diffvae_residual_rms_norm_modulate_is_bit_exact( + residual_count: int, + isolated_runtime_state, +) -> None: + operators = _selected_operators() + generator = torch.Generator(device="cuda").manual_seed(20260829) + shape = (2, 3, 7, 7, 256) + x = torch.randn(shape, device="cuda", dtype=torch.bfloat16, generator=generator) + residual_a = torch.randn_like(x) + residual_b = torch.randn_like(x) if residual_count == 2 else None + norm_weight = torch.randn((256,), device="cuda", dtype=torch.bfloat16, generator=generator) + scale = torch.randn((2, 1, 1, 1, 256), device="cuda", dtype=torch.bfloat16, generator=generator) + shift = torch.randn_like(scale) + hidden_states = x + residual_a + if residual_b is not None: + hidden_states = hidden_states + residual_b + expected = torch.nn.functional.rms_norm(hidden_states, (256,), norm_weight, eps=1e-6) * (1 + scale) + shift + + with torch.inference_mode(): + actual = operators.residual_norm( + x, + residual_a, + residual_b, + norm_weight, + scale, + shift, + 1e-6, + ) + + assert actual is not None + assert torch.equal(actual, expected) + + +def test_diffvae_residual_add3_is_bit_exact(isolated_runtime_state) -> None: + operators = _selected_operators() + generator = torch.Generator(device="cuda").manual_seed(20260829) + tensors = [ + torch.randn((2, 3, 7, 7, 256), device="cuda", dtype=torch.bfloat16, generator=generator) for _ in range(4) + ] + expected = tensors[0] + tensors[1] + expected = expected + tensors[2] + expected = expected + tensors[3] + + with torch.inference_mode(): + actual = operators.residual_add(*tensors) + + assert actual is not None + assert torch.equal(actual, expected) + + +def test_diffvae_residual_adaln_permanently_falls_back_after_failure( + monkeypatch: pytest.MonkeyPatch, + isolated_runtime_state, +) -> None: + operators = _selected_operators() + shape = (1, 3, 7, 7, 256) + x = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + residual = torch.randn_like(x) + weight = torch.randn(256, device="cuda", dtype=torch.bfloat16) + scale = torch.randn(1, 1, 1, 1, 256, device="cuda", dtype=torch.bfloat16) + shift = torch.randn_like(scale) + calls = 0 + + def fail_launch(*args, **kwargs): + nonlocal calls + calls += 1 + raise RuntimeError("injected launch failure") + + monkeypatch.setattr(residual_adaln_ops, "_launch_adaln", fail_launch) + with torch.inference_mode(): + for _ in range(2): + assert operators.residual_norm(x, residual, None, weight, scale, shift, 1e-6) is None + + assert calls == 1 + + +def _tables( + frames: int, + height: int, + width: int, + dim_split: tuple[int, int, int] = _DIM_SPLIT, + base: float = 10000.0, +) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: + tables = [] + for length, dim in zip((frames, height, width), dim_split, strict=True): + exponents = torch.arange(0, dim, 2, dtype=torch.float64, device="cuda") / dim + inv_freqs = (1.0 / base**exponents).to(torch.float32) + positions = torch.arange(length, dtype=torch.float32, device="cuda") + angles = positions[:, None] * inv_freqs[None, :] + tables.append((angles.cos(), angles.sin())) + return tuple(tables) + + +def _apply_rope_reference( + hidden_states: torch.Tensor, + tables: tuple[tuple[torch.Tensor, torch.Tensor], ...], + dim_split: tuple[int, int, int] = _DIM_SPLIT, +) -> torch.Tensor: + outputs = [] + offset = 0 + for axis, (dim, (cos, sin)) in enumerate(zip(dim_split, tables, strict=True), 1): + chunk = hidden_states[..., offset : offset + dim] + pairs = chunk.reshape(*chunk.shape[:-1], dim // 2, 2) + even = pairs[..., 0].float() + odd = pairs[..., 1].float() + shape = [1, 1, 1, 1, 1, dim // 2] + shape[axis] = cos.shape[0] + cos = cos.reshape(shape) + sin = sin.reshape(shape) + rotated = torch.stack([even * cos - odd * sin, even * sin + odd * cos], dim=-1) + outputs.append(rotated.reshape(chunk.shape).to(hidden_states.dtype)) + offset += dim + return torch.cat(outputs, dim=-1) + + +@pytest.mark.parametrize( + ("frames", "height", "width", "heads"), + [(3, 7, 7, 1), (31, 136, 192, 4)], +) +def test_diffvae_qk_rms_norm_scale_rope_3d_is_bit_exact( + frames: int, + height: int, + width: int, + heads: int, + isolated_runtime_state, +) -> None: + operators = _selected_operators() + generator = torch.Generator(device="cuda").manual_seed(20260829) + shape = (1, frames, height, width, heads, 64) + query = torch.randn(shape, device="cuda", dtype=torch.bfloat16, generator=generator) + key = torch.randn_like(query) + query_weight = torch.randn((64,), device="cuda", dtype=torch.bfloat16, generator=generator) + key_weight = torch.randn((64,), device="cuda", dtype=torch.bfloat16, generator=generator) + tables = _tables(frames, height, width) + expected = ( + _apply_rope_reference(torch.nn.functional.rms_norm(query, (64,), query_weight, 1e-6) * 0.125, tables), + _apply_rope_reference(torch.nn.functional.rms_norm(key, (64,), key_weight, 1e-6), tables), + ) + + with torch.inference_mode(): + actual = operators.qk_norm_rope( + query, + key, + query_weight, + key_weight, + 1e-6, + 0.125, + _DIM_SPLIT, + 10000.0, + ) + + assert actual is not None + assert torch.equal(actual[0], expected[0]) + assert torch.equal(actual[1], expected[1]) + + +def test_diffvae_qk_rms_norm_scale_rope_3d_permanently_falls_back_after_failure( + monkeypatch: pytest.MonkeyPatch, + isolated_runtime_state, +) -> None: + operators = _selected_operators() + query = torch.randn(1, 3, 7, 7, 1, 64, device="cuda", dtype=torch.bfloat16) + key = torch.randn_like(query) + weight = torch.ones(64, device="cuda", dtype=torch.bfloat16) + calls = 0 + + def fail_launch(*args, **kwargs): + nonlocal calls + calls += 1 + raise RuntimeError("injected launch failure") + + monkeypatch.setattr(qk_rms_norm_ops, "_launch_combined", fail_launch) + with torch.inference_mode(): + for _ in range(2): + assert operators.qk_norm_rope(query, key, weight, weight, 1e-6, 0.125, _DIM_SPLIT, 10000.0) is None + + assert calls == 1 diff --git a/tests/diffusion/models/ltx2/test_ltx2_ops_platform.py b/tests/diffusion/models/ltx2/test_ltx2_ops_platform.py new file mode 100644 index 00000000000..6ffff3136f4 --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_ops_platform.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import sys +from dataclasses import dataclass + +import pytest +import torch +from vllm.platforms.interface import DeviceCapability + +from vllm_omni.diffusion.models.ltx2 import ops as ltx2_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@dataclass(frozen=True) +class _TensorProbe: + is_cuda: bool + device: torch.device + + +@pytest.fixture(autouse=True) +def clear_platform_cache(): + ltx2_platform._cuda_compute_capability.cache_clear() + try: + yield + finally: + ltx2_platform._cuda_compute_capability.cache_clear() + + +@pytest.mark.parametrize( + ("capability", "triton", "tilelang", "fusion_expected", "fna_expected"), + [ + (80, True, True, False, False), + (90, True, True, True, True), + (90, True, False, True, False), + (90, False, True, False, False), + (100, True, True, False, False), + (103, True, True, False, False), + (110, True, True, False, False), + (120, True, True, False, False), + (121, True, True, False, False), + ], +) +def test_production_resolver_and_tensor_eligibility_agree( + monkeypatch, capability, triton, tilelang, fusion_expected, fna_expected +): + monkeypatch.setattr(ltx2_platform, "HAS_TRITON", triton) + monkeypatch.setattr(ltx2_platform, "has_tilelang", lambda: tilelang) + monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_cuda", lambda: True) + monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_available", lambda: True) + monkeypatch.setattr( + ltx2_platform.current_omni_platform, + "get_device_capability", + lambda device_id=0: DeviceCapability(major=capability // 10, minor=capability % 10), + ) + # Resolution must be lazy even when FNA is selected. + monkeypatch.setitem(sys.modules, "vllm_omni.diffusion.models.ltx2.ops.fna", None) + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: False) + device = torch.device("cuda:0") + operators = ltx2_platform.resolve_ltx2_vae_operators(device) + assert (operators is not None) is fusion_expected + assert (operators is not None and operators.fna is not None) is fna_expected + if operators is not None: + from vllm_omni.diffusion.models.ltx2.ops import qk_rms_norm, residual_adaln, swiglu + + assert operators.qk_norm_rope is qk_rms_norm.try_qk_rms_norm_scale_rope_3d_exact + assert operators.residual_norm is residual_adaln.try_residual_rms_norm_modulate_exact + assert operators.residual_add is residual_adaln.try_residual_add3_exact + assert operators.swiglu is swiglu.try_swiglu_tiled_exact + + tensor = _TensorProbe(is_cuda=True, device=device) + with torch.inference_mode(): + assert ltx2_platform.is_ltx2_fusion_eligible(tensor) is fusion_expected + assert ltx2_platform.is_ltx2_fna_eligible(tensor) is fna_expected + + +@pytest.mark.parametrize( + ("device", "is_cuda", "available"), + [("cpu", True, True), ("cuda:0", False, True), ("cuda:0", True, False)], +) +def test_resolver_rejects_unsupported_platform_before_capability_lookup(monkeypatch, device, is_cuda, available): + monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_cuda", lambda: is_cuda) + monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_available", lambda: available) + + def unexpected_lookup(_index): + raise AssertionError("Unsupported platforms must not query CUDA capability") + + monkeypatch.setattr(ltx2_platform, "_cuda_compute_capability", unexpected_lookup) + assert ltx2_platform.resolve_ltx2_vae_operators(torch.device(device)) is None + + +@pytest.mark.parametrize(("compiling", "grad_enabled"), [(True, False), (False, True)]) +def test_runtime_guards_bypass_operator_resolution(monkeypatch, compiling, grad_enabled): + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: compiling) + monkeypatch.setattr(torch, "is_grad_enabled", lambda: grad_enabled) + + def unexpected_resolve(_device): + raise AssertionError("Ineligible runtime must not select eager operators") + + monkeypatch.setattr(ltx2_platform, "resolve_ltx2_vae_operators", unexpected_resolve) + tensor = _TensorProbe(is_cuda=True, device=torch.device("cuda:0")) + assert not ltx2_platform.is_ltx2_fusion_eligible(tensor) + assert not ltx2_platform.is_ltx2_fna_eligible(tensor) + + +def test_device_capability_is_cached(monkeypatch): + calls = 0 + monkeypatch.setattr(ltx2_platform, "HAS_TRITON", True) + + def get_device_capability(device_id=0): + nonlocal calls + calls += 1 + return DeviceCapability(major=9, minor=0) + + monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_cuda", lambda: True) + monkeypatch.setattr(ltx2_platform.current_omni_platform, "is_available", lambda: True) + monkeypatch.setattr(ltx2_platform.current_omni_platform, "get_device_capability", get_device_capability) + assert ltx2_platform._cuda_compute_capability(0) == 90 + assert ltx2_platform._cuda_compute_capability(0) == 90 + assert calls == 1 diff --git a/tests/diffusion/models/ltx2/test_ltx2_vae.py b/tests/diffusion/models/ltx2/test_ltx2_vae.py index b9bc52d9452..e387eb7a4ce 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_vae.py +++ b/tests/diffusion/models/ltx2/test_ltx2_vae.py @@ -3,6 +3,8 @@ """Unit tests for LTX video VAE tiling and distributed decode behavior.""" +from collections.abc import Callable +from dataclasses import dataclass from types import SimpleNamespace from typing import Any @@ -35,6 +37,22 @@ def test_ltx_vocoder_deterministic_context(monkeypatch, initial_deterministic, f assert {name: getattr(torch.backends.cudnn, name) for name in settings} == settings +@dataclass(frozen=True) +class _OperatorSetProbe: + fna: Callable[..., torch.Tensor] | None + + +@dataclass +class _AttentionProbe: + kernel_size: tuple[int, int, int] + head_dim: int + to_out: tuple[torch.nn.Module, ...] + projected: torch.Tensor + + def project_qkv(self, _hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return self.projected, self.projected, self.projected + + def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch): import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime @@ -61,9 +79,164 @@ def forward(self, value): class TestLTXDiffusionDecoder: + def test_diffusion_decoder_reuses_diffusers_with_scoped_overrides(self, monkeypatch): + from diffusers.models.autoencoders.ltx2_diffusion_decoder import ( + LTX2VideoDiffusionDecoderModel as DiffusersModel, + ) + + from vllm_omni.diffusion.models.ltx2.vae.decoder import ( + LTX2VideoDiffusionDecoderModel, + LTX2VideoVaeDiffusionNABlock, + LTX2VideoVaeNeighborhoodAttention, + LTX2VideoVaeSwiGLU, + ) + + original_init = DiffusersModel.__init__ + original_parameters = {} + + def capture_parameters(self, *args, **kwargs): + original_init(self, *args, **kwargs) + original_parameters.update(self.named_parameters()) + + with torch.device("meta"): + reference = DiffusersModel() + monkeypatch.setattr(DiffusersModel, "__init__", capture_parameters) + model = LTX2VideoDiffusionDecoderModel() + assert isinstance(model, DiffusersModel) + assert all(parameter is original_parameters[name] for name, parameter in model.named_parameters()) + assert {name: tensor.shape for name, tensor in model.state_dict().items()} == { + name: tensor.shape for name, tensor in reference.state_dict().items() + } + + modules = tuple(model.decoder.modules()) + expected_blocks = sum(model.config.decoder_stage_depths) + assert sum(isinstance(module, LTX2VideoVaeNeighborhoodAttention) for module in modules) == expected_blocks + assert sum(isinstance(module, LTX2VideoVaeSwiGLU) for module in modules) == expected_blocks + assert ( + sum(isinstance(module, LTX2VideoVaeDiffusionNABlock) for module in modules) + == model.config.decoder_stage_depths[-1] + ) + + assert type(model.decoder).__module__.startswith("vllm_omni.") + assert model.tiled_decode.__func__ is DiffusersModel.tiled_decode + assert type(reference.decoder) is type(model.decoder).__bases__[0] + + def test_natten_processor_routes_stage5_through_tilelang_fna(self, monkeypatch): + from vllm_omni.diffusion.models.ltx2.vae import decoder as diffvae_modules + + processor = object.__new__(diffvae_modules.LTX2VideoVaeNeighborhoodNattenProcessor) + processor.backend = None + fallback_calls = [] + processor._na3d = lambda *args, **kwargs: fallback_calls.append((args, kwargs)) + + hidden_states = torch.randn(1, 11, 11, 11, 256, dtype=torch.bfloat16) + projected = hidden_states.view(1, 11, 11, 11, 4, 64) + + class FakeAttention: + kernel_size = (11, 11, 11) + head_dim = 64 + to_out = (torch.nn.Identity(),) + + @staticmethod + def project_qkv(_hidden_states): + return projected, projected, projected + + def token_permute(value, _tile, *, flip_tiled_dims): + assert flip_tiled_dims + return value, value.shape, None + + def token_unpermute(value, *, token_layout_shape, tile_shape, flip_tiled_dims): + assert token_layout_shape == projected.shape + assert tile_shape == (4, 4, 4) + assert flip_tiled_dims + return value + + fast_path_calls = [] + + def fast_path(query, key, value, *, shape, window): + fast_path_calls.append((query, key, value, shape, window)) + return query + + monkeypatch.setattr(diffvae_modules, "is_ltx2_fna_eligible", lambda _tensor: True) + monkeypatch.setattr( + diffvae_modules, "resolve_ltx2_vae_operators", lambda _device: _OperatorSetProbe(fna=fast_path) + ) + monkeypatch.setattr( + diffvae_modules, + "_load_natten_token_permutation", + lambda: (token_permute, token_unpermute), + ) + + output = processor(FakeAttention(), hidden_states) + + assert torch.equal(output, hidden_states) + assert len(fast_path_calls) == 1 + assert fast_path_calls[0][3:] == ((11, 11, 11), (11, 11, 11)) + assert not fallback_calls + + @pytest.mark.parametrize("reason", ["device", "head_dim", "channels", "dtype", "tilelang_missing", "kernel"]) + def test_natten_processor_keeps_natten_when_fast_path_does_not_apply(self, monkeypatch, reason): + from vllm_omni.diffusion.models.ltx2.vae import decoder as ops + + processor = object.__new__(ops.LTX2VideoVaeNeighborhoodNattenProcessor) + processor.backend = None + channels = 128 if reason == "channels" else 256 + head_dim = 32 if reason == "head_dim" else 64 + x = torch.zeros(1, 11, 11, 11, channels, dtype=torch.float32 if reason == "dtype" else torch.bfloat16) + projected = x.view(1, 11, 11, 11, channels // head_dim, head_dim) + attn = _AttentionProbe( + kernel_size=(3, 5, 5) if reason == "kernel" else (11, 11, 11), + head_dim=head_dim, + to_out=(torch.nn.Identity(),), + projected=projected, + ) + calls = [] + + def natten(q, k, v, **kwargs): + calls.append(kwargs) + return q + + def unexpected_native_load(*_args, **_kwargs): + raise AssertionError("Native FNA must not be loaded for an unsupported configuration") + + processor._na3d = natten + monkeypatch.setattr(ops, "is_ltx2_fna_eligible", lambda _tensor: reason != "device") + monkeypatch.setattr( + ops, + "resolve_ltx2_vae_operators", + lambda _device: _OperatorSetProbe(fna=None if reason == "tilelang_missing" else unexpected_native_load), + ) + assert torch.equal(processor(attn, x), x) + assert calls == [{"kernel_size": attn.kernel_size, "scale": 1.0, "backend": None}] + + def test_diffusion_decoder_noneligible_block_uses_diffusers_forward(self, monkeypatch): + from vllm_omni.diffusion.models.ltx2.vae import decoder as diffvae_modules + + block = object.__new__(diffvae_modules.LTX2VideoVaeDiffusionNABlock) + torch.nn.Module.__init__(block) + expected = torch.randn(1, 1, 1, 1, 1) + calls = [] + + def upstream_forward(_self, hidden_states, latent_context, modulation, block_mask): + assert block_mask is None + calls.append((hidden_states, latent_context, modulation)) + return expected + + monkeypatch.setattr( + diffvae_modules.diffusers_decoder.LTX2VideoVaeDiffusionNABlock, + "forward", + upstream_forward, + ) + hidden_states = torch.randn(1, 1, 1, 1, 1) + latent_context = torch.randn_like(hidden_states) + modulation = (torch.randn_like(hidden_states),) + + assert block.forward(hidden_states, latent_context, modulation) is expected + assert calls == [(hidden_states, latent_context, modulation)] + @pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width"]) def test_distributed_diffusion_decoder_rejects_non_tile_parallel_modes(self, mode): - from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import ( + from vllm_omni.diffusion.models.ltx2.vae.distributed import ( DistributedLTX2VideoDiffusionDecoderModel, ) @@ -75,7 +248,7 @@ def test_distributed_diffusion_decoder_rejects_non_tile_parallel_modes(self, mod model.set_parallel_size(2, mode=mode) def test_distributed_diffusion_decoder_accepts_tile_parallel_mode(self): - from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import ( + from vllm_omni.diffusion.models.ltx2.vae.distributed import ( DistributedLTX2VideoDiffusionDecoderModel, ) @@ -91,7 +264,7 @@ def test_distributed_diffusion_decoder_accepts_tile_parallel_mode(self): assert calls == [(4, "tile")] def test_short_clip_keeps_stage5_temporal_context_then_crops_output(self): - from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import ( + from vllm_omni.diffusion.models.ltx2.vae.decoder import ( LTX2VideoDiffusionDecoder3d, LTX2VideoDiffusionDecoderModel, ) @@ -166,7 +339,7 @@ def test_conv_vae_opt_in_rejects_non_boolean_model_extra(self): _ltx2_use_diffusion_decoder(SimpleNamespace(extras={"ltx2_use_conv_vae": "true"}), "2.5") def test_native_diffusion_decoder_conversion_splits_qkv_and_folds_gates(self): - from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import ( + from vllm_omni.diffusion.models.ltx2.vae.decoder import ( convert_ltx25_native_diffusion_decoder_state_dict, ) @@ -213,7 +386,7 @@ def test_native_diffusion_decoder_conversion_splits_qkv_and_folds_gates(self): assert not any("type_emb" in key or "coarse" in key or "gate_msa" in key for key in converted) def test_native_diffusion_decoder_conversion_rejects_invalid_fused_qkv(self): - from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import ( + from vllm_omni.diffusion.models.ltx2.vae.decoder import ( convert_ltx25_native_diffusion_decoder_state_dict, ) @@ -411,8 +584,8 @@ def decode(self, *_args, **_kwargs): assert output.output[1].numel() == 0 def test_diffusion_decoder_patch_parallel_size_one_uses_native_tiling(self, monkeypatch): - from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel - from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import ( + from vllm_omni.diffusion.models.ltx2.vae.decoder import LTX2VideoDiffusionDecoderModel + from vllm_omni.diffusion.models.ltx2.vae.distributed import ( DistributedLTX2VideoDiffusionDecoderModel, ) @@ -447,7 +620,7 @@ def native_tiled_decode(self, z, generator=None, num_inference_steps=None): def test_distributed_diffusion_tiles_preserve_serial_noise_order(self, monkeypatch): from diffusers.utils.torch_utils import randn_tensor - from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import ( + from vllm_omni.diffusion.models.ltx2.vae.distributed import ( DistributedLTX2VideoDiffusionDecoderModel, ) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py index 72e63a22f3b..f1dcc38bb2d 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_components.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_components.py @@ -15,7 +15,6 @@ import torch from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler -from diffusers.models.autoencoders.ltx2_diffusion_decoder import LTX2VideoVaeNeighborhoodNattenProcessor from diffusers.pipelines.ltx2 import LTX2TextConnectors from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder @@ -36,12 +35,6 @@ if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig -from .ltx2_diffusion_decoder import ( - LTX25_NATIVE_ARTIFACT_REVISION, - LTX25_NATIVE_DIFFUSION_DECODER_FILENAME, - LTX25_NATIVE_DIFFUSION_DECODER_REPO_ID, -) -from .ltx2_diffusion_decoder_distributed import DistributedLTX2VideoDiffusionDecoderModel from .ltx2_request import LTXCheckpointKind, validate_ltx_checkpoint from .ltx2_transformer import ( LTX2VideoTransformer3DModel, @@ -49,6 +42,13 @@ apply_split_rotary_emb, to_ltx_padding_mask, ) +from .vae.decoder import ( + LTX25_NATIVE_ARTIFACT_REVISION, + LTX25_NATIVE_DIFFUSION_DECODER_FILENAME, + LTX25_NATIVE_DIFFUSION_DECODER_REPO_ID, + LTX2VideoVaeNeighborhoodNattenProcessor, +) +from .vae.distributed import DistributedLTX2VideoDiffusionDecoderModel try: from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE diff --git a/vllm_omni/diffusion/models/ltx2/ops/__init__.py b/vllm_omni/diffusion/models/ltx2/ops/__init__.py new file mode 100644 index 00000000000..93f2ff67291 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ops/__init__.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Shared operator selection for LTX-2 eager kernels and their tests.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from functools import cache + +import torch +from vllm.triton_utils import HAS_TRITON +from vllm.utils.import_utils import has_tilelang + +from vllm_omni.platforms import current_omni_platform + +# Pointwise fusions are launch/memory-traffic reductions whose first execution +# for every runtime key is checked bit-for-bit; failures use the ordinary model +# math. Both paths remain SM90-only until other architectures are validated. +# The FNA schedule is performance-tuned for SM90 and has no silent fallback. +_CUDA_COMPUTE_CAPABILITIES = frozenset({90}) + + +@dataclass(frozen=True) +class LTX2VAEOperatorSet: + qk_norm_rope: Callable[..., tuple[torch.Tensor, torch.Tensor] | None] + residual_norm: Callable[..., torch.Tensor | None] + residual_add: Callable[..., torch.Tensor | None] + swiglu: Callable[..., torch.Tensor | None] + fna: Callable[..., torch.Tensor] | None + + +def _fna3d_tilelang(*args, **kwargs): + # Selecting an operator set must not import the optional compiler backend. + from .fna import fna3d_tilelang + + return fna3d_tilelang(*args, **kwargs) + + +@cache +def _operator_set(with_fna: bool) -> LTX2VAEOperatorSet: + from .qk_rms_norm import try_qk_rms_norm_scale_rope_3d_exact + from .residual_adaln import try_residual_add3_exact, try_residual_rms_norm_modulate_exact + from .swiglu import try_swiglu_tiled_exact + + return LTX2VAEOperatorSet( + qk_norm_rope=try_qk_rms_norm_scale_rope_3d_exact, + residual_norm=try_residual_rms_norm_modulate_exact, + residual_add=try_residual_add3_exact, + swiglu=try_swiglu_tiled_exact, + fna=_fna3d_tilelang if with_fna else None, + ) + + +@cache +def _cuda_compute_capability(device_index: int) -> int | None: + if not HAS_TRITON or not current_omni_platform.is_cuda() or not current_omni_platform.is_available(): + return None + capability = current_omni_platform.get_device_capability(device_id=device_index) + return capability.to_int() if capability is not None else None + + +def resolve_ltx2_vae_operators(device: torch.device) -> LTX2VAEOperatorSet | None: + """Select the production operators, or return None for unsupported devices.""" + if device.type != "cuda" or not current_omni_platform.is_cuda() or not current_omni_platform.is_available(): + return None + device_index = device.index if device.index is not None else torch.accelerator.current_device_index() + if _cuda_compute_capability(int(device_index)) not in _CUDA_COMPUTE_CAPABILITIES: + return None + return _operator_set(has_tilelang()) + + +def _eligible_operators(tensor: torch.Tensor) -> LTX2VAEOperatorSet | None: + if not tensor.is_cuda or torch.compiler.is_compiling() or torch.is_grad_enabled(): + return None + return resolve_ltx2_vae_operators(tensor.device) + + +def is_ltx2_fusion_eligible(tensor: torch.Tensor) -> bool: + """Allow self-verifying pointwise fusions on validated SM90 devices.""" + + return _eligible_operators(tensor) is not None + + +def is_ltx2_fna_eligible(tensor: torch.Tensor) -> bool: + """Keep the strict, Hopper-tuned TileLang FNA schedule on SM90.""" + + operators = _eligible_operators(tensor) + return operators is not None and operators.fna is not None + + +__all__ = ["LTX2VAEOperatorSet", "resolve_ltx2_vae_operators", "is_ltx2_fna_eligible", "is_ltx2_fusion_eligible"] diff --git a/vllm_omni/diffusion/models/ltx2/ops/_numerics.py b/vllm_omni/diffusion/models/ltx2/ops/_numerics.py new file mode 100644 index 00000000000..bdc22b14244 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ops/_numerics.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Numerical primitives shared by exact LTX-2 Triton kernels.""" + +from vllm.triton_utils import tl, triton + + +@triton.jit +def round_bf16_to_fp32(value): + """RNE-round FP32 to BF16 precision while retaining an FP32 register.""" + + bits = value.to(tl.int32, bitcast=True) + rounding_bias = 0x7FFF + ((bits >> 16) & 1) + rounded_bits = (bits + rounding_bias) & -65536 + return rounded_bits.to(tl.float32, bitcast=True) + + +@triton.jit +def add_rn_f32(x, y): + return tl.inline_asm_elementwise( + asm="add.rn.f32 $0, $1, $2;", + constraints="=f,f,f", + args=[x, y], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def mul_rn_f32(x, y): + return tl.inline_asm_elementwise( + asm="mul.rn.f32 $0, $1, $2;", + constraints="=f,f,f", + args=[x, y], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def fma_rn_f32(x, y, accumulator): + return tl.inline_asm_elementwise( + asm="fma.rn.f32 $0, $1, $2, $3;", + constraints="=f,f,f,f", + args=[x, y, accumulator], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def rsqrt_approx_f32(x): + return tl.inline_asm_elementwise( + asm="rsqrt.approx.f32 $0, $1;", + constraints="=f,f", + args=[x], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def shfl_down_f32(value, delta: tl.constexpr): + return tl.inline_asm_elementwise( + asm="shfl.sync.down.b32 $0, $1, $2, 0x1f, 0xffffffff;", + constraints="=f,f,n", + args=[value, delta], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +__all__ = [ + "add_rn_f32", + "fma_rn_f32", + "mul_rn_f32", + "round_bf16_to_fp32", + "rsqrt_approx_f32", + "shfl_down_f32", +] diff --git a/vllm_omni/diffusion/models/ltx2/ops/fna.py b/vllm_omni/diffusion/models/ltx2/ops/fna.py new file mode 100644 index 00000000000..d42477d9ac2 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ops/fna.py @@ -0,0 +1,525 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""TileLang FNA for the LTX-2.5 DiffVAE stage-5 schedule. + +This intentionally reuses NATTEN's token permutation. It only replaces the +middle BF16 attention kernel for Q tiles (4, 4, 4), KV tiles (4, 4, 8), +head_dim=64, stride=dilation=1, and scale=1.0. +""" + +import functools +import math +from collections import OrderedDict +from dataclasses import dataclass + +import torch +from vllm.tilelang_utils import T, tilelang + +Q_TILE = (4, 4, 4) +KV_TILE = (4, 4, 8) +Q_TOKENS = math.prod(Q_TILE) +KV_TOKENS = math.prod(KV_TILE) +HEAD_DIM = 64 +HEADS = 4 +LOG2_E = 1.4426950408889634 +KV_BUCKET_LIMITS = (32, 48, 60, 75) +PATTERN_TABLE_SIZE = 16 +KV_TILE_ID_BITS = 15 +SHAPE_CACHE_SIZE = 16 + + +@dataclass(frozen=True) +class Problem: + frames: int + height: int + width: int + window: tuple[int, int, int] + + def __post_init__(self) -> None: + if min(self.frames, self.height, self.width) <= 0: + raise ValueError(f"invalid FNA problem: {self}") + if len(self.window) != 3 or any(x <= 0 for x in self.window): + raise ValueError(f"invalid FNA window: {self.window}") + if any(w > n for w, n in zip(self.window, self.shape)): + raise ValueError(f"window {self.window} exceeds shape {self.shape}") + if math.prod(self.kv_grid) > 1 << KV_TILE_ID_BITS: + raise ValueError(f"KV tile grid exceeds {KV_TILE_ID_BITS}-bit metadata encoding: {self.kv_grid}") + + @property + def shape(self) -> tuple[int, int, int]: + return self.frames, self.height, self.width + + @property + def q_padded(self) -> tuple[int, int, int]: + return tuple(_round_up(n, tile) for n, tile in zip(self.shape, Q_TILE)) + + @property + def kv_padded(self) -> tuple[int, int, int]: + return tuple(_round_up(n, tile) for n, tile in zip(self.shape, KV_TILE)) + + @property + def q_grid(self) -> tuple[int, int, int]: + return tuple(n // tile for n, tile in zip(self.q_padded, Q_TILE)) + + @property + def kv_grid(self) -> tuple[int, int, int]: + return tuple(n // tile for n, tile in zip(self.kv_padded, KV_TILE)) + + @property + def q_length(self) -> int: + return math.prod(self.q_padded) + + @property + def kv_length(self) -> int: + return math.prod(self.kv_padded) + + +def _round_up(value: int, multiple: int) -> int: + return (value + multiple - 1) // multiple * multiple + + +def _window_start(index: int, window: int, length: int) -> int: + return min(max(index - window // 2, 0), length - window) + + +@functools.lru_cache(maxsize=32) +def _dimension_pattern_table( + length: int, + q_tile_size: int, + kv_tile_size: int, + window: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Map (Q tile, KV tile) pairs to a small table of packed 1D masks.""" + q_tiles = _round_up(length, q_tile_size) // q_tile_size + kv_tiles = _round_up(length, kv_tile_size) // kv_tile_size + zero = (0,) * q_tile_size + patterns: list[tuple[int, ...]] = [zero] + pattern_ids = {zero: 0} + ids = torch.zeros((q_tiles, kv_tiles), dtype=torch.int64) + + for q_tile in range(q_tiles): + for kv_tile in range(kv_tiles): + rows: list[int] = [] + for q_inner in range(q_tile_size): + q = q_tile * q_tile_size + q_inner + bits = 0 + if q < length: + start = _window_start(q, window, length) + for kv_inner in range(kv_tile_size): + kv = kv_tile * kv_tile_size + kv_inner + if kv < length and start <= kv < start + window: + bits |= 1 << kv_inner + rows.append(bits) + signature = tuple(rows) + pattern_id = pattern_ids.get(signature) + if pattern_id is None: + pattern_id = len(patterns) + pattern_ids[signature] = pattern_id + patterns.append(signature) + ids[q_tile, kv_tile] = pattern_id + + if len(patterns) > PATTERN_TABLE_SIZE: + raise AssertionError(f"4-bit pattern id overflow for length={length}: {len(patterns)}") + packed = [] + for pattern in patterns: + value = sum(bits << (q * kv_tile_size) for q, bits in enumerate(pattern)) + # Store uint32 bit patterns in an int32 tensor. Arithmetic right shift + # followed by `& 1` still extracts the desired bit when bit 31 is set. + packed.append(value if value < (1 << 31) else value - (1 << 32)) + # The four-bit pattern ids address at most 16 entries. Pad to that fixed + # size so T/H/W-specific table lengths cannot specialize the GPU kernel. + packed.extend([0] * (PATTERN_TABLE_SIZE - len(packed))) + return ids, torch.tensor(packed, dtype=torch.int32) + + +@functools.lru_cache(maxsize=SHAPE_CACHE_SIZE) +def build_kv_tile_map(problem: Problem) -> tuple[torch.Tensor, torch.Tensor]: + """Return packed KV tile ids and valid counts for each permuted Q tile.""" + q_grid_t, q_grid_h, q_grid_w = problem.q_grid + kv_grid_t, kv_grid_h, _ = problem.kv_grid + rows: list[list[int]] = [] + + # Token permutation flattens outer tile modes as (W, H, T), with T + # changing fastest. Keep KV tiles in the same increasing linear order. + for q_w_tile in range(q_grid_w): + for q_h_tile in range(q_grid_h): + for q_t_tile in range(q_grid_t): + ranges: list[range] = [] + for q_tile_idx, q_tile_size, kv_tile_size, window, length in zip( + (q_t_tile, q_h_tile, q_w_tile), + Q_TILE, + KV_TILE, + problem.window, + problem.shape, + ): + first_q = q_tile_idx * q_tile_size + last_q = min(first_q + q_tile_size, length) - 1 + first_k = _window_start(first_q, window, length) + last_k = _window_start(last_q, window, length) + window - 1 + ranges.append(range(first_k // kv_tile_size, last_k // kv_tile_size + 1)) + + t_tiles, h_tiles, w_tiles = ranges + row = [ + (kv_w * kv_grid_h + kv_h) * kv_grid_t + kv_t + for kv_w in w_tiles + for kv_h in h_tiles + for kv_t in t_tiles + ] + if not row: + raise AssertionError("every Q tile must overlap at least one KV tile") + rows.append(row) + + max_tiles = max(map(len, rows)) + # Duplicate a valid tile in padding slots. The kernel masks those slots + # with `counts`; avoiding -1 keeps every indirect global load in bounds. + indices = torch.tensor( + [row + [row[0]] * (max_tiles - len(row)) for row in rows], + dtype=torch.int32, + ) + counts = torch.tensor([len(row) for row in rows], dtype=torch.int32) + return indices, counts + + +@functools.lru_cache(maxsize=SHAPE_CACHE_SIZE) +def build_kv_tile_metadata( + problem: Problem, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Pack KV tile id and three separable mask pattern ids into int32.""" + indices, counts = build_kv_tile_map(problem) + q_grid_t, q_grid_h, _ = problem.q_grid + kv_grid_t, kv_grid_h, _ = problem.kv_grid + t_ids, t_patterns = _dimension_pattern_table(problem.frames, Q_TILE[0], KV_TILE[0], problem.window[0]) + h_ids, h_patterns = _dimension_pattern_table(problem.height, Q_TILE[1], KV_TILE[1], problem.window[1]) + w_ids, w_patterns = _dimension_pattern_table(problem.width, Q_TILE[2], KV_TILE[2], problem.window[2]) + + q_blocks = torch.arange(indices.shape[0], dtype=torch.int64)[:, None] + q_t = q_blocks % q_grid_t + q_h = (q_blocks // q_grid_t) % q_grid_h + q_w = q_blocks // (q_grid_t * q_grid_h) + kv = indices.to(torch.int64) + kv_t = kv % kv_grid_t + kv_h = (kv // kv_grid_t) % kv_grid_h + kv_w = kv // (kv_grid_t * kv_grid_h) + t_pattern_id = t_ids[q_t, kv_t] + h_pattern_id = h_ids[q_h, kv_h] + w_pattern_id = w_ids[q_w, kv_w] + metadata = kv | (t_pattern_id << 15) | (h_pattern_id << 19) | (w_pattern_id << 23) + active = torch.arange(indices.shape[1])[None, :] < counts[:, None] + metadata = torch.where(active, metadata, 0).to(torch.int32) + return metadata, t_patterns, h_patterns, w_patterns + + +@functools.lru_cache(maxsize=SHAPE_CACHE_SIZE) +def build_kv_tile_buckets( + problem: Problem, +) -> tuple[tuple[torch.Tensor, torch.Tensor, int], ...]: + """Group Q tiles by an upper bound on their useful KV tile count.""" + metadata, *_ = build_kv_tile_metadata(problem) + _, counts = build_kv_tile_map(problem) + maximum = int(counts.max()) + if maximum > KV_BUCKET_LIMITS[-1]: + raise ValueError( + f"FNA neighborhood needs {maximum} KV tiles, exceeding the supported ceiling {KV_BUCKET_LIMITS[-1]}" + ) + buckets = [] + lower = 0 + for limit in KV_BUCKET_LIMITS: + q_tile_ids = torch.nonzero((counts > lower) & (counts <= limit), as_tuple=False).flatten() + if q_tile_ids.numel(): + # Small grids can have fewer columns than the first bucket limit. + # Zero metadata encodes an all-masked tile, preserving the fixed ABI. + bucket_metadata = torch.zeros((q_tile_ids.numel(), limit), dtype=torch.int32) + columns = min(limit, metadata.shape[1]) + bucket_metadata[:, :columns] = metadata[q_tile_ids, :columns] + buckets.append( + ( + q_tile_ids.to(torch.int32), + bucket_metadata, + limit, + ) + ) + lower = limit + if limit >= maximum: + break + return tuple(buckets) + + +@functools.lru_cache(maxsize=32) +def compile_kernel( + max_kv_tiles: int, + *, + target: str = "auto", +): + """Compile one T/H/W-dynamic kernel for a fixed KV-count ceiling.""" + if max_kv_tiles not in KV_BUCKET_LIMITS: + raise ValueError(f"unsupported KV bucket ceiling: {max_kv_tiles}") + + q_length = T.dynamic("q_length") + kv_length = T.dynamic("kv_length") + q_tile_count = T.dynamic("q_tile_count") + + # Batch is fixed to one by the public guard. Excluding it from the + # kernel ABI keeps every physical stride static even when token lengths + # vary, preserving vectorized global/shared-memory copies. + q_shape = [q_length, HEADS, HEAD_DIM] + kv_shape = [kv_length, HEADS, HEAD_DIM] + q_tile_id_shape = [q_tile_count] + map_shape = [q_tile_count, max_kv_tiles] + pattern_shape = [PATTERN_TABLE_SIZE] + dtype = T.bfloat16 + accum_dtype = T.float32 + + @T.prim_func + def main( + query: T.Tensor(q_shape, dtype), + key: T.Tensor(kv_shape, dtype), + value: T.Tensor(kv_shape, dtype), + q_tile_ids: T.Tensor(q_tile_id_shape, T.int32), + kv_tile_metadata: T.Tensor(map_shape, T.int32), + t_mask_patterns: T.Tensor(pattern_shape, T.int32), + h_mask_patterns: T.Tensor(pattern_shape, T.int32), + w_mask_patterns: T.Tensor(pattern_shape, T.int32), + output: T.Tensor(q_shape, dtype), + ): + with T.Kernel(q_tile_count, HEADS, threads=128) as (bx, by): + q_shared = T.alloc_shared([Q_TOKENS, HEAD_DIM], dtype) + k_shared = T.alloc_shared([KV_TOKENS, HEAD_DIM], dtype) + v_shared = T.alloc_shared([KV_TOKENS, HEAD_DIM], dtype) + o_shared = T.alloc_shared([Q_TOKENS, HEAD_DIM], dtype) + scores = T.alloc_fragment([Q_TOKENS, KV_TOKENS], accum_dtype) + probs = T.alloc_fragment([Q_TOKENS, KV_TOKENS], dtype) + acc_o = T.alloc_fragment([Q_TOKENS, HEAD_DIM], accum_dtype) + row_max = T.alloc_fragment([Q_TOKENS], accum_dtype) + row_max_prev = T.alloc_fragment([Q_TOKENS], accum_dtype) + row_scale = T.alloc_fragment([Q_TOKENS], accum_dtype) + row_sum = T.alloc_fragment([Q_TOKENS], accum_dtype) + row_denom = T.alloc_fragment([Q_TOKENS], accum_dtype) + + q_block = q_tile_ids[bx] + # CPU metadata only emits padded Q/KV tiles that are fully backed + # by the permuted tensors. Expose those invariants so dynamic + # lengths do not add per-copy boundary predicates. + T.assume(q_block >= 0) + T.assume((q_block + 1) * Q_TOKENS <= q_length) + T.copy( + query[ + q_block * Q_TOKENS : (q_block + 1) * Q_TOKENS, + by, + :, + ], + q_shared, + ) + T.fill(acc_o, 0) + T.fill(row_denom, 0) + T.fill(row_max, -T.infinity(accum_dtype)) + + for block_iter in T.Pipelined(max_kv_tiles, num_stages=1): + metadata = kv_tile_metadata[bx, block_iter] + kv_block = metadata & 0x7FFF + T.assume((kv_block + 1) * KV_TOKENS <= kv_length) + t_mask = t_mask_patterns[(metadata >> 15) & 0xF] + h_mask = h_mask_patterns[(metadata >> 19) & 0xF] + w_mask = w_mask_patterns[(metadata >> 23) & 0xF] + T.copy( + key[ + kv_block * KV_TOKENS : (kv_block + 1) * KV_TOKENS, + by, + :, + ], + k_shared, + prefer_instruction="cp_async", + ) + + for i, j in T.Parallel(Q_TOKENS, KV_TOKENS): + q_t = i % Q_TILE[0] + q_h = (i // Q_TILE[0]) % Q_TILE[1] + q_w = i // (Q_TILE[0] * Q_TILE[1]) + k_t = j % KV_TILE[0] + k_h = (j // KV_TILE[0]) % KV_TILE[1] + k_w = j // (KV_TILE[0] * KV_TILE[1]) + valid = ( + (((t_mask >> (q_t * KV_TILE[0] + k_t)) & 1) != 0) + & (((h_mask >> (q_h * KV_TILE[1] + k_h)) & 1) != 0) + & (((w_mask >> (q_w * KV_TILE[2] + k_w)) & 1) != 0) + ) + scores[i, j] = T.if_then_else(valid, 0, -T.infinity(accum_dtype)) + + T.gemm( + q_shared, + k_shared, + scores, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + T.copy(row_max, row_max_prev) + T.fill(row_max, -T.infinity(accum_dtype)) + T.reduce_max(scores, row_max, dim=1, clear=False) + # A KV tile is selected for the union of all 64 queries in + # this Q tile. It can therefore be fully masked for one + # individual query. Preserve that row's previous online + # softmax state instead of evaluating (-inf) - (-inf). + T.copy(row_max, row_sum) + for i in T.Parallel(Q_TOKENS): + has_score = row_sum[i] > -T.infinity(accum_dtype) + merged_max = T.max(row_sum[i], row_max_prev[i]) + row_max[i] = T.if_then_else(has_score, merged_max, row_max_prev[i]) + row_scale[i] = T.if_then_else( + has_score, + T.exp2((row_max_prev[i] - merged_max) * LOG2_E), + 1, + ) + # Match Hopper NATTEN's two separately rounded FP32 + # products before exp2. + row_sum[i] = T.if_then_else( + has_score, + merged_max * LOG2_E, + -T.infinity(accum_dtype), + ) + for i, j in T.Parallel(Q_TOKENS, KV_TOKENS): + scores[i, j] = T.if_then_else( + row_sum[i] > -T.infinity(accum_dtype), + T.exp2(scores[i, j] * LOG2_E - row_sum[i]), + 0, + ) + T.reduce_sum(scores, row_sum, dim=1) + for i in T.Parallel(Q_TOKENS): + row_denom[i] = row_denom[i] * row_scale[i] + row_sum[i] + T.copy(scores, probs) + for i, j in T.Parallel(Q_TOKENS, HEAD_DIM): + acc_o[i, j] *= row_scale[i] + + T.copy( + value[ + kv_block * KV_TOKENS : (kv_block + 1) * KV_TOKENS, + by, + :, + ], + v_shared, + prefer_instruction="cp_async", + ) + T.gemm(probs, v_shared, acc_o, policy=T.GemmWarpPolicy.FullRow) + + # Match NATTEN's epilogue: compute one reciprocal per query row, + # then multiply every output accumulator by that shared alpha. + for i in T.Parallel(Q_TOKENS): + row_scale[i] = T.if_then_else( + row_denom[i] > 0, + T.call_extern("float32", "__frcp_rn", row_denom[i]), + 0, + ) + for i, j in T.Parallel(Q_TOKENS, HEAD_DIM): + acc_o[i, j] *= row_scale[i] + T.copy(acc_o, o_shared) + T.copy( + o_shared, + output[ + q_block * Q_TOKENS : (q_block + 1) * Q_TOKENS, + by, + :, + ], + ) + + return tilelang.compile( + main, + out_idx=None, + target="auto" if target == "auto" else {"kind": "cuda", "arch": target}, + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + # Keeping all work in the 128-thread consumer warpgroup avoids a + # separate producer warpgroup. The generated 240-register + # consumer can then admit two CTAs per SM instead of one. + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + }, + ) + + +_DEVICE_MAP_CACHE: OrderedDict[tuple[Problem, torch.device], tuple[object, ...]] = OrderedDict() + + +def _fna3d_tilelang( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + shape: tuple[int, int, int], + window: tuple[int, int, int], +) -> torch.Tensor: + problem = Problem(*shape, window=window) + if query.shape[1] != problem.q_length or key.shape[1] != problem.kv_length: + raise ValueError(f"permuted lengths do not match {problem}: Q={query.shape[1]}, KV={key.shape[1]}") + cache_key = (problem, query.device) + device_inputs = _DEVICE_MAP_CACHE.get(cache_key) + if device_inputs is None: + _, t_patterns, h_patterns, w_patterns = build_kv_tile_metadata(problem) + pattern_tables = tuple( + tensor.to(query.device, non_blocking=True) for tensor in (t_patterns, h_patterns, w_patterns) + ) + buckets = tuple( + ( + q_tile_ids.to(query.device, non_blocking=True), + metadata.to(query.device, non_blocking=True), + limit, + ) + for q_tile_ids, metadata, limit in build_kv_tile_buckets(problem) + ) + device_inputs = (pattern_tables, buckets) + _DEVICE_MAP_CACHE[cache_key] = device_inputs + if len(_DEVICE_MAP_CACHE) > SHAPE_CACHE_SIZE: + _DEVICE_MAP_CACHE.popitem(last=False) + else: + _DEVICE_MAP_CACHE.move_to_end(cache_key) + + pattern_tables, buckets = device_inputs + output = torch.empty_like(query) + for q_tile_ids, metadata, limit in buckets: + from tilelang.backend.target import determine_target + + target = str(determine_target("auto", return_object=True).attrs["arch"]) + kernel = compile_kernel(limit, target=target) + for batch in range(query.shape[0]): + kernel(query[batch], key[batch], value[batch], q_tile_ids, metadata, *pattern_tables, output[batch]) + return output + + +def fna3d_tilelang( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + shape: tuple[int, int, int], + window: tuple[int, int, int], +) -> torch.Tensor: + """Run native stage-5 FNA, or raise for unsupported inputs/launch failures. + + Q and KV use the documented token-tiled layout. Batch members share spatial + metadata; compilation is cached by target architecture and KV-count bucket. + """ + if not ( + query.is_cuda + and query.dtype is torch.bfloat16 + and query.ndim == key.ndim == value.ndim == 4 + and query.shape[0] > 0 + and query.shape[-2:] == (HEADS, HEAD_DIM) + and query.numel() > 0 + and query.is_contiguous() + and key.device == value.device == query.device + and key.dtype is value.dtype is query.dtype + and key.is_contiguous() + and value.is_contiguous() + and value.shape == key.shape + and key.shape[0] == query.shape[0] + and key.shape[-2:] == query.shape[-2:] + and len(shape) == 3 + and tuple(window) == (11, 11, 11) + ): + raise ValueError( + "LTX DiffVAE FNA requires contiguous CUDA BF16 [B, tokens, 4, 64] Q/K/V and an 11x11x11 window." + ) + if torch.is_grad_enabled() and any(x.requires_grad for x in (query, key, value)): + raise RuntimeError("LTX DiffVAE FNA is inference-only.") + with torch.accelerator.device_index(query.device.index): + return _fna3d_tilelang(query, key, value, shape=tuple(shape), window=tuple(window)) + + +__all__ = ["fna3d_tilelang"] diff --git a/vllm_omni/diffusion/models/ltx2/ops/qk_rms_norm.py b/vllm_omni/diffusion/models/ltx2/ops/qk_rms_norm.py new file mode 100644 index 00000000000..be89aeca293 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ops/qk_rms_norm.py @@ -0,0 +1,377 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Bit-exact multi-row Q/K RMSNorm for the LTX diffusion VAE decoder.""" + +from __future__ import annotations + +import logging +import math + +import torch +import torch.nn.functional as F +from vllm.triton_utils import tl, triton + +from . import is_ltx2_fusion_eligible +from ._numerics import ( + add_rn_f32, + fma_rn_f32, + mul_rn_f32, + round_bf16_to_fp32, + rsqrt_approx_f32, + shfl_down_f32, +) +from .rope_tables import get_rope_tables + +_HEAD_DIM = 64 +_ROWS_PER_PROGRAM = 2 +_VERIFY_ROWS = 256 +_FAILED_ROPE_KEYS: set[tuple[int | None, float, float, tuple[int, int, int], float]] = set() +_VERIFIED_ROPE_KEYS: set[tuple[int | None, float, float, tuple[int, int, int], float]] = set() + +logger = logging.getLogger(__name__) + + +@triton.jit +def _shuffle_lane_zero(value): + return tl.inline_asm_elementwise( + asm="shfl.sync.idx.b32 $0, $1, 0, 0x1f, 0xffffffff;", + constraints="=f,f", + args=[value], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _warp_sum(value): + value = add_rn_f32(value, shfl_down_f32(value, 16)) + value = add_rn_f32(value, shfl_down_f32(value, 8)) + value = add_rn_f32(value, shfl_down_f32(value, 4)) + value = add_rn_f32(value, shfl_down_f32(value, 2)) + value = add_rn_f32(value, shfl_down_f32(value, 1)) + return _shuffle_lane_zero(value) + + +@triton.jit +def _qk_rms_norm_scale_rope_3d_kernel( + query_output_ptr, + key_output_ptr, + query_ptr, + key_ptr, + query_weight_ptr, + key_weight_ptr, + cos_t_ptr, + sin_t_ptr, + cos_h_ptr, + sin_h_ptr, + cos_w_ptr, + sin_w_ptr, + rows, + tokens_per_batch, + height, + width, + heads, + eps, + query_scale, + rows_per_program: tl.constexpr, + head_dim: tl.constexpr, + pairs_t: tl.constexpr, + pairs_h: tl.constexpr, +): + threads = tl.arange(0, rows_per_program * 32) + local_row = threads // 32 + lane = threads % 32 + row = tl.program_id(0).to(tl.int64) * rows_per_program + local_row + valid_row = row < rows + reduction_valid = valid_row & (lane < 16) + reduction_base = row * head_dim + lane * 4 + + query_accumulator = tl.zeros((rows_per_program * 32,), dtype=tl.float32) + key_accumulator = tl.zeros((rows_per_program * 32,), dtype=tl.float32) + for element in tl.static_range(4): + query = tl.load(query_ptr + reduction_base + element, mask=reduction_valid, other=0.0).to(tl.float32) + key = tl.load(key_ptr + reduction_base + element, mask=reduction_valid, other=0.0).to(tl.float32) + query_accumulator = fma_rn_f32(query, query, query_accumulator) + key_accumulator = fma_rn_f32(key, key, key_accumulator) + query_rstd = rsqrt_approx_f32(_warp_sum(query_accumulator) / head_dim + eps) + key_rstd = rsqrt_approx_f32(_warp_sum(key_accumulator) / head_dim + eps) + + tokens = (row // heads) % tokens_per_batch + pos_w = tokens % width + pos_h = (tokens // width) % height + pos_t = tokens // (height * width) + pairs_per_head: tl.constexpr = head_dim // 2 + pairs_w: tl.constexpr = pairs_per_head - pairs_t - pairs_h + is_t = lane < pairs_t + is_h = (lane >= pairs_t) & (lane < pairs_t + pairs_h) + is_w = ~(is_t | is_h) + pair_t = lane + pair_h = lane - pairs_t + pair_w = lane - pairs_t - pairs_h + cos_t = tl.load(cos_t_ptr + pos_t * pairs_t + pair_t, mask=valid_row & is_t, other=0.0) + sin_t = tl.load(sin_t_ptr + pos_t * pairs_t + pair_t, mask=valid_row & is_t, other=0.0) + cos_h = tl.load(cos_h_ptr + pos_h * pairs_h + pair_h, mask=valid_row & is_h, other=0.0) + sin_h = tl.load(sin_h_ptr + pos_h * pairs_h + pair_h, mask=valid_row & is_h, other=0.0) + cos_w = tl.load(cos_w_ptr + pos_w * pairs_w + pair_w, mask=valid_row & is_w, other=0.0) + sin_w = tl.load(sin_w_ptr + pos_w * pairs_w + pair_w, mask=valid_row & is_w, other=0.0) + cos = cos_t + cos_h + cos_w + sin = sin_t + sin_h + sin_w + + even_column = lane * 2 + odd_column = even_column + 1 + even_offset = row * head_dim + even_column + odd_offset = row * head_dim + odd_column + query_even = tl.load(query_ptr + even_offset, mask=valid_row).to(tl.float32) + query_odd = tl.load(query_ptr + odd_offset, mask=valid_row).to(tl.float32) + key_even = tl.load(key_ptr + even_offset, mask=valid_row).to(tl.float32) + key_odd = tl.load(key_ptr + odd_offset, mask=valid_row).to(tl.float32) + query_even_weight = tl.load(query_weight_ptr + even_column).to(tl.float32) + query_odd_weight = tl.load(query_weight_ptr + odd_column).to(tl.float32) + key_even_weight = tl.load(key_weight_ptr + even_column).to(tl.float32) + key_odd_weight = tl.load(key_weight_ptr + odd_column).to(tl.float32) + + query_even = round_bf16_to_fp32(mul_rn_f32(query_even_weight, mul_rn_f32(query_rstd, query_even))) + query_odd = round_bf16_to_fp32(mul_rn_f32(query_odd_weight, mul_rn_f32(query_rstd, query_odd))) + query_even = round_bf16_to_fp32(mul_rn_f32(query_even, query_scale)) + query_odd = round_bf16_to_fp32(mul_rn_f32(query_odd, query_scale)) + key_even = round_bf16_to_fp32(mul_rn_f32(key_even_weight, mul_rn_f32(key_rstd, key_even))) + key_odd = round_bf16_to_fp32(mul_rn_f32(key_odd_weight, mul_rn_f32(key_rstd, key_odd))) + + tl.store( + query_output_ptr + even_offset, + add_rn_f32(mul_rn_f32(query_even, cos), -mul_rn_f32(query_odd, sin)), + mask=valid_row, + ) + tl.store( + query_output_ptr + odd_offset, + add_rn_f32(mul_rn_f32(query_even, sin), mul_rn_f32(query_odd, cos)), + mask=valid_row, + ) + tl.store( + key_output_ptr + even_offset, + add_rn_f32(mul_rn_f32(key_even, cos), -mul_rn_f32(key_odd, sin)), + mask=valid_row, + ) + tl.store( + key_output_ptr + odd_offset, + add_rn_f32(mul_rn_f32(key_even, sin), mul_rn_f32(key_odd, cos)), + mask=valid_row, + ) + + +def _weight_matches(value: torch.Tensor, weight: torch.Tensor) -> bool: + return ( + weight.is_cuda + and weight.device == value.device + and weight.dtype is torch.bfloat16 + and weight.shape == (_HEAD_DIM,) + and weight.is_contiguous() + ) + + +def _supported_inputs( + query: torch.Tensor, + key: torch.Tensor, + query_weight: torch.Tensor, + key_weight: torch.Tensor, + eps: float, + query_scale: float, +) -> bool: + return ( + is_ltx2_fusion_eligible(query) + and query.dtype is torch.bfloat16 + and query.ndim == 6 + and query.shape[-1] == _HEAD_DIM + and query.numel() > 0 + and query.is_contiguous() + and key.is_cuda + and key.device == query.device + and key.dtype is query.dtype + and key.shape == query.shape + and key.is_contiguous() + and _weight_matches(query, query_weight) + and _weight_matches(query, key_weight) + and math.isfinite(eps) + and eps > 0 + and math.isfinite(query_scale) + ) + + +def _combined_supported_inputs( + query: torch.Tensor, + key: torch.Tensor, + query_weight: torch.Tensor, + key_weight: torch.Tensor, + eps: float, + query_scale: float, + dim_split: tuple[int, int, int], + base: float, +) -> bool: + return ( + _supported_inputs(query, key, query_weight, key_weight, eps, query_scale) + and len(dim_split) == 3 + and sum(dim_split) == _HEAD_DIM + and all(dim > 0 and dim % 2 == 0 for dim in dim_split) + and math.isfinite(base) + and base > 0 + ) + + +def _launch_combined( + query: torch.Tensor, + key: torch.Tensor, + query_weight: torch.Tensor, + key_weight: torch.Tensor, + eps: float, + query_scale: float, + dim_split: tuple[int, int, int], + base: float, +) -> tuple[torch.Tensor, torch.Tensor]: + query_output = torch.empty_like(query) + key_output = torch.empty_like(key) + rows = query.numel() // _HEAD_DIM + tables = get_rope_tables(query, dim_split, base) + cos_t, sin_t = tables[0] + cos_h, sin_h = tables[1] + cos_w, sin_w = tables[2] + with torch.accelerator.device_index(query.device.index): + _qk_rms_norm_scale_rope_3d_kernel[(triton.cdiv(rows, _ROWS_PER_PROGRAM),)]( + query_output, + key_output, + query, + key, + query_weight, + key_weight, + cos_t, + sin_t, + cos_h, + sin_h, + cos_w, + sin_w, + rows, + query.shape[1] * query.shape[2] * query.shape[3], + query.shape[2], + query.shape[3], + query.shape[4], + eps, + query_scale, + rows_per_program=_ROWS_PER_PROGRAM, + head_dim=_HEAD_DIM, + pairs_t=dim_split[0] // 2, + pairs_h=dim_split[1] // 2, + num_warps=_ROWS_PER_PROGRAM, + ) + return query_output, key_output + + +def _verify_combined_prefix( + output: tuple[torch.Tensor, torch.Tensor], + query: torch.Tensor, + key: torch.Tensor, + query_weight: torch.Tensor, + key_weight: torch.Tensor, + eps: float, + query_scale: float, + dim_split: tuple[int, int, int], + base: float, +) -> bool: + rows = min(query.numel() // _HEAD_DIM, _VERIFY_ROWS) + query_rows = query.reshape(-1, _HEAD_DIM)[:rows] + key_rows = key.reshape(-1, _HEAD_DIM)[:rows] + query_rows = F.rms_norm(query_rows, (_HEAD_DIM,), query_weight, eps) * query_scale + key_rows = F.rms_norm(key_rows, (_HEAD_DIM,), key_weight, eps) + tables = get_rope_tables(query, dim_split, base) + row_indices = torch.arange(rows, device=query.device) + tokens_per_batch = query.shape[1] * query.shape[2] * query.shape[3] + tokens = (row_indices // query.shape[4]) % tokens_per_batch + positions = ( + tokens // (query.shape[2] * query.shape[3]), + (tokens // query.shape[3]) % query.shape[2], + tokens % query.shape[3], + ) + cos = torch.cat([axis_tables[0][position] for axis_tables, position in zip(tables, positions, strict=True)], -1) + sin = torch.cat([axis_tables[1][position] for axis_tables, position in zip(tables, positions, strict=True)], -1) + + def rotate(value: torch.Tensor) -> torch.Tensor: + pairs = value.reshape(rows, _HEAD_DIM // 2, 2) + even = pairs[..., 0].float() + odd = pairs[..., 1].float() + return ( + torch.stack((even * cos - odd * sin, even * sin + odd * cos), -1).reshape(rows, _HEAD_DIM).to(value.dtype) + ) + + reference = rotate(query_rows), rotate(key_rows) + actual = tuple(value.reshape(-1, _HEAD_DIM)[:rows] for value in output) + return all(torch.equal(a, b) for a, b in zip(actual, reference, strict=True)) + + +def try_qk_rms_norm_scale_rope_3d_exact( + query: torch.Tensor, + key: torch.Tensor, + query_weight: torch.Tensor, + key_weight: torch.Tensor, + eps: float, + query_scale: float, + dim_split: tuple[int, int, int], + base: float, +) -> tuple[torch.Tensor, torch.Tensor] | None: + """Normalize Q/K, scale Q, and apply paired 3D RoPE in one kernel.""" + + if not _combined_supported_inputs( + query, + key, + query_weight, + key_weight, + eps, + query_scale, + dim_split, + base, + ): + return None + runtime_key = (query.device.index, float(eps), float(query_scale), dim_split, float(base)) + if runtime_key in _FAILED_ROPE_KEYS: + return None + try: + output = _launch_combined( + query, + key, + query_weight, + key_weight, + eps, + query_scale, + dim_split, + base, + ) + if runtime_key not in _VERIFIED_ROPE_KEYS: + if not _verify_combined_prefix( + output, + query, + key, + query_weight, + key_weight, + eps, + query_scale, + dim_split, + base, + ): + _FAILED_ROPE_KEYS.add(runtime_key) + logger.warning( + "Disabling LTX DiffVAE Q/K RMSNorm+3D RoPE fusion on %s after a bit-exactness mismatch", + query.device, + ) + return None + _VERIFIED_ROPE_KEYS.add(runtime_key) + except Exception as exc: # noqa: BLE001 - fail closed after optimized-path failure + _FAILED_ROPE_KEYS.add(runtime_key) + logger.warning( + "Disabling LTX DiffVAE Q/K RMSNorm+3D RoPE fusion on %s after failure: %s", + query.device, + exc, + ) + return None + return output + + +__all__ = ["try_qk_rms_norm_scale_rope_3d_exact"] diff --git a/vllm_omni/diffusion/models/ltx2/ops/residual_adaln.py b/vllm_omni/diffusion/models/ltx2/ops/residual_adaln.py new file mode 100644 index 00000000000..16a4aefe6f0 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ops/residual_adaln.py @@ -0,0 +1,489 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Exact LTX DiffVAE residual, RMSNorm, and AdaLN Triton fusions.""" + +from __future__ import annotations + +import logging +import math + +import torch +import torch.nn.functional as F +from vllm.triton_utils import tl, triton + +from . import is_ltx2_fusion_eligible +from ._numerics import ( + add_rn_f32, + fma_rn_f32, + mul_rn_f32, + round_bf16_to_fp32, + rsqrt_approx_f32, + shfl_down_f32, +) + +_HIDDEN_SIZE = 256 +_POINTWISE_BLOCK = 1024 +_VERIFY_ROWS = 256 +_VERIFY_ELEMENTS = 1 << 20 +_FAILED_ADALN_KEYS: set[tuple[int | None, int, float]] = set() +_VERIFIED_ADALN_KEYS: set[tuple[int | None, int, float]] = set() +_FAILED_ADD_DEVICES: set[int | None] = set() +_VERIFIED_ADD_DEVICES: set[int | None] = set() + +logger = logging.getLogger(__name__) + + +@triton.jit +def _two_warp_sum(values): + values = add_rn_f32(values, shfl_down_f32(values, 16)) + values = add_rn_f32(values, shfl_down_f32(values, 8)) + values = add_rn_f32(values, shfl_down_f32(values, 4)) + values = add_rn_f32(values, shfl_down_f32(values, 2)) + values = add_rn_f32(values, shfl_down_f32(values, 1)) + threads = tl.arange(0, 64) + warp_0 = tl.sum(tl.where(threads == 0, values, 0.0)) + warp_1 = tl.sum(tl.where(threads == 32, values, 0.0)) + return add_rn_f32(warp_0, warp_1) + + +@triton.jit +def _load_ordered_residual( + x_ptr, + residual_a_ptr, + residual_b_ptr, + offset, + valid, + residual_count: tl.constexpr, +): + value = add_rn_f32( + tl.load(x_ptr + offset, mask=valid, other=0.0).to(tl.float32), + tl.load(residual_a_ptr + offset, mask=valid, other=0.0).to(tl.float32), + ) + value = round_bf16_to_fp32(value) + if residual_count == 2: + value = add_rn_f32( + value, + tl.load(residual_b_ptr + offset, mask=valid, other=0.0).to(tl.float32), + ) + value = round_bf16_to_fp32(value) + return value + + +@triton.jit +def _store_modulated( + output_ptr, + norm_weight_ptr, + scale_ptr, + shift_ptr, + value, + offset, + column, + modulation_base, + reciprocal_rms, + valid, +): + normalized = mul_rn_f32(value, reciprocal_rms) + norm_weight = tl.load(norm_weight_ptr + column, mask=valid, other=0.0).to(tl.float32) + normalized = round_bf16_to_fp32(mul_rn_f32(normalized, norm_weight)) + scale = tl.load(scale_ptr + modulation_base + column, mask=valid, other=0.0).to(tl.float32) + shift = tl.load(shift_ptr + modulation_base + column, mask=valid, other=0.0).to(tl.float32) + one_plus_scale = round_bf16_to_fp32(add_rn_f32(1.0, scale)) + product = round_bf16_to_fp32(mul_rn_f32(normalized, one_plus_scale)) + tl.store(output_ptr + offset, add_rn_f32(product, shift), mask=valid) + + +@triton.jit +def _residual_rms_norm_modulate_kernel( + output_ptr, + x_ptr, + residual_a_ptr, + residual_b_ptr, + norm_weight_ptr, + scale_ptr, + shift_ptr, + rows_per_batch, + eps, + residual_count: tl.constexpr, + hidden_size: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + batch = row // rows_per_batch + threads = tl.arange(0, 64) + row_base = row * hidden_size + + # PyTorch's BF16 RMSNorm kernel uses four contiguous values per thread for + # hidden=256. Keep the four ordered BF16 residual values in registers so + # the output pass does not read the large activations a second time. + column_0 = threads * 4 + column_1 = column_0 + 1 + column_2 = column_0 + 2 + column_3 = column_0 + 3 + valid_0 = column_0 < hidden_size + valid_1 = column_1 < hidden_size + valid_2 = column_2 < hidden_size + valid_3 = column_3 < hidden_size + offset_0 = row_base + column_0 + offset_1 = row_base + column_1 + offset_2 = row_base + column_2 + offset_3 = row_base + column_3 + value_0 = _load_ordered_residual(x_ptr, residual_a_ptr, residual_b_ptr, offset_0, valid_0, residual_count) + value_1 = _load_ordered_residual(x_ptr, residual_a_ptr, residual_b_ptr, offset_1, valid_1, residual_count) + value_2 = _load_ordered_residual(x_ptr, residual_a_ptr, residual_b_ptr, offset_2, valid_2, residual_count) + value_3 = _load_ordered_residual(x_ptr, residual_a_ptr, residual_b_ptr, offset_3, valid_3, residual_count) + accumulator = fma_rn_f32(value_0, value_0, tl.zeros((64,), dtype=tl.float32)) + accumulator = fma_rn_f32(value_1, value_1, accumulator) + accumulator = fma_rn_f32(value_2, value_2, accumulator) + accumulator = fma_rn_f32(value_3, value_3, accumulator) + + total = _two_warp_sum(accumulator) + reciprocal_rms = rsqrt_approx_f32(total / hidden_size + eps) + modulation_base = batch * hidden_size + _store_modulated( + output_ptr, + norm_weight_ptr, + scale_ptr, + shift_ptr, + value_0, + offset_0, + column_0, + modulation_base, + reciprocal_rms, + valid_0, + ) + _store_modulated( + output_ptr, + norm_weight_ptr, + scale_ptr, + shift_ptr, + value_1, + offset_1, + column_1, + modulation_base, + reciprocal_rms, + valid_1, + ) + _store_modulated( + output_ptr, + norm_weight_ptr, + scale_ptr, + shift_ptr, + value_2, + offset_2, + column_2, + modulation_base, + reciprocal_rms, + valid_2, + ) + _store_modulated( + output_ptr, + norm_weight_ptr, + scale_ptr, + shift_ptr, + value_3, + offset_3, + column_3, + modulation_base, + reciprocal_rms, + valid_3, + ) + + +@triton.jit +def _residual_add3_kernel( + output_ptr, + x_ptr, + residual_a_ptr, + residual_b_ptr, + residual_c_ptr, + elements, + block: tl.constexpr, +): + offsets = tl.program_id(0).to(tl.int64) * block + tl.arange(0, block) + valid = offsets < elements + value = add_rn_f32( + tl.load(x_ptr + offsets, mask=valid, other=0.0).to(tl.float32), + tl.load(residual_a_ptr + offsets, mask=valid, other=0.0).to(tl.float32), + ) + value = round_bf16_to_fp32(value) + value = add_rn_f32( + value, + tl.load(residual_b_ptr + offsets, mask=valid, other=0.0).to(tl.float32), + ) + value = round_bf16_to_fp32(value) + value = add_rn_f32( + value, + tl.load(residual_c_ptr + offsets, mask=valid, other=0.0).to(tl.float32), + ) + tl.store(output_ptr + offsets, value, mask=valid) + + +def _same_activation(x: torch.Tensor, residual: torch.Tensor) -> bool: + return ( + residual.dtype is torch.bfloat16 + and residual.is_cuda + and residual.device == x.device + and residual.shape == x.shape + and residual.is_contiguous() + ) + + +def _modulation_matches(x: torch.Tensor, value: torch.Tensor) -> bool: + return ( + value.dtype is torch.bfloat16 + and value.is_cuda + and value.device == x.device + and value.shape[-1] == _HIDDEN_SIZE + and value.numel() == x.shape[0] * _HIDDEN_SIZE + and value.is_contiguous() + ) + + +def _norm_weight_matches(x: torch.Tensor, weight: torch.Tensor) -> bool: + return ( + weight.dtype is torch.bfloat16 + and weight.is_cuda + and weight.device == x.device + and weight.shape == (_HIDDEN_SIZE,) + and weight.is_contiguous() + ) + + +def _adaln_inputs_supported( + x: torch.Tensor, + residual_a: torch.Tensor, + residual_b: torch.Tensor | None, + norm_weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + eps: float, +) -> bool: + return ( + is_ltx2_fusion_eligible(x) + and x.dtype is torch.bfloat16 + and x.ndim >= 2 + and x.shape[-1] == _HIDDEN_SIZE + and x.numel() > 0 + and x.is_contiguous() + and _same_activation(x, residual_a) + and (residual_b is None or _same_activation(x, residual_b)) + and _norm_weight_matches(x, norm_weight) + and _modulation_matches(x, scale) + and _modulation_matches(x, shift) + and math.isfinite(eps) + and eps > 0 + ) + + +def _residual_reference( + x: torch.Tensor, + residual_a: torch.Tensor, + residual_b: torch.Tensor | None, +) -> torch.Tensor: + hidden_states = x + residual_a + if residual_b is not None: + hidden_states = hidden_states + residual_b + return hidden_states + + +def _adaln_reference( + x: torch.Tensor, + residual_a: torch.Tensor, + residual_b: torch.Tensor | None, + norm_weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + eps: float, +) -> torch.Tensor: + hidden_states = _residual_reference(x, residual_a, residual_b) + normalized = F.rms_norm(hidden_states, (_HIDDEN_SIZE,), weight=norm_weight, eps=eps) + return normalized * (1 + scale) + shift + + +def _launch_adaln( + x: torch.Tensor, + residual_a: torch.Tensor, + residual_b: torch.Tensor | None, + norm_weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + eps: float, +) -> torch.Tensor: + output = torch.empty_like(x) + rows = x.numel() // _HIDDEN_SIZE + rows_per_batch = rows // x.shape[0] + with torch.accelerator.device_index(x.device.index): + _residual_rms_norm_modulate_kernel[(rows,)]( + output, + x, + residual_a, + x if residual_b is None else residual_b, + norm_weight, + scale, + shift, + rows_per_batch, + eps, + residual_count=1 if residual_b is None else 2, + hidden_size=_HIDDEN_SIZE, + num_warps=2, + ) + return output + + +def _verify_adaln_prefix( + output: torch.Tensor, + x: torch.Tensor, + residual_a: torch.Tensor, + residual_b: torch.Tensor | None, + norm_weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + eps: float, +) -> bool: + rows_per_batch = x.numel() // (x.shape[0] * _HIDDEN_SIZE) + rows = min(rows_per_batch, _VERIFY_ROWS) + x_rows = x.reshape(x.shape[0], rows_per_batch, _HIDDEN_SIZE)[0, :rows] + residual_a_rows = residual_a.reshape(x.shape[0], rows_per_batch, _HIDDEN_SIZE)[0, :rows] + residual_b_rows = ( + None if residual_b is None else residual_b.reshape(x.shape[0], rows_per_batch, _HIDDEN_SIZE)[0, :rows] + ) + scale_row = scale.reshape(x.shape[0], 1, _HIDDEN_SIZE)[0] + shift_row = shift.reshape(x.shape[0], 1, _HIDDEN_SIZE)[0] + reference = _adaln_reference( + x_rows, + residual_a_rows, + residual_b_rows, + norm_weight, + scale_row, + shift_row, + eps, + ) + actual = output.reshape(x.shape[0], rows_per_batch, _HIDDEN_SIZE)[0, :rows] + return torch.equal(actual, reference) + + +def try_residual_rms_norm_modulate_exact( + x: torch.Tensor, + residual_a: torch.Tensor, + residual_b: torch.Tensor | None, + norm_weight: torch.Tensor, + scale: torch.Tensor, + shift: torch.Tensor, + eps: float, +) -> torch.Tensor | None: + """Fuse one or two ordered BF16 residuals with RMSNorm and AdaLN.""" + + if not _adaln_inputs_supported(x, residual_a, residual_b, norm_weight, scale, shift, eps): + return None + residual_count = 1 if residual_b is None else 2 + runtime_key = (x.device.index, residual_count, float(eps)) + if runtime_key in _FAILED_ADALN_KEYS: + return None + try: + output = _launch_adaln(x, residual_a, residual_b, norm_weight, scale, shift, eps) + if runtime_key not in _VERIFIED_ADALN_KEYS: + if not _verify_adaln_prefix( + output, + x, + residual_a, + residual_b, + norm_weight, + scale, + shift, + eps, + ): + _FAILED_ADALN_KEYS.add(runtime_key) + logger.warning( + "Disabling LTX DiffVAE residual-AdaLN fusion on %s after a bit-exactness mismatch", + x.device, + ) + return None + _VERIFIED_ADALN_KEYS.add(runtime_key) + except Exception as exc: # noqa: BLE001 - fail closed after optimized-path failure + _FAILED_ADALN_KEYS.add(runtime_key) + logger.warning( + "Disabling LTX DiffVAE residual-AdaLN fusion on %s after failure: %s", + x.device, + exc, + ) + return None + return output + + +def _add_inputs_supported(x: torch.Tensor, residuals: tuple[torch.Tensor, ...]) -> bool: + return ( + is_ltx2_fusion_eligible(x) + and x.dtype is torch.bfloat16 + and x.numel() > 0 + and x.is_contiguous() + and all(_same_activation(x, residual) for residual in residuals) + ) + + +def _launch_add3( + x: torch.Tensor, + residual_a: torch.Tensor, + residual_b: torch.Tensor, + residual_c: torch.Tensor, +) -> torch.Tensor: + output = torch.empty_like(x) + with torch.accelerator.device_index(x.device.index): + _residual_add3_kernel[(triton.cdiv(x.numel(), _POINTWISE_BLOCK),)]( + output, + x, + residual_a, + residual_b, + residual_c, + x.numel(), + block=_POINTWISE_BLOCK, + num_warps=4, + ) + return output + + +def try_residual_add3_exact( + x: torch.Tensor, + residual_a: torch.Tensor, + residual_b: torch.Tensor, + residual_c: torch.Tensor, +) -> torch.Tensor | None: + """Fuse three ordered BF16 residual additions into one pointwise pass.""" + + residuals = (residual_a, residual_b, residual_c) + if not _add_inputs_supported(x, residuals) or x.device.index in _FAILED_ADD_DEVICES: + return None + try: + output = _launch_add3(x, *residuals) + if x.device.index not in _VERIFIED_ADD_DEVICES: + elements = min(x.numel(), _VERIFY_ELEMENTS) + reference = ( + _residual_reference( + x.reshape(-1)[:elements], + residual_a.reshape(-1)[:elements], + residual_b.reshape(-1)[:elements], + ) + + residual_c.reshape(-1)[:elements] + ) + if not torch.equal(output.reshape(-1)[:elements], reference): + _FAILED_ADD_DEVICES.add(x.device.index) + logger.warning( + "Disabling LTX DiffVAE residual-add fusion on %s after a bit-exactness mismatch", + x.device, + ) + return None + _VERIFIED_ADD_DEVICES.add(x.device.index) + except Exception as exc: # noqa: BLE001 - fail closed after optimized-path failure + _FAILED_ADD_DEVICES.add(x.device.index) + logger.warning( + "Disabling LTX DiffVAE residual-add fusion on %s after failure: %s", + x.device, + exc, + ) + return None + return output + + +__all__ = [ + "try_residual_add3_exact", + "try_residual_rms_norm_modulate_exact", +] diff --git a/vllm_omni/diffusion/models/ltx2/ops/rope_tables.py b/vllm_omni/diffusion/models/ltx2/ops/rope_tables.py new file mode 100644 index 00000000000..78a2ad2ee3f --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ops/rope_tables.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Cached 3D RoPE tables for LTX DiffVAE kernels.""" + +from __future__ import annotations + +import torch + +_TABLE_CACHE_MAX = 16 +_TABLE_CACHE: dict[ + tuple[torch.device, int, int, int, tuple[int, int, int], float], + tuple[tuple[torch.Tensor, torch.Tensor], ...], +] = {} + + +def _axis_tables( + length: int, + dim: int, + device: torch.device, + base: float, +) -> tuple[torch.Tensor, torch.Tensor]: + exponents = torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim + inv_freqs = (1.0 / base**exponents).to(torch.float32) + positions = torch.arange(length, dtype=torch.float32, device=device) + angles = positions[:, None] * inv_freqs[None, :] + return angles.cos().contiguous(), angles.sin().contiguous() + + +def get_rope_tables( + query: torch.Tensor, + dim_split: tuple[int, int, int], + base: float, +) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: + num_frames, height, width = query.shape[1:4] + cache_key = (query.device, num_frames, height, width, dim_split, base) + cached = _TABLE_CACHE.get(cache_key) + if cached is not None: + return cached + tables = tuple( + _axis_tables(length, dim, query.device, base) + for length, dim in zip((num_frames, height, width), dim_split, strict=True) + ) + if len(_TABLE_CACHE) >= _TABLE_CACHE_MAX: + _TABLE_CACHE.pop(next(iter(_TABLE_CACHE))) + _TABLE_CACHE[cache_key] = tables + return tables + + +__all__ = ["get_rope_tables"] diff --git a/vllm_omni/diffusion/models/ltx2/ops/swiglu.py b/vllm_omni/diffusion/models/ltx2/ops/swiglu.py new file mode 100644 index 00000000000..6c332b7edc8 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ops/swiglu.py @@ -0,0 +1,206 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Bit-exact workspace-reusing SwiGLU for the LTX diffusion VAE decoder.""" + +from __future__ import annotations + +import logging + +import torch +import torch.nn.functional as F +from vllm.triton_utils import tl, triton + +from . import is_ltx2_fusion_eligible +from ._numerics import mul_rn_f32, round_bf16_to_fp32 + +_POINTWISE_BLOCK = 1024 +_BASE_TILE_ROWS = 16384 +_MAX_TILE_HIDDEN_ELEMENTS = _BASE_TILE_ROWS * 4 * 2048 +_FAILED_KEYS: set[tuple[int | None, int, int, int, int]] = set() +_VERIFIED_KEYS: set[tuple[int | None, int, int, int, int]] = set() + +logger = logging.getLogger(__name__) + + +def _tile_rows(rows: int, hidden: int) -> int: + """Bound projection storage while giving narrow MLPs larger GEMMs.""" + + budget_rows = _MAX_TILE_HIDDEN_ELEMENTS // hidden + aligned_rows = max(_BASE_TILE_ROWS, budget_rows // _BASE_TILE_ROWS * _BASE_TILE_ROWS) + return min(rows, aligned_rows) + + +@triton.jit +def _silu_mul_kernel( + gate_ptr, + up_ptr, + elements, + block: tl.constexpr, +): + offsets = tl.program_id(0).to(tl.int64) * block + tl.arange(0, block) + valid = offsets < elements + gate = tl.load(gate_ptr + offsets, mask=valid).to(tl.float32) + up = tl.load(up_ptr + offsets, mask=valid).to(tl.float32) + # Eager rounds SiLU to BF16 before the separate BF16 multiply. + silu = round_bf16_to_fp32(mul_rn_f32(gate, tl.sigmoid(gate))) + tl.store(gate_ptr + offsets, mul_rn_f32(silu, up), mask=valid) + + +def _weight_matches( + hidden_states: torch.Tensor, + weight: torch.Tensor, + shape: tuple[int, int], +) -> bool: + return ( + weight.is_cuda + and weight.device == hidden_states.device + and weight.dtype is torch.bfloat16 + and weight.shape == shape + and weight.is_contiguous() + ) + + +def _supported_inputs( + hidden_states: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, +) -> bool: + if hidden_states.ndim < 3: + return False + dim = hidden_states.shape[-1] + hidden = gate_weight.shape[0] if gate_weight.ndim == 2 else 0 + return ( + is_ltx2_fusion_eligible(hidden_states) + and hidden_states.dtype is torch.bfloat16 + and hidden_states.shape[0] == 1 + and hidden_states.numel() > 0 + and hidden_states.is_contiguous() + and 256 <= dim <= 2048 + and hidden == 4 * dim + and _weight_matches(hidden_states, gate_weight, (hidden, dim)) + and _weight_matches(hidden_states, up_weight, (hidden, dim)) + and _weight_matches(hidden_states, down_weight, (dim, hidden)) + ) + + +def _silu_mul_inplace(gate: torch.Tensor, up: torch.Tensor) -> None: + with torch.accelerator.device_index(gate.device.index): + _silu_mul_kernel[(triton.cdiv(gate.numel(), _POINTWISE_BLOCK),)]( + gate, + up, + gate.numel(), + block=_POINTWISE_BLOCK, + num_warps=4, + ) + + +def _reference_tile( + hidden_states: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, +) -> torch.Tensor: + return F.linear( + F.silu(F.linear(hidden_states, gate_weight)) * F.linear(hidden_states, up_weight), + down_weight, + ) + + +def _launch( + hidden_states: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, +) -> torch.Tensor: + dim = hidden_states.shape[-1] + hidden = gate_weight.shape[0] + flat = hidden_states.reshape(-1, dim) + tile_rows = _tile_rows(flat.shape[0], hidden) + output = torch.empty_like(flat) + workspace_gate = torch.empty((tile_rows, hidden), device=flat.device, dtype=flat.dtype) + for start in range(0, flat.shape[0], tile_rows): + end = min(start + tile_rows, flat.shape[0]) + tokens = end - start + tile = flat[start:end] + gate = workspace_gate[:tokens] + torch.mm(tile, gate_weight.t(), out=gate) + up = torch.mm(tile, up_weight.t()) + _silu_mul_inplace(gate, up) + torch.mm(gate, down_weight.t(), out=output[start:end]) + return output.reshape(hidden_states.shape) + + +def _verify_edge_tiles( + output: torch.Tensor, + hidden_states: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, +) -> bool: + from diffusers.models.autoencoders.ltx2_diffusion_decoder import _SWIGLU_TILE_SIZE + + dim = hidden_states.shape[-1] + flat = hidden_states.reshape(-1, dim) + actual = output.reshape(-1, dim) + tile_rows = _tile_rows(flat.shape[0], gate_weight.shape[0]) + starts = {0, ((flat.shape[0] - 1) // tile_rows) * tile_rows} + for start in starts: + end = min(start + tile_rows, flat.shape[0]) + # Match upstream's token tiling, not the optimized GEMM shape: changing + # the matmul row count can change its reduction order and BF16 output. + for reference_start in range(start, end, _SWIGLU_TILE_SIZE): + reference_end = min(reference_start + _SWIGLU_TILE_SIZE, end) + if not torch.equal( + actual[reference_start:reference_end], + _reference_tile(flat[reference_start:reference_end], gate_weight, up_weight, down_weight), + ): + return False + return True + + +def try_swiglu_tiled_exact( + hidden_states: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, +) -> torch.Tensor | None: + """Run workspace-reusing tiled SwiGLU, or expose the original module.""" + + if not _supported_inputs(hidden_states, gate_weight, up_weight, down_weight): + return None + rows = hidden_states.numel() // hidden_states.shape[-1] + tile_rows = _tile_rows(rows, gate_weight.shape[0]) + runtime_key = ( + hidden_states.device.index, + hidden_states.shape[-1], + gate_weight.shape[0], + tile_rows, + (rows - 1) % tile_rows + 1, + ) + if runtime_key in _FAILED_KEYS: + return None + try: + output = _launch(hidden_states, gate_weight, up_weight, down_weight) + if runtime_key not in _VERIFIED_KEYS: + if not _verify_edge_tiles(output, hidden_states, gate_weight, up_weight, down_weight): + _FAILED_KEYS.add(runtime_key) + logger.warning( + "Disabling LTX DiffVAE tiled SwiGLU fusion on %s after a bit-exactness mismatch", + hidden_states.device, + ) + return None + _VERIFIED_KEYS.add(runtime_key) + except Exception as exc: # noqa: BLE001 - fail closed after optimized-path failure + _FAILED_KEYS.add(runtime_key) + logger.warning( + "Disabling LTX DiffVAE tiled SwiGLU fusion on %s after failure: %s", + hidden_states.device, + exc, + ) + return None + return output + + +__all__ = ["try_swiglu_tiled_exact"] diff --git a/vllm_omni/diffusion/models/ltx2/vae/__init__.py b/vllm_omni/diffusion/models/ltx2/vae/__init__.py new file mode 100644 index 00000000000..9381766fc29 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/vae/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Ordinary LTX-2.5 DiffVAE, tiled decoding and acceleration operators.""" diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder.py b/vllm_omni/diffusion/models/ltx2/vae/decoder.py similarity index 57% rename from vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder.py rename to vllm_omni/diffusion/models/ltx2/vae/decoder.py index 73db2ea6b67..31257899007 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder.py +++ b/vllm_omni/diffusion/models/ltx2/vae/decoder.py @@ -7,19 +7,24 @@ """LTX-2.5 Native-weight adapter for Diffusers' diffusion VAE decoder.""" +import math from collections.abc import Mapping +from functools import cache from typing import Any import torch from diffusers.models.autoencoders import ( LTX2VideoDiffusionDecoderModel as DiffusersLTX2VideoDiffusionDecoderModel, ) +from diffusers.models.autoencoders import ltx2_diffusion_decoder as diffusers_decoder from diffusers.models.autoencoders.ltx2_diffusion_decoder import ( LTX2VideoDiffusionDecoder3d as DiffusersLTX2VideoDiffusionDecoder3d, ) from diffusers.models.autoencoders.vae import DecoderOutput from safetensors import safe_open +from ..ops import is_ltx2_fna_eligible, is_ltx2_fusion_eligible, resolve_ltx2_vae_operators + LTX25_NATIVE_DIFFUSION_DECODER_REPO_ID = "Lightricks/LTX-2.5" LTX25_NATIVE_ARTIFACT_REVISION = "8a4ff96f581e72bedc1b44367581c49d544a05f1" LTX25_NATIVE_DIFFUSION_DECODER_FILENAME = "vae/ltx-2.5-video-vae-bf16.safetensors" @@ -131,6 +136,177 @@ def load_ltx25_native_diffusion_decoder_state_dict(path: str) -> dict[str, torch return convert_ltx25_native_diffusion_decoder_state_dict(native_state_dict) +@cache +def _load_natten_token_permutation(): + from kernels import get_kernel + + natten = get_kernel("shi-labs/natten", version=1) + return natten.token_permute.token_permute_operation, natten.token_permute.token_unpermute_operation + + +class LTX2VideoVaeNeighborhoodNattenProcessor(diffusers_decoder.LTX2VideoVaeNeighborhoodNattenProcessor): + """NATTEN dispatch with Omni FNA for eligible stage-5 grids.""" + + def __call__( + self, attn: "LTX2VideoVaeNeighborhoodAttention", hidden_states: torch.Tensor, block_mask=None + ) -> torch.Tensor: + batch_size, num_frames, height, width, channels = hidden_states.shape + query, key, value = attn.project_qkv(hidden_states) + # NATTEN's CUTLASS kernels silently produce wrong output for non-contiguous inputs. + query, key, value = query.contiguous(), key.contiguous(), value.contiguous() + operators = resolve_ltx2_vae_operators(hidden_states.device) if is_ltx2_fna_eligible(hidden_states) else None + # `scale=1.0`: the query is already scaled in `project_qkv`, as in the reference. + if ( + tuple(attn.kernel_size) == (11, 11, 11) + and hidden_states.dtype == torch.bfloat16 + and channels == 256 + and attn.head_dim == 64 + and operators is not None + and operators.fna is not None + # The native metadata encodes KV tile indices in 15 bits. + and math.prod((n + tile - 1) // tile for n, tile in zip((num_frames, height, width), (4, 4, 8))) <= 1 << 15 + ): + token_permute, token_unpermute = _load_natten_token_permutation() + query_permuted, token_shape, _ = token_permute(query, (4, 4, 4), flip_tiled_dims=True) + key_permuted, _, _ = token_permute(key, (4, 4, 8), flip_tiled_dims=True) + value_permuted, _, _ = token_permute(value, (4, 4, 8), flip_tiled_dims=True) + output = operators.fna( + query_permuted, + key_permuted, + value_permuted, + shape=(num_frames, height, width), + window=tuple(attn.kernel_size), + ) + hidden_states = token_unpermute( + output, + token_layout_shape=token_shape, + tile_shape=(4, 4, 4), + flip_tiled_dims=True, + ) + else: + hidden_states = self._na3d( + query, + key, + value, + kernel_size=attn.kernel_size, + scale=1.0, + backend=self.backend, + ) + hidden_states = hidden_states.reshape(batch_size, num_frames, height, width, channels) + return attn.to_out[0](hidden_states) + + +class LTX2VideoVaeNeighborhoodAttention(diffusers_decoder.LTX2VideoVaeNeighborhoodAttention): + _available_processors = [ + *diffusers_decoder.LTX2VideoVaeNeighborhoodAttention._available_processors, + LTX2VideoVaeNeighborhoodNattenProcessor, + ] + + def project_qkv(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch_size, num_frames, height, width, _ = hidden_states.shape + shape = (batch_size, num_frames, height, width, self.heads, self.head_dim) + query = self.to_q(hidden_states).view(shape) + key = self.to_k(hidden_states).view(shape) + value = self.to_v(hidden_states).view(shape) + + operators = resolve_ltx2_vae_operators(query.device) if is_ltx2_fusion_eligible(query) else None + optimized = ( + None + if operators is None + else operators.qk_norm_rope( + query, + key, + self.norm_q.weight, + self.norm_k.weight, + self.norm_q.eps, + self.scale, + self.rope.rope_dim_split, + self.rope.base, + ) + ) + if optimized is not None: + query, key = optimized + return query, key, value + + query = self.norm_q(query) + key = self.norm_k(key) + query = query * self.scale + return self.rope(query), self.rope(key), value + + +class LTX2VideoVaeSwiGLU(diffusers_decoder.LTX2VideoVaeSwiGLU): + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + operators = resolve_ltx2_vae_operators(hidden_states.device) if is_ltx2_fusion_eligible(hidden_states) else None + optimized = ( + None + if operators is None + else operators.swiglu(hidden_states, self.w_gate.weight, self.w_up.weight, self.w_down.weight) + ) + return super().forward(hidden_states) if optimized is None else optimized + + +class LTX2VideoVaeDiffusionNABlock(diffusers_decoder.LTX2VideoVaeDiffusionNABlock): + def forward( + self, + hidden_states: torch.Tensor, + latent_context: torch.Tensor, + modulation: tuple[torch.Tensor, ...], + block_mask=None, + ) -> torch.Tensor: + if hidden_states.dtype is not torch.bfloat16 or not is_ltx2_fusion_eligible(hidden_states): + return super().forward(hidden_states, latent_context, modulation, block_mask) + operators = resolve_ltx2_vae_operators(hidden_states.device) + assert operators is not None + + scale_msa, shift_msa, _, scale_mlp, shift_mlp, _, _ = [ + modulation[i] + self.scale_shift_table[i].view(1, 1, 1, 1, -1) for i in range(self.num_mod_params) + ] + context_output = self.context_proj(latent_context) + attention_input = operators.residual_norm( + hidden_states, + context_output, + None, + self.norm1.weight, + scale_msa, + shift_msa, + self.norm1.eps, + ) + if attention_input is None: + hidden_states = hidden_states + context_output + hidden_states = hidden_states + self.attn( + self.norm1(hidden_states) * (1 + scale_msa) + shift_msa, block_mask + ) + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states) * (1 + scale_mlp) + shift_mlp) + return hidden_states + + attention_output = self.attn(attention_input, block_mask) + mlp_input = operators.residual_norm( + hidden_states, + context_output, + attention_output, + self.norm2.weight, + scale_mlp, + shift_mlp, + self.norm2.eps, + ) + if mlp_input is None: + residual = hidden_states + context_output + residual = residual + attention_output + mlp_input = self.norm2(residual) * (1 + scale_mlp) + shift_mlp + mlp_output = self.mlp(mlp_input) + output = operators.residual_add( + hidden_states, + context_output, + attention_output, + mlp_output, + ) + if output is not None: + return output + residual = hidden_states + context_output + residual = residual + attention_output + return residual + mlp_output + + class LTX2VideoDiffusionDecoder3d(DiffusersLTX2VideoDiffusionDecoder3d): """Diffusers decoder core with the short-clip NATTEN context fix.""" @@ -200,6 +376,15 @@ def __init__( # to this behavior-only subclass preserves parameters and state-dict keys. self.decoder.__class__ = LTX2VideoDiffusionDecoder3d self.decoder.stage5_kernel = tuple(decoder_stage5_kernel) + overrides = { + diffusers_decoder.LTX2VideoVaeNeighborhoodAttention: LTX2VideoVaeNeighborhoodAttention, + diffusers_decoder.LTX2VideoVaeSwiGLU: LTX2VideoVaeSwiGLU, + diffusers_decoder.LTX2VideoVaeDiffusionNABlock: LTX2VideoVaeDiffusionNABlock, + } + # Preserve every module and parameter object; only execution methods change. + for module in self.decoder.modules(): + if replacement := overrides.get(type(module)): + module.__class__ = replacement @classmethod def from_ltx25_native_checkpoint( diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder_distributed.py b/vllm_omni/diffusion/models/ltx2/vae/distributed.py similarity index 99% rename from vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder_distributed.py rename to vllm_omni/diffusion/models/ltx2/vae/distributed.py index 7f486dbdf79..9a98ba73998 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder_distributed.py +++ b/vllm_omni/diffusion/models/ltx2/vae/distributed.py @@ -25,7 +25,7 @@ TileTask, ) -from .ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel +from .decoder import LTX2VideoDiffusionDecoderModel logger = logging.get_logger(__name__) # pylint: disable=invalid-name From ecdd762ff57640b2007ed84e9f681c18a872e28d Mon Sep 17 00:00:00 2001 From: Anjie Hou <149605198+specture724@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:40:32 +0800 Subject: [PATCH 27/93] [2/N] Add a minimal temporal chunk callback for MiniMax-H3 (#7017) Signed-off-by: specture724 Co-authored-by: Claude Sonnet 5 --- .../minimax_h3/test_minimax_h3_vae_chunks.py | 357 ++++++++++++++++++ .../autoencoders/autoencoder_kl_wan.py | 6 +- vllm_omni/diffusion/models/interface.py | 8 + .../models/minimax_h3/chunked_decode.py | 85 +++++ .../models/minimax_h3/temporal_chunks.py | 122 ++++++ vllm_omni/diffusion/models/minimax_h3/vae.py | 90 +++-- 6 files changed, 641 insertions(+), 27 deletions(-) create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/chunked_decode.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/temporal_chunks.py diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py new file mode 100644 index 00000000000..fa7410b46de --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_chunks.py @@ -0,0 +1,357 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import gc +import tempfile +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist + +from vllm_omni.diffusion.models.minimax_h3.chunked_decode import decode_h3_chunks +from vllm_omni.diffusion.models.minimax_h3.temporal_chunks import decode_temporal_chunks + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +@pytest.fixture(scope="module") +def cpu_process_group(): + if dist.is_initialized(): + yield dist.group.WORLD + return + + with tempfile.NamedTemporaryFile(prefix="h3_chunk_dist_") as rendezvous: + init_method = f"file://{rendezvous.name}" + dist.init_process_group("gloo", rank=0, world_size=1, init_method=init_method) + try: + yield dist.group.WORLD + finally: + dist.destroy_process_group() + gc.collect() + + +class _FakeTemporalModel: + use_3d_conv = True + token_drop = 3 + tokens_chunk_size = 5 + token_overlap = 2 + vae_ratio_t = 4 + frame_pre_padding = 3 + frame_overlap = 5 + isolated_first_frame = False + isolated_last_frame = False + + def _decode_temporal_output_frame_plan(self, z, z_head, z_tail, num_chunks, pad_tokens): + del z, z_head, z_tail, num_chunks, pad_tokens + return 35, 0, 35 + + def _adaptive_decode(self, clip): + value = float(clip[:, :, 0].mean()) + return torch.full((1, 1, 24, 2, 2), value) + + @staticmethod + def blend(overlap, part, frame_overlap, dim): + del overlap, frame_overlap, dim + return part + + +class _FakeHost: + def __init__(self): + self.model = _FakeTemporalModel() + + @staticmethod + def _denormalize_latent(latent): + return latent + + @staticmethod + def _normalize_decoded_frames(frames): + return frames.float() + + +def test_temporal_chunks_emit_ordered_frames_and_collect_when_unconsumed(): + model = _FakeTemporalModel() + latent = torch.arange(8, dtype=torch.float32).view(1, 1, 8, 1, 1) + chunks = [] + marker = decode_temporal_chunks(model, latent, chunks.append) + + assert marker.shape == (0,) + assert [chunk.shape[2] for chunk in chunks] == [17, 17, 1] + assert torch.equal(torch.cat(chunks, dim=2), decode_temporal_chunks(model, latent, None)) + + +def test_h3_callback_failure_is_deferred_until_temporal_decode_finishes(): + host = _FakeHost() + latent = torch.zeros(1, 1, 8, 1, 1) + seen = [] + + def fail_once(frames): + seen.append(frames.shape[2]) + raise RuntimeError("sink failed") + + with pytest.raises(RuntimeError, match="sink failed"): + decode_h3_chunks(host, latent, fail_once, group=None) + assert seen == [17] + + +def test_h3_without_callback_is_a_plain_full_decode_on_a_vae_group(cpu_process_group): + host = _FakeHost() + latent = torch.arange(8, dtype=torch.float32).view(1, 1, 8, 1, 1) + + full = decode_h3_chunks(host, latent, None, group=cpu_process_group) + + chunks = [] + decode_h3_chunks(host, latent, chunks.append, group=cpu_process_group) + assert torch.equal(torch.cat(chunks, dim=2), full) + + +class _SimulatedGroup: + """Stand in for a VAE group so one process can check every rank's verdict. + + ``all_reduce`` is replaced with the sum the real collective would produce + for a group of ``world_size`` ranks, ``supplied`` of which handed in a + callback. That sum is what makes a rank-local verdict observable from a + single process. + """ + + def __init__(self, world_size: int, supplied: int): + self.world_size = world_size + self.supplied = supplied + + +def _patch_simulated_ranks(monkeypatch, simulated: _SimulatedGroup, rank: int) -> None: + from vllm_omni.diffusion.models.minimax_h3 import chunked_decode as mod + + def fake_all_reduce(tensor, group=None): + del group # the simulated census below stands in for the real reduction + tensor.copy_(torch.tensor([simulated.supplied], dtype=tensor.dtype)) + + monkeypatch.setattr(mod.dist, "get_rank", lambda _group: rank) + monkeypatch.setattr(mod.dist, "get_world_size", lambda _group: simulated.world_size) + monkeypatch.setattr(mod.dist, "all_reduce", fake_all_reduce) + monkeypatch.setattr(mod.dist, "broadcast", lambda tensor, src=0, group=None: None) + monkeypatch.setattr(mod.dist, "get_global_rank", lambda _group, _rank: 0) + + +@pytest.mark.parametrize("rank", [0, 1, 2]) +def test_h3_every_rank_rejects_a_partially_supplied_callback(monkeypatch, rank): + """Ranks that disagree must all refuse, not let one decode alone. + + A rank-local verdict let a peer fall through into the decoder collectives + while its peers raised, hanging the stage instead of failing it. + """ + host = _FakeHost() + latent = torch.zeros(1, 1, 8, 1, 1) + # Two ranks of three handed in a callback. + group = _SimulatedGroup(world_size=3, supplied=2) + _patch_simulated_ranks(monkeypatch, group, rank) + + with pytest.raises(ValueError, match="every rank"): + decode_h3_chunks(host, latent, lambda _frames: None, group=object()) + + +@pytest.mark.parametrize("rank", [0, 1, 2]) +def test_h3_only_rank_zero_publishes_when_every_rank_supplies(monkeypatch, rank): + """Ownership is positional: all ranks stream, rank 0 publishes.""" + host = _FakeHost() + latent = torch.zeros(1, 1, 8, 1, 1) + group = _SimulatedGroup(world_size=3, supplied=3) + _patch_simulated_ranks(monkeypatch, group, rank) + + seen: list[int] = [] + decode_h3_chunks(host, latent, lambda f: seen.append(int(f.shape[2])), group=object()) + + # Peers run the loop to stay in the collectives but publish nothing. + assert (seen != []) is (rank == 0) + + +def _stub_video_vae(*, parallel_size: int, tile_count: int): + """A MiniMaxH3VideoVAE with only the pieces the decode entry points read.""" + from contextlib import contextmanager + + from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE + + vae = object.__new__(MiniMaxH3VideoVAE) + vae.parallel_size = parallel_size + vae.model = SimpleNamespace(_adaptive_decode=lambda clip: clip) + vae.entered_rank_local = False + + @contextmanager + def rank_local(): + vae.entered_rank_local = True + yield + + vae._decoder_tile_count = lambda latent: tile_count + vae._rank_local_tiling = rank_local + vae.is_distributed_enabled = lambda: False + return vae + + +def test_chunked_decode_falls_back_when_tiles_are_fewer_than_ranks(monkeypatch): + """Without the fallback, tileless ranks raise while their peers block.""" + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_mod + + vae = _stub_video_vae(parallel_size=4, tile_count=1) + monkeypatch.setattr(vae_mod, "decode_h3_chunks", lambda *args, **kwargs: torch.zeros(1)) + + vae.decode_with_chunks(torch.zeros(1, 1, 8, 1, 1), on_chunk=lambda _frames: None) + + assert vae.entered_rank_local, "chunked decode must share decode_latent's too-few-tiles fallback" + + +def test_chunked_decode_keeps_shared_tiling_when_every_rank_has_a_tile(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_mod + + vae = _stub_video_vae(parallel_size=4, tile_count=8) + monkeypatch.setattr(vae_mod, "decode_h3_chunks", lambda *args, **kwargs: torch.zeros(1)) + + vae.decode_with_chunks(torch.zeros(1, 1, 8, 1, 1), on_chunk=lambda _frames: None) + + assert not vae.entered_rank_local + + +def test_h3_vae_declares_the_chunked_decode_capability(): + """A consumer must be able to detect the capability and the pixel range.""" + from vllm_omni.diffusion.models.interface import supports_chunked_vae_decode + from vllm_omni.diffusion.models.minimax_h3.vae import MiniMaxH3VideoVAE + + vae = object.__new__(MiniMaxH3VideoVAE) + + assert supports_chunked_vae_decode(vae) + # H3 reverts through the checkpoint's processor, which lands in [0, 1]; + # a Wan VAE publishes [-1, 1], so the range cannot be assumed. + assert vae.chunk_value_range == (0.0, 1.0) + + +def _native_decode_temporal_streaming(model, z, z_head, z_tail, num_chunks, pad_tokens): + """The released ``_decode_temporal_streaming`` loop, as the parity oracle. + + Transcribed from the checkpoint's ``AutoencoderKLLegacy`` so the fork can be + compared against the control flow it mirrors, including the isolated + head/tail extraction. + """ + total_frames, pad_frames, output_frames = model._decode_temporal_output_frame_plan( + z, z_head, z_tail, num_chunks, pad_tokens + ) + chunk_dec = model.tokens_chunk_size * model.vae_ratio_t + split_count = int(model.token_drop > 0) + 1 + dec = None + dec_overlap = None + write_pos = 0 + + def write_part(part): + nonlocal dec, write_pos + part_frames = int(part.shape[2]) + if part_frames <= 0: + return + if dec is None: + out_shape = list(part.shape) + out_shape[2] = output_frames + dec = torch.empty(out_shape, dtype=part.dtype, device=part.device) + remaining = int(dec.shape[2]) - write_pos + copy_frames = min(part_frames, max(0, remaining)) + if copy_frames > 0: + dec[:, :, write_pos : write_pos + copy_frames].copy_(part[:, :, :copy_frames]) + write_pos += copy_frames + + for i in range(num_chunks): + t_start = i * model.tokens_chunk_size + clip_z = z[:, :, t_start : t_start + model.tokens_chunk_size + model.token_overlap] + if i == 0 and z_head is not None: + clip_z = torch.cat([z_head, clip_z], dim=2) + if i == num_chunks - 1 and z_tail is not None: + clip_z = torch.cat([clip_z, z_tail], dim=2) + + clip_dec = model._adaptive_decode(clip_z) + + dec_tail = None + if i == 0 and z_head is not None: + write_part(clip_dec[:, :, model.vae_ratio_t - 1 : model.vae_ratio_t]) + clip_dec = clip_dec[:, :, model.vae_ratio_t :] + if i == num_chunks - 1 and z_tail is not None: + dec_tail = clip_dec[:, :, -1:] + clip_dec = clip_dec[:, :, : -model.vae_ratio_t] + + for j in range(split_count): + f_start = j * chunk_dec + f_end = min(f_start + chunk_dec, clip_dec.shape[2]) + chunk = clip_dec[:, :, f_start:f_end][:, :, model.frame_pre_padding :] + if j == 0: + if dec_overlap is not None: + chunk = model.blend(dec_overlap, chunk, model.frame_overlap, dim=-3) + dec_overlap = None + write_part(chunk) + else: + dec_overlap = chunk.contiguous() + + if i == num_chunks - 1: + if dec_overlap is not None: + write_part(dec_overlap) + dec_overlap = None + if dec_tail is not None: + write_part(dec_tail) + + return dec + + +class _IsolatedFrameModel(_FakeTemporalModel): + """Fake decoder with the isolated boundaries the released config can set.""" + + frame_pre_padding = 0 + + def __init__(self, *, first: bool, last: bool): + self.isolated_first_frame = first + self.isolated_last_frame = last + + def _decode_temporal_output_frame_plan(self, z, z_head, z_tail, num_chunks, pad_tokens): + del pad_tokens + frames = num_chunks * self.tokens_chunk_size * self.vae_ratio_t + frames += int(z_head is not None) + int(z_tail is not None) + del z + return frames, 0, frames + + def _adaptive_decode(self, clip): + # Every latent token decodes to a distinct, ordered block, so a shifted + # or reordered emission is visible in the values themselves. + base = float(clip[:, :, 0].mean()) + frames = int(clip.shape[2]) * self.vae_ratio_t + return torch.arange(frames, dtype=torch.float32).view(1, 1, frames, 1, 1) + base + + +@pytest.mark.parametrize( + ("first", "last"), + [(True, False), (False, True), (True, True)], +) +def test_temporal_chunks_match_the_native_loop_at_isolated_boundaries(first, last): + """The fork must extract the isolated head/tail the way the native loop does.""" + model = _IsolatedFrameModel(first=first, last=last) + latent = torch.arange(12, dtype=torch.float32).view(1, 1, 12, 1, 1) + + chunks: list[torch.Tensor] = [] + decode_temporal_chunks(model, latent, chunks.append) + streamed = torch.cat(chunks, dim=2) + + # Feed the oracle exactly what the released ``decode_temporal`` would: + # split the isolated tokens off, then pad the remainder to a whole chunk. + isolated = int(first) + int(last) + pseudo_tokens = int(latent.shape[2]) - isolated + model.token_drop + remainder = pseudo_tokens % model.tokens_chunk_size + pad_tokens = (model.tokens_chunk_size - remainder) if remainder else 0 + num_chunks = (pseudo_tokens + pad_tokens) // model.tokens_chunk_size - int(model.token_drop > 0) + + body = latent + z_head = None + if first: + z_head = body[:, :, :1] + body = body[:, :, 1:] + z_tail = None + if last: + z_tail = body[:, :, -1:] + body = body[:, :, :-1] + if pad_tokens: + body = torch.cat([body, body[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)], dim=2) + + expected = _native_decode_temporal_streaming(model, body, z_head, z_tail, num_chunks, pad_tokens) + + assert streamed.shape == expected.shape + assert torch.equal(streamed, expected) diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py index eddcb1ad650..9ab28d09c05 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from contextlib import nullcontext -from typing import Any +from typing import Any, ClassVar import torch import torch.distributed as dist @@ -62,6 +62,10 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True): return self._decode_temporal(z, return_dict=return_dict) + # Chunks are published straight from the decoder, clamped to the + # checkpoint's output range. + chunk_value_range: ClassVar[tuple[float, float]] = (-1.0, 1.0) + @apply_forward_hook def decode_with_chunks(self, z: torch.Tensor, *, on_chunk: DecodedChunkConsumer) -> None: """Decode ``z`` while synchronously delivering temporal chunks. diff --git a/vllm_omni/diffusion/models/interface.py b/vllm_omni/diffusion/models/interface.py index dac9a0b1d73..92216d19e92 100644 --- a/vllm_omni/diffusion/models/interface.py +++ b/vllm_omni/diffusion/models/interface.py @@ -58,8 +58,16 @@ class SupportsChunkedVAEDecode(Protocol): Implementations must drain decoding before surfacing callback errors. In a distributed VAE, every rank invokes the method so collectives stay in lockstep, while ``on_chunk`` is called only on the rank that owns output. + + ``chunk_value_range`` is the closed interval the delivered floats occupy, + which differs per checkpoint: a Wan VAE emits ``(-1.0, 1.0)`` while a + MiniMax-H3 VAE reverts through its processor to ``(0.0, 1.0)``. A consumer + cannot quantize chunks to pixels without it, so it is part of the + capability rather than knowledge each consumer hard-codes per model. """ + chunk_value_range: ClassVar[tuple[float, float]] + def decode_with_chunks( self, z: torch.Tensor, diff --git a/vllm_omni/diffusion/models/minimax_h3/chunked_decode.py b/vllm_omni/diffusion/models/minimax_h3/chunked_decode.py new file mode 100644 index 00000000000..08ce0e64bbd --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/chunked_decode.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Frames-only callback coordination for MiniMax-H3 VAE decoding.""" + +from __future__ import annotations + +from collections.abc import Callable + +import torch +import torch.distributed as dist + +from .temporal_chunks import decode_temporal_chunks + +MiniMaxH3VideoChunkCallback = Callable[[torch.Tensor], None] + + +def decode_h3_chunks( + host, + latent: torch.Tensor, + callback: MiniMaxH3VideoChunkCallback | None, + *, + group: dist.ProcessGroup | None, +) -> torch.Tensor: + """Run the local temporal loop, synchronizing callback errors across ranks.""" + streaming = callback is not None + owner = streaming + if group is not None: + rank = dist.get_rank(group) + world_size = dist.get_world_size(group) + # Ownership is positional: every rank supplies the callback and rank 0 + # publishes. Reduce only whether the ranks agree that this decode + # streams, so the verdict is identical everywhere -- a rank-local + # verdict would let a peer fall through and enter the decoder + # collectives alone, hanging the stage instead of failing it. + census = torch.tensor([int(streaming)], dtype=torch.int64, device=latent.device) + dist.all_reduce(census, group=group) + supplied = int(census.item()) + if supplied not in (0, world_size): + raise ValueError( + "MiniMax-H3 chunk callback must be supplied on every rank of the VAE " + f"group or on none of them; got {supplied} of {world_size}" + ) + # No rank supplied a callback: every rank runs a plain full decode. + streaming = supplied == world_size + owner = streaming and rank == 0 + + error: BaseException | None = None + + def publish(raw: torch.Tensor) -> None: + nonlocal error + if not owner or error is not None: + return + try: + processor = getattr(host.model, "processor", None) + decoded = raw if processor is None else processor.revert_tensor(raw) + frames = host._normalize_decoded_frames(decoded).contiguous() + assert callback is not None + callback(frames) + except BaseException as exc: # noqa: BLE001 + error = exc + + # Peer ranks run the temporal loop without materializing a second video. + sink = publish if owner else ((lambda _frames: None) if streaming else None) + result = decode_temporal_chunks( + host.model, + host._denormalize_latent(latent), + sink, + ) + + if group is not None: + failed = torch.tensor([int(rank == 0 and error is not None)], dtype=torch.int32, device=latent.device) + dist.broadcast(failed, src=dist.get_global_rank(group, 0), group=group) + if int(failed.item()): + if rank == 0: + assert error is not None + raise error + raise RuntimeError("MiniMax-H3 video chunk callback failed on rank zero") + elif error is not None: + raise error + + if not result.numel(): + return result + processor = getattr(host.model, "processor", None) + decoded = result if processor is None else processor.revert_tensor(result) + return host._normalize_decoded_frames(decoded) diff --git a/vllm_omni/diffusion/models/minimax_h3/temporal_chunks.py b/vllm_omni/diffusion/models/minimax_h3/temporal_chunks.py new file mode 100644 index 00000000000..a5154525dc8 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/temporal_chunks.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Small compatibility loop for the released MiniMax-H3 temporal decoder.""" + +from __future__ import annotations + +from collections.abc import Callable + +import torch +import torch.nn as nn + + +def decode_temporal_chunks( + model: nn.Module, + latent: torch.Tensor, + callback: Callable[[torch.Tensor], None] | None, +) -> torch.Tensor: + """Decode H3 temporal clips and synchronously publish committed frames. + + The released VAE exposes the temporal clip primitives but materializes the + final tensor. This is the same loop with the write point replaced by a + frames-only callback. ``callback=None`` retains the complete-output path. + """ + if latent.ndim != 5 or not bool(getattr(model, "use_3d_conv", False)): + raise ValueError("MiniMax-H3 temporal chunk decode requires a rank-5 3D latent") + + token_drop = int(model.token_drop) + chunk_size = int(model.tokens_chunk_size) + overlap_tokens = int(model.token_overlap) + ratio_t = int(model.vae_ratio_t) + pre_padding = int(model.frame_pre_padding) + + isolated_first = bool(model.isolated_first_frame and pre_padding == 0) + isolated_last = bool(model.isolated_last_frame) + z_head = latent[:, :, :1] if isolated_first else None + z_tail = latent[:, :, -1:] if isolated_last else None + start = int(isolated_first) + stop = int(latent.shape[2]) - int(isolated_last) + latent = latent[:, :, start:stop] + + pseudo_tokens = int(latent.shape[2]) + token_drop + pad_tokens = (-pseudo_tokens) % chunk_size + if pad_tokens: + latent = torch.cat((latent, latent[:, :, -1:].repeat(1, 1, pad_tokens, 1, 1)), dim=2) + num_chunks = (pseudo_tokens + pad_tokens) // chunk_size - int(token_drop > 0) + if num_chunks <= 0: + raise ValueError("MiniMax-H3 temporal chunk plan is empty") + + total_frames, pad_frames, output_frames = model._decode_temporal_output_frame_plan( + latent, z_head, z_tail, num_chunks, pad_tokens + ) + output_frames = int(output_frames) + if output_frames <= 0: + raise ValueError("MiniMax-H3 temporal chunk plan has no output frames") + + collected: list[torch.Tensor] = [] + overlap: torch.Tensor | None = None + written = 0 + + def emit(part: torch.Tensor) -> None: + nonlocal written + frames = int(part.shape[2]) + if frames <= 0 or written >= output_frames: + return + frames = min(frames, output_frames - written) + part = part[:, :, :frames] + if callback is None: + collected.append(part) + else: + callback(part) + written += frames + + chunk_frames = chunk_size * ratio_t + split_count = int(token_drop > 0) + 1 + for index in range(num_chunks): + begin = index * chunk_size + end = begin + chunk_size + overlap_tokens + clip = latent[:, :, begin:end] + if index == 0 and z_head is not None: + clip = torch.cat((z_head, clip), dim=2) + if index == num_chunks - 1 and z_tail is not None: + clip = torch.cat((clip, z_tail), dim=2) + + decoded = model._adaptive_decode(clip) + + # The isolated boundary latents were prepended/appended to the clip, so + # the decoder returns their frames inside this clip's block. Publish + # them the way the released loop does and drop their temporal block + # before splitting, or every later frame shifts by one block. + decoded_tail = None + if index == 0 and z_head is not None: + emit(decoded[:, :, ratio_t - 1 : ratio_t]) + decoded = decoded[:, :, ratio_t:] + if index == num_chunks - 1 and z_tail is not None: + decoded_tail = decoded[:, :, -1:] + decoded = decoded[:, :, :-ratio_t] + + for split in range(split_count): + begin = split * chunk_frames + end = min(begin + chunk_frames, int(decoded.shape[2])) + part = decoded[:, :, begin:end] + part = part[:, :, pre_padding:] + if split == 0: + if overlap is not None: + part = model.blend(overlap, part, int(model.frame_overlap), dim=-3) + overlap = None + emit(part) + else: + overlap = part.contiguous() + + if index == num_chunks - 1: + if overlap is not None: + emit(overlap) + overlap = None + if decoded_tail is not None: + emit(decoded_tail) + + if written != output_frames: + raise RuntimeError(f"MiniMax-H3 temporal decode emitted {written}/{output_frames} frames") + if callback is not None: + return latent.new_empty((0,)) + return torch.cat(collected, dim=2) if collected else latent.new_empty((0,)) diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index e986594099d..ca8359904ea 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -9,7 +9,7 @@ from collections.abc import Iterator from contextlib import AbstractContextManager, contextmanager, nullcontext from pathlib import Path -from typing import Any +from typing import Any, ClassVar import torch import torch.distributed as dist @@ -22,11 +22,13 @@ DistributedVaeMixin, ) from vllm_omni.diffusion.distributed.parallel_state import get_world_group +from vllm_omni.diffusion.models.interface import DecodedChunkConsumer from vllm_omni.diffusion.offloader.module_residency import ( BoundedAllocatorCache, PinnedModuleStager, ) +from .chunked_decode import decode_h3_chunks from .ops import install_h3_vae_optimizations from .packed_tokens import minimax_h3_patchify_video_latent @@ -377,25 +379,16 @@ def encode_video( ).float() return rows, shape - @torch.inference_mode() - def decode_latent(self, latent: torch.Tensor) -> torch.Tensor: - channels = int(self.config_dict["latent_channels"]) - mean = torch.tensor( - self.config_dict["latents_mean"], - device=latent.device, - dtype=latent.dtype, - ).view(1, channels, 1, 1, 1) - std = torch.tensor( - self.config_dict["latents_std"], - device=latent.device, - dtype=latent.dtype, - ).view(1, channels, 1, 1, 1) - # The checkpoint hands rank r the tiles ``range(r, num_tiles, sp_size)`` - # and then rejects an empty share inside the gather. A rank with no - # tiles raises and leaves the collective while the others block in it - # forever, so too few tiles hangs the whole stage rather than failing - # it. Tile count depends only on the latent shape, so every rank takes - # this branch together. + def _decode_tiling_context(self, latent: torch.Tensor) -> AbstractContextManager: + """Pick the tiling mode a decode of ``latent`` can safely use. + + The checkpoint hands rank r the tiles ``range(r, num_tiles, sp_size)`` + and then rejects an empty share inside the gather. A rank with no + tiles raises and leaves the collective while the others block in it + forever, so too few tiles hangs the whole stage rather than failing + it. Tile count depends only on the latent shape, so every rank takes + this branch together. + """ num_tiles = self._decoder_tile_count(latent) if self.parallel_size > 1 and num_tiles < self.parallel_size: logger.warning_once( @@ -405,19 +398,64 @@ def decode_latent(self, latent: torch.Tensor) -> torch.Tensor: num_tiles, self.parallel_size, ) - tiling_context: AbstractContextManager = self._rank_local_tiling() - else: - tiling_context = nullcontext() + return self._rank_local_tiling() + return nullcontext() + + @torch.inference_mode() + def decode_latent(self, latent: torch.Tensor) -> torch.Tensor: + with self._decode_tiling_context(latent): + decoded = self.model.decode_base(self._denormalize_latent(latent)) + return self._normalize_decoded_frames(self.model.processor.revert_tensor(decoded)) - with tiling_context: - decoded = self.model.decode_base(latent * std + mean) - frames = self.model.processor.revert_tensor(decoded) + def _denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor: + channels = int(self.config_dict["latent_channels"]) + mean = torch.tensor(self.config_dict["latents_mean"], device=latent.device, dtype=latent.dtype) + std = torch.tensor(self.config_dict["latents_std"], device=latent.device, dtype=latent.dtype) + shape = (1, channels, 1, 1, 1) + return latent * std.view(shape) + mean.view(shape) + + @staticmethod + def _normalize_decoded_frames(decoded: torch.Tensor) -> torch.Tensor: + """Canonicalize remote H3 decoder output to [B,C,T,H,W].""" + frames = decoded if frames.ndim == 4: frames = frames.unsqueeze(0).transpose(1, 2) if frames.ndim != 5: raise ValueError(f"unexpected decoded video shape {tuple(frames.shape)}") return frames.float() + # Chunks are reverted through the checkpoint's processor, which + # denormalizes and clamps into the unit interval. + chunk_value_range: ClassVar[tuple[float, float]] = (0.0, 1.0) + + @torch.inference_mode() + def decode_with_chunks(self, z: torch.Tensor, *, on_chunk: DecodedChunkConsumer) -> None: + """Decode temporal clips and synchronously publish frames-only chunks. + + Implements :class:`SupportsChunkedVAEDecode`. Every rank participating + in distributed VAE execution must invoke this method with a callback so + the temporal collectives stay in lockstep; ``on_chunk`` is called only + on the rank that owns output. Chunks arrive as ``[B, C, T, H, W]`` + float frames, normalized through the checkpoint's processor to match + the complete decode path. After a callback failure, the remaining + chunks are decoded and discarded before the exception is re-raised. + """ + if not callable(on_chunk): + raise TypeError("on_chunk must be callable") + if not callable(getattr(self.model, "_adaptive_decode", None)): + raise RuntimeError("Loaded MiniMax-H3 VAE does not expose temporal decode primitives") + group = None + if self.is_distributed_enabled(): + group = self._native_parallel_state().get("sp_process_group") + if group is None or dist.get_world_size(group) != self.parallel_size: + raise RuntimeError("MiniMax-H3 VAE chunk decode has an invalid spatial-parallel group") + # Native H3 tiling performs its own collectives for every temporal clip, + # so this path needs the same too-few-tiles fallback as the complete + # decode: without it a shape that leaves some ranks tileless hangs the + # gather instead of decoding rank-locally. + with self._decode_tiling_context(z): + decode_h3_chunks(self, z, on_chunk, group=group) + class MiniMaxH3AudioVAE(nn.Module): def __init__( From cdd4638b17f86db33a3f798dba468a7aef7bc76b Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Wed, 9 Sep 2026 17:58:42 +0800 Subject: [PATCH 28/93] [Feature][Diffusion] Expose detailed pipeline timings (#6822) Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- .../models/minimax_h3/test_minimax_h3_offload.py | 4 +++- tests/diffusion/test_diffusion_engine_metrics.py | 1 + tests/metrics/test_modality.py | 2 ++ vllm_omni/diffusion/diffusion_engine.py | 8 +++++++- .../diffusion/models/minimax_h3/pipeline_minimax_h3.py | 2 ++ vllm_omni/metrics/stats.py | 1 + 6 files changed, 16 insertions(+), 2 deletions(-) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py index dade4437cfa..865b2248a10 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py @@ -13,12 +13,14 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] -def test_h3_profiler_targets_cover_aggregate_reference_encoders(): +def test_h3_profiler_targets_cover_reference_encoders_and_vae_decoders(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline targets = MiniMaxH3Pipeline._PROFILER_TARGETS assert "_encode_visual_conditions" in targets assert "_encode_reference_audio_conditions" in targets + assert "video_vae.decode_latent" in targets + assert "audio_vae.decode_latent" in targets assert "_encode_video_conditions" not in targets assert "_encode_video_audio_conditions" not in targets assert "_encode_audio_conditions" not in targets diff --git a/tests/diffusion/test_diffusion_engine_metrics.py b/tests/diffusion/test_diffusion_engine_metrics.py index dc909a2b920..1dad9facef9 100644 --- a/tests/diffusion/test_diffusion_engine_metrics.py +++ b/tests/diffusion/test_diffusion_engine_metrics.py @@ -140,6 +140,7 @@ def test_timing_metric_key_naming_consistency(self) -> None: raise AssertionError("diffusion_engine_total_time_ms should be attached in step_streaming()") assert '"diffusion_engine_exec_time_ms": exec_total_time * 1000' in step_streaming_source + assert '"output_ready_wait_time_ms": output_ready_wait_time * 1000' in step_streaming_source # step_total_ms is no longer emitted as a metric (no family consumes # it); it lives only in the debug log breakdown below. assert '"diffusion_engine_total_time_ms"' not in step_streaming_source diff --git a/tests/metrics/test_modality.py b/tests/metrics/test_modality.py index 694808a3123..14bb4f309ca 100644 --- a/tests/metrics/test_modality.py +++ b/tests/metrics/test_modality.py @@ -532,6 +532,7 @@ def test_accumulator_converts_all_engine_milliseconds_to_seconds(self): metrics={ "preprocess_time_ms": 50.0, "diffusion_engine_exec_time_ms": 1500.0, + "output_ready_wait_time_ms": 250.0, "postprocess_time_ms": 20.0, "vae_decode_time_ms": 300.0, "forward_time_ms": 800.0, @@ -545,6 +546,7 @@ def test_accumulator_converts_all_engine_milliseconds_to_seconds(self): assert dict(aggregator.diffusion_metrics["req-units"]) == { "preprocess_time_s": pytest.approx(0.05), "diffusion_engine_exec_time_s": pytest.approx(1.5), + "output_ready_wait_time_s": pytest.approx(0.25), "postprocess_time_s": pytest.approx(0.02), "vae_decode_time_s": pytest.approx(0.3), "forward_time_s": pytest.approx(0.8), diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 52a22c0b973..c36ac0459db 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -439,8 +439,10 @@ async def step_streaming(self, request: OmniDiffusionRequest) -> AsyncGenerator[ generator = self.get_output_stream(request_id) async for output in generator: exec_total_time = time.perf_counter() - exec_start_time + output_ready_wait_time = 0.0 # Async mode: wait for background D2H/SHM to complete. if output.async_output_id: + output_ready_wait_start_time = time.perf_counter() fut = self.executor.wait_output_ready(output.async_output_id) timeout = _async_output_timeout() try: @@ -455,6 +457,7 @@ async def step_streaming(self, request: OmniDiffusionRequest) -> AsyncGenerator[ describe(output.async_output_id) if describe else "unavailable", ) raise + output_ready_wait_time = time.perf_counter() - output_ready_wait_start_time postprocess_start_time = time.perf_counter() scheduler_metrics = diffusion_scheduler_waiting_metrics(getattr(self, "_scheduler_num_waiting_reqs", 0)) try: @@ -468,9 +471,11 @@ async def step_streaming(self, request: OmniDiffusionRequest) -> AsyncGenerator[ step_total_ms = (time.perf_counter() - diffusion_engine_start_time) * 1000 logger.debug( "DiffusionEngine.step_streaming breakdown: preprocess=%.2f ms, " - "add_req_and_wait=%.2f ms, postprocess=%.2f ms, total=%.2f ms", + "add_req_and_wait=%.2f ms, output_ready_wait=%.2f ms, " + "postprocess=%.2f ms, total=%.2f ms", preprocess_time * 1000, exec_total_time * 1000, + output_ready_wait_time * 1000, postprocess_time * 1000, step_total_ms, ) @@ -478,6 +483,7 @@ async def step_streaming(self, request: OmniDiffusionRequest) -> AsyncGenerator[ metrics_update = { "preprocess_time_ms": preprocess_time * 1000, "diffusion_engine_exec_time_ms": exec_total_time * 1000, + "output_ready_wait_time_ms": output_ready_wait_time * 1000, "postprocess_time_ms": postprocess_time * 1000, **scheduler_metrics, } diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index 794d2dc098f..ff1f8a3e5e7 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -683,6 +683,8 @@ class MiniMaxH3Pipeline( "_encode_reference_audio_conditions", "diffuse", "decode", + "video_vae.decode_latent", + "audio_vae.decode_latent", "prepare_encode", "denoise_step", "post_decode", diff --git a/vllm_omni/metrics/stats.py b/vllm_omni/metrics/stats.py index 9e9191eaa1b..ef86a5a3599 100644 --- a/vllm_omni/metrics/stats.py +++ b/vllm_omni/metrics/stats.py @@ -656,6 +656,7 @@ def stage_postprocess_timer(self, stage_id: int, req_id: Any): _MS_TO_S: dict[str, str] = { "preprocess_time_ms": "preprocess_time_s", "diffusion_engine_exec_time_ms": "diffusion_engine_exec_time_s", + "output_ready_wait_time_ms": "output_ready_wait_time_s", "postprocess_time_ms": "postprocess_time_s", "vae_decode_time_ms": "vae_decode_time_s", "forward_time_ms": "forward_time_s", From 038c9b948fbf10e99d3167c18a6a22c20049f7f5 Mon Sep 17 00:00:00 2001 From: NumberWan Date: Wed, 9 Sep 2026 18:20:34 +0800 Subject: [PATCH 29/93] [Bugfix] Resolve #6931 hub FA3 on torch 2.13 via kernels 0.16.1 (#7185) Signed-off-by: NumberWan --- .../attention_backends/huggingface_hub.md | 2 +- recipes/Boogu/Boogu-Image.md | 2 +- requirements/common.txt | 10 ++- tests/diffusion/attention/test_kernels_hub.py | 74 +++++++++++++++++ .../models/ltx2/test_ltx25_pipeline.py | 4 +- tests/e2e/accuracy/test_qwen_image.py | 83 ++++++++----------- .../attention/backends/flash_attn_hub.py | 6 +- .../diffusion/models/ltx2/ltx2_components.py | 4 +- vllm_omni/platforms/cuda/platform.py | 20 +++++ 9 files changed, 142 insertions(+), 63 deletions(-) diff --git a/docs/user_guide/diffusion/attention_backends/huggingface_hub.md b/docs/user_guide/diffusion/attention_backends/huggingface_hub.md index 574ff625a0b..cfd220e829c 100644 --- a/docs/user_guide/diffusion/attention_backends/huggingface_hub.md +++ b/docs/user_guide/diffusion/attention_backends/huggingface_hub.md @@ -12,7 +12,7 @@ implementations to reduce numerical drift and sampling divergence. ## Installation ```bash -pip install kernels==0.15.2 +pip install kernels==0.16.1 ``` If `kernels` is unavailable, an explicit `FLASH_ATTN_HUB` or diff --git a/recipes/Boogu/Boogu-Image.md b/recipes/Boogu/Boogu-Image.md index 0d832b0ae93..0353b977802 100644 --- a/recipes/Boogu/Boogu-Image.md +++ b/recipes/Boogu/Boogu-Image.md @@ -152,7 +152,7 @@ The following configuration was validated: - vLLM-Omni version or commit: Use a commit that contains TorchAO FP8 checkpoint loading for diffusion models - TorchAO version: 0.17.0 -- `kernels` / `kernels-data`: 0.15.2 / 0.16.1 +- `kernels` / `kernels-data`: 0.16.1 / matching kernels-data #### Command diff --git a/requirements/common.txt b/requirements/common.txt index bb79e0a6a9a..855fb6b0a30 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -9,13 +9,15 @@ # That breakage is now handled in code instead of by pinning: # MossTTSNanoForGeneration wraps weight loading in # transformers_keys_to_ignore_compat (vllm_omni/model_executor/models/utils.py). -transformers >= 5.10.1, < 5.15 +# Floor is 5.13.0 so transformers.integrations.hub_kernels LayerRepository +# entries always pass version=/revision= required by kernels>=0.16 (#6971). +transformers >= 5.13.0, < 5.15 av>=14.0.0 omegaconf>=2.3.0 diffusers==0.40.0 -# Hub kernels, including LTX-2.5 NATTEN. Keep this pin in sync with the -# supported transformers range above. -kernels==0.15.2 +# Hub kernels, including LTX-2.5 NATTEN. 0.16+ resolves torch-stable-abi* +# FA3 builds on torch 2.13 / cu130; keep in sync with transformers above. +kernels==0.16.1 safetensors>=0.8.0 accelerate==1.12.0 soundfile>=0.13.1 diff --git a/tests/diffusion/attention/test_kernels_hub.py b/tests/diffusion/attention/test_kernels_hub.py index 3f7451ca3f3..c85fb58f3a8 100644 --- a/tests/diffusion/attention/test_kernels_hub.py +++ b/tests/diffusion/attention/test_kernels_hub.py @@ -13,6 +13,12 @@ from vllm_omni.platforms import current_omni_platform +def _clear_hub_module_cache(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.attention.backends import flash_attn_hub + + monkeypatch.setattr(flash_attn_hub, "_hub_modules", {}) + + @pytest.mark.core_model @pytest.mark.cpu def test_explicit_kernels_hub_selection_does_not_fallback(monkeypatch: pytest.MonkeyPatch): @@ -54,6 +60,74 @@ def test_explicit_kernels_hub_selection_does_not_fallback(monkeypatch: pytest.Mo CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64) +@pytest.mark.core_model +@pytest.mark.cpu +def test_explicit_kernels_hub_variant_miss_does_not_fallback(monkeypatch: pytest.MonkeyPatch): + """Explicit Hub backends must fail loudly when no build variant resolves (#6971).""" + from vllm.platforms.interface import DeviceCapability + + from vllm_omni.diffusion.envs import PACKAGES_CHECKER + from vllm_omni.platforms.cuda.platform import CudaOmniPlatform + + monkeypatch.setattr( + CudaOmniPlatform, + "get_device_capability", + classmethod(lambda cls, device_id=0: DeviceCapability(9, 0)), + ) + monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True}) + _clear_hub_module_cache(monkeypatch) + + kernels_module = types.ModuleType("kernels") + + def _missing_variant(repo_id, version=None, **kwargs): + raise FileNotFoundError(f"Cannot find a build variant for this system in {repo_id}") + + kernels_module.get_kernel = _missing_variant # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "kernels", kernels_module) + + with pytest.raises(RuntimeError, match="no compatible build"): + CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64) + with pytest.raises(RuntimeError, match="no compatible build"): + CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_explicit_kernels_hub_preflight_uses_v2_when_v1_misses(monkeypatch: pytest.MonkeyPatch): + """Preflight must share _load_hub_module's (1, 2) policy, not pin version=1.""" + from vllm.platforms.interface import DeviceCapability + + from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum + from vllm_omni.diffusion.envs import PACKAGES_CHECKER + from vllm_omni.platforms.cuda.platform import CudaOmniPlatform + + monkeypatch.setattr( + CudaOmniPlatform, + "get_device_capability", + classmethod(lambda cls, device_id=0: DeviceCapability(9, 0)), + ) + monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True}) + _clear_hub_module_cache(monkeypatch) + + kernels_module = types.ModuleType("kernels") + versions_tried: list[int | None] = [] + + def _v1_miss_v2_hit(repo_id, version=None, **kwargs): + versions_tried.append(version) + if version == 1: + raise FileNotFoundError(f"Cannot find a build variant for this system in {repo_id}") + return types.SimpleNamespace(name=f"{repo_id}@v{version}") + + kernels_module.get_kernel = _v1_miss_v2_hit # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "kernels", kernels_module) + + fa3_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64) + assert fa3_path == DiffusionAttentionBackendEnum.FLASH_ATTN_3_HUB.get_path() + fa2_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64) + assert fa2_path == DiffusionAttentionBackendEnum.FLASH_ATTN_HUB.get_path() + assert versions_tried == [1, 2, 1, 2] + + @pytest.mark.core_model @pytest.mark.cpu def test_explicit_flash_attention_unavailable_does_not_fallback(monkeypatch: pytest.MonkeyPatch): diff --git a/tests/diffusion/models/ltx2/test_ltx25_pipeline.py b/tests/diffusion/models/ltx2/test_ltx25_pipeline.py index ae90ea1a369..a96fa8a4eb2 100644 --- a/tests/diffusion/models/ltx2/test_ltx25_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx25_pipeline.py @@ -186,7 +186,7 @@ def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatc ltx2_components.initialize_pipeline_components(pipe, od_config) assert str(exc_info.value) == ( - "LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.10.1,<5.15." + "LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.13.0,<5.15." ) @@ -387,7 +387,7 @@ def __init__(self): MissingNattenProcessor, ) - with pytest.raises(RuntimeError, match=r"kernels==0\.15\.2.*supported GPU.*allow Hub access"): + with pytest.raises(RuntimeError, match=r"kernels==0\.16\.1.*supported GPU.*allow Hub access"): ltx2_components._create_ltx25_natten_processor() diff --git a/tests/e2e/accuracy/test_qwen_image.py b/tests/e2e/accuracy/test_qwen_image.py index a09014c8d2b..45fd20934da 100644 --- a/tests/e2e/accuracy/test_qwen_image.py +++ b/tests/e2e/accuracy/test_qwen_image.py @@ -34,11 +34,7 @@ SEED = 42 SSIM_THRESHOLD = 0.94 # With --fa-deterministic, H100 nightlies stably land near SSIM 0.958 / PSNR 27.8 -# vs Diffusers (see #5734 / #5963), so the primary gates are 0.94 / 27.0. -# The SDPA-matched fallback (hub kernel unavailable for the image's torch) -# landed at SSIM 0.9687 in build 2954 — keep a slightly tighter SSIM gate for -# that path only. -SSIM_THRESHOLD_SDPA_FALLBACK = 0.96 +# vs Diffusers on the FA3 hub path (see #5734 / #5963 / #6971). PSNR_THRESHOLD = 27.0 MODEL_2512_ID = "Qwen/Qwen-Image-2512" @@ -65,44 +61,6 @@ P99_ABS_DIFF_THRESHOLD_2512 = 4e-1 -_FLASH_ATTN3_HUB_AVAILABLE: bool | None = None - - -def _flash_attn3_hub_available() -> bool: - """Whether kernels-community/flash-attn3 has a build for this torch/CUDA. - - The hub repo only ships prebuilt variants for specific torch versions; - a base-image torch bump can outrun it (build 2952 failed with "Cannot - find a build variant for this system"). When it is unavailable, both - sides of the comparison fall back to torch SDPA together — the test - compares implementations, so the attention math must stay matched. - """ - global _FLASH_ATTN3_HUB_AVAILABLE - if _FLASH_ATTN3_HUB_AVAILABLE is None: - try: - from kernels import get_kernel - - get_kernel("kernels-community/flash-attn3") - _FLASH_ATTN3_HUB_AVAILABLE = True - except Exception as exc: - print(f"kernels-community/flash-attn3 unavailable ({exc}); using matched torch SDPA on both sides.") - _FLASH_ATTN3_HUB_AVAILABLE = False - return _FLASH_ATTN3_HUB_AVAILABLE - - -def _omni_server_env() -> dict[str, str] | None: - if _flash_attn3_hub_available(): - return None - return {"DIFFUSION_ATTENTION_BACKEND": "TORCH_SDPA"} - - -def _set_reference_attention_backend(pipe: DiffusionPipeline) -> None: - if _flash_attn3_hub_available(): - pipe.transformer.set_attention_backend("_flash_3_hub") - else: - pipe.transformer.set_attention_backend("native") - - def _model_name() -> str: return os.environ.get(MODEL_ENV_VAR, MODEL_ID) @@ -115,7 +73,22 @@ def _local_files_only(model: str) -> bool: return Path(model).exists() +_OMNI_FA3_HUB_BACKEND = "FLASH_ATTN_3_HUB" +_OMNI_FA3_HUB_ENV = {"DIFFUSION_ATTENTION_BACKEND": _OMNI_FA3_HUB_BACKEND} + + +def _assert_omni_fa3_hub_resolves() -> None: + from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum + from vllm_omni.platforms.cuda.platform import CudaOmniPlatform + + path = CudaOmniPlatform.get_diffusion_attn_backend_cls(_OMNI_FA3_HUB_BACKEND, head_size=128) + expected = DiffusionAttentionBackendEnum.FLASH_ATTN_3_HUB.get_path() + assert path == expected, path + print(f"Resolved Omni attention backend: {_OMNI_FA3_HUB_BACKEND} -> {path}") + + def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> Image.Image: + _assert_omni_fa3_hub_resolves() server_args = [ "--num-gpus", "1", @@ -124,8 +97,10 @@ def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> Image.Image: "--init-timeout", "900", "--fa-deterministic", + "--diffusion-attention-backend", + _OMNI_FA3_HUB_BACKEND, ] - with OmniServer(model, server_args, use_omni=True, env_dict=_omni_server_env()) as omni_server: + with OmniServer(model, server_args, use_omni=True, env_dict=_OMNI_FA3_HUB_ENV) as omni_server: response = requests.post( f"http://{omni_server.host}:{omni_server.port}/v1/images/generations", json={ @@ -161,7 +136,7 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image: trust_remote_code=True, local_files_only=_local_files_only(model), ).to("cuda") - _set_reference_attention_backend(pipe) + pipe.transformer.set_attention_backend("_flash_3_hub") generator = torch.Generator(device="cuda").manual_seed(SEED) result = pipe( # pyright: ignore[reportCallIssue] prompt=PROMPT, @@ -186,8 +161,18 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image: def _run_vllm_omni_qwen_image_2512(*, model: str, output_path: Path) -> Image.Image: - server_args = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"] - with OmniServer(model, server_args, use_omni=True, env_dict=_omni_server_env()) as omni_server: + _assert_omni_fa3_hub_resolves() + server_args = [ + "--num-gpus", + "1", + "--stage-init-timeout", + "300", + "--init-timeout", + "900", + "--diffusion-attention-backend", + _OMNI_FA3_HUB_BACKEND, + ] + with OmniServer(model, server_args, use_omni=True, env_dict=_OMNI_FA3_HUB_ENV) as omni_server: response = requests.post( f"http://{omni_server.host}:{omni_server.port}/v1/images/generations", json={ @@ -223,7 +208,7 @@ def _run_diffusers_qwen_image_2512(*, model: str, output_path: Path) -> Image.Im trust_remote_code=True, local_files_only=_local_files_only(model), ).to("cuda") - _set_reference_attention_backend(pipe) + pipe.transformer.set_attention_backend("_flash_3_hub") generator = torch.Generator(device="cuda").manual_seed(SEED_2512) result = pipe( # pyright: ignore[reportCallIssue] prompt=PROMPT_2512, @@ -262,7 +247,7 @@ def test_qwen_image_matches_diffusers(accuracy_artifact_root: Path) -> None: diffusers_image=diffusers_output, width=WIDTH, height=HEIGHT, - ssim_threshold=SSIM_THRESHOLD if _flash_attn3_hub_available() else SSIM_THRESHOLD_SDPA_FALLBACK, + ssim_threshold=SSIM_THRESHOLD, psnr_threshold=PSNR_THRESHOLD, ) diff --git a/vllm_omni/diffusion/attention/backends/flash_attn_hub.py b/vllm_omni/diffusion/attention/backends/flash_attn_hub.py index 7f9c78bf74c..b24c9969492 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn_hub.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn_hub.py @@ -28,11 +28,9 @@ def _load_hub_module(repo_id: str): logger.info("Loading %s kernel from HuggingFace Hub...", repo_id) last_error = None - for version in (1, 2, None): + for version in (1, 2): try: - if version is not None: - return get_kernel(repo_id, version=version) - return get_kernel(repo_id) + return get_kernel(repo_id, version=version) except Exception as exc: logger.info("Failed to load %s version %s: %s", repo_id, version, exc) last_error = exc diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py index f1dcc38bb2d..5826f273e64 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_components.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_components.py @@ -298,7 +298,7 @@ def _create_ltx25_natten_processor() -> LTX2VideoVaeNeighborhoodNattenProcessor: except (ImportError, OSError, RuntimeError, ValueError) as exc: raise RuntimeError( "LTX-2.5 DiffVAE requires the shi-labs/natten Hub kernel. " - "Install kernels==0.15.2, use a supported GPU, leave " + "Install kernels==0.16.1, use a supported GPU, leave " "DIFFUSERS_DISABLE_REMOTE_CODE unset, and allow Hub access during kernel initialization." ) from exc @@ -662,7 +662,7 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: revision=revision, ) if profile.text_encoder_cls is None: - raise ImportError("LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.10.1,<5.15.") + raise ImportError("LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.13.0,<5.15.") with torch.device("cpu"): pipeline.text_encoder = _load_component( profile.text_encoder_cls, diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index c635877a92d..25124b4e0b9 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/vllm_omni/platforms/cuda/platform.py @@ -231,6 +231,26 @@ def get_diffusion_attn_backend_cls( "with compute capability 9.x. Select a compatible backend." ) + if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"): + # Same loader/cache as execution (Hub versions 1 then 2). A missing + # torch/CUDA variant used to degrade silently to native/SDPA (#6971). + hub_repo = ( + "kernels-community/flash-attn3" + if backend_upper == "FLASH_ATTN_3_HUB" + else "kernels-community/flash-attn2" + ) + try: + from vllm_omni.diffusion.attention.backends.flash_attn_hub import _get_hub_module + + _get_hub_module(hub_repo) + except Exception as e: + raise RuntimeError( + f"{backend_upper} was explicitly selected, but {hub_repo} has no " + "compatible build for this torch/CUDA (or failed to load). " + "Install kernels>=0.16.1 so stable-ABI variants resolve, or select " + "a different attention backend." + ) from e + if backend_upper == "FLASH_ATTN" and not flash_attn_supported: if is_blackwell: reason = ( From 6e79d582e3c4a4da4a159904125f3de7a6cbae2b Mon Sep 17 00:00:00 2001 From: zyz111222 Date: Wed, 9 Sep 2026 19:03:09 +0800 Subject: [PATCH 30/93] [Bugfix][Ascend] fix npu 310/a5 bugs (#6685) Signed-off-by: zouyizhou --- tests/platforms/npu/test_310p_patches.py | 69 ++--- tests/platforms/npu/test_a5_patches.py | 259 ++++++++++++++++++ .../models/qwen3_tts/prompt_embeds_builder.py | 10 +- vllm_omni/platforms/npu/_310p/__init__.py | 7 +- .../platforms/npu/_310p/patch/qwen3_tts.py | 45 +-- vllm_omni/platforms/npu/__init__.py | 16 ++ .../{qwen3_tts_code2wav.py => qwen3_tts.py} | 59 +++- vllm_omni/platforms/npu/platform.py | 6 +- 8 files changed, 381 insertions(+), 90 deletions(-) create mode 100644 tests/platforms/npu/test_a5_patches.py rename vllm_omni/platforms/npu/models/{qwen3_tts_code2wav.py => qwen3_tts.py} (54%) diff --git a/tests/platforms/npu/test_310p_patches.py b/tests/platforms/npu/test_310p_patches.py index 08cbb2adc8d..c19dcd99782 100644 --- a/tests/platforms/npu/test_310p_patches.py +++ b/tests/platforms/npu/test_310p_patches.py @@ -431,9 +431,17 @@ def maybe_trans_nz(weight: torch.Tensor) -> torch.Tensor: trans_calls.append(weight.detach().clone()) return weight + 1.0 - def npu_rotary_mul(hidden_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: - rotary_calls.append((hidden_states.detach().clone(), cos.detach().clone(), sin.detach().clone())) - return hidden_states + cos - sin + def npu_apply_rotary_pos_emb( + query: torch.Tensor, + key: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + *, + layout: str = "BSND", + rotary_mode: str = "half", + ): + rotary_calls.append((query.detach().clone(), key.detach().clone(), cos.detach().clone(), sin.detach().clone())) + return query + cos - sin, key + cos - sin def npu_flash_attention(**kwargs) -> None: flash_calls.append( @@ -443,7 +451,7 @@ def npu_flash_attention(**kwargs) -> None: monkeypatch.setattr(module, "maybe_trans_nz", maybe_trans_nz) monkeypatch.setattr(module, "aligned_16", lambda tensor: tensor) - monkeypatch.setattr(module.torch_npu, "npu_rotary_mul", npu_rotary_mul, raising=False) + monkeypatch.setattr(module.torch_npu, "npu_apply_rotary_pos_emb", npu_apply_rotary_pos_emb, raising=False) monkeypatch.setattr(module.torch_npu, "_npu_flash_attention", npu_flash_attention, raising=False) config = SimpleNamespace( @@ -503,7 +511,12 @@ def npu_flash_attention(**kwargs) -> None: expected_v_f = expected_v.transpose(1, 2).reshape(3, attention.num_kv_heads, attention.head_dim) expected_flash = (expected_q_f + expected_k_f + expected_v_f).to(torch.float16).to(torch.float32) - assert len(rotary_calls) == 2 + # One fused aclnn rope call covers q and k in the 4D BSND layout. + assert len(rotary_calls) == 1 + assert rotary_calls[0][0].shape == (1, 3, attention.num_heads, attention.head_dim) + assert rotary_calls[0][1].shape == (1, 3, attention.num_kv_heads, attention.head_dim) + assert rotary_calls[0][2].shape == (1, 3, 1, attention.head_dim) + assert rotary_calls[0][3].shape == (1, 3, 1, attention.head_dim) assert len(flash_calls) == 1 torch.testing.assert_close(flash_calls[0]["query"], expected_q_f) torch.testing.assert_close(flash_calls[0]["key"], expected_k_f) @@ -535,39 +548,16 @@ def test_qwen3_tts_talker_patch_uses_fp16_runtime_dtype(monkeypatch: pytest.Monk assert codes[0].shape == (4, 2) -def test_qwen3_tts_prompt_patch_runs_stft_frontend_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None: +def test_qwen3_tts_prompt_patch_sets_cpu_mel_front_end_flag(monkeypatch: pytest.MonkeyPatch) -> None: module, _ = _load_qwen3_tts_patch(monkeypatch) - captured = {} - - def fake_mel_spectrogram(wav_tensor, **kwargs): - captured["wav_device"] = wav_tensor.device - captured["wav_dtype"] = wav_tensor.dtype - captured["kwargs"] = kwargs - return torch.ones(1, 128, 3, dtype=torch.float32) - - class FakeSpeakerEncoder(torch.nn.Module): - def __init__(self): - super().__init__() - self.param = torch.nn.Parameter(torch.zeros(1, dtype=torch.float16)) - - def forward(self, mels): - captured["speaker_input_dtype"] = mels.dtype - return (torch.ones(4, dtype=mels.dtype),) - monkeypatch.setattr(module.prompt_embeds_builder, "mel_spectrogram", fake_mel_spectrogram) - builder = object.__new__(module._Qwen3TTSPromptEmbedsBuilder310P) - builder._device = lambda: torch.device("cpu") - builder._embedding_dtype = torch.float16 - builder._speaker_encoder = FakeSpeakerEncoder() - builder._config = SimpleNamespace(speaker_encoder_config=SimpleNamespace(sample_rate=24000)) - - speaker = builder.extract_speaker_embedding(np.zeros(16, dtype=np.float32), 24000) - - assert captured["wav_device"] == torch.device("cpu") - assert captured["wav_dtype"] is torch.float32 - assert captured["kwargs"]["sampling_rate"] == 24000 - assert captured["speaker_input_dtype"] is torch.float16 - assert speaker.dtype is torch.float16 + # The 310P builder keeps the base implementation and only opts into the + # CPU mel front-end, because the NPU has no torch.stft. + assert issubclass( + module._Qwen3TTSPromptEmbedsBuilder310P, + module.prompt_embeds_builder.Qwen3TTSPromptEmbedsBuilder, + ) + assert module._Qwen3TTSPromptEmbedsBuilder310P._mel_spectrogram_on_cpu is True def test_qwen3_tts_tokenizer_npu_patch_dispatches_fused_ops(monkeypatch: pytest.MonkeyPatch) -> None: @@ -706,11 +696,10 @@ def load_weights(self, weights): _install_fake_module(monkeypatch, "vllm_ascend.utils", maybe_trans_nz=maybe_trans_nz) _install_fake_module(monkeypatch, "vllm_omni") _install_fake_module(monkeypatch, "vllm_omni.platforms", current_omni_platform=current_platform) - _install_fake_module(monkeypatch, "vllm_omni.model_executor") - _install_fake_module(monkeypatch, "vllm_omni.model_executor.models") - _install_fake_module(monkeypatch, "vllm_omni.model_executor.models.qwen3_tts") + # Not A5: keep the FRACTAL_Z conv weight layout for 310P. + _install_fake_module(monkeypatch, "vllm_omni.platforms.npu", is_a5=lambda: False) - path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "models" / "qwen3_tts_code2wav.py" + path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "models" / "qwen3_tts.py" module = _load_source_module("vllm_omni_test_qwen3_tts_code2wav_npu_patch", path) module.apply_qwen3_tts_code2wav_patch() diff --git a/tests/platforms/npu/test_a5_patches.py b/tests/platforms/npu/test_a5_patches.py new file mode 100644 index 00000000000..0d6648235b0 --- /dev/null +++ b/tests/platforms/npu/test_a5_patches.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for A5 (Ascend 950 PR) patch wiring. + +The A5 logic lives in the shared NPU platform files (``npu/__init__.py`` for +device detection, ``npu/models/qwen3_tts.py`` for the Qwen3-TTS +patch). The tests load those modules from source with fake Qwen3-TTS +dependencies, so they validate the patch contract without loading real model +or NPU kernels. A5 device detection is exercised with a fake +``vllm_ascend.utils`` module. +""" + +from __future__ import annotations + +import importlib.util +import sys +import types +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _repo_root() -> Path: + marker = Path("vllm_omni") / "platforms" / "npu" / "models" + for parent in Path(__file__).resolve().parents: + if (parent / marker).is_dir(): + return parent + raise FileNotFoundError(f"could not locate repo root containing {marker}") + + +def _load_source_module(name: str, path: Path): + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot load import spec for {path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _install_fake_module(monkeypatch: pytest.MonkeyPatch, name: str, **attrs): + module = types.ModuleType(name) + for key, value in attrs.items(): + setattr(module, key, value) + monkeypatch.setitem(sys.modules, name, module) + return module + + +def _load_npu_init(monkeypatch: pytest.MonkeyPatch): + _install_fake_module(monkeypatch, "vllm") + _install_fake_module( + monkeypatch, + "vllm.logger", + init_logger=lambda _name: SimpleNamespace(debug=lambda *_: None), + ) + path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "__init__.py" + return _load_source_module("vllm_omni_test_npu_init", path) + + +def _load_qwen3_tts_code2wav_patch( + monkeypatch: pytest.MonkeyPatch, + *, + a5: bool, +): + class FakePromptEmbedsBuilder: + pass + + class FakeCode2Wav: + def __init__(self, *, vllm_config, prefix=""): + pass + + def load_weights(self, weights): + pass + + fake_prompt_builder = _install_fake_module( + monkeypatch, + "vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder", + Qwen3TTSPromptEmbedsBuilder=FakePromptEmbedsBuilder, + mel_spectrogram=lambda *_args, **_kwargs: torch.empty(0), + ) + fake_talker = _install_fake_module( + monkeypatch, + "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker", + Qwen3TTSPromptEmbedsBuilder=FakePromptEmbedsBuilder, + ) + _install_fake_module( + monkeypatch, + "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav", + Qwen3TTSCode2Wav=FakeCode2Wav, + ) + _install_fake_module(monkeypatch, "vllm_omni") + _install_fake_module(monkeypatch, "vllm_omni.model_executor") + _install_fake_module(monkeypatch, "vllm_omni.model_executor.models") + pkg_qwen3_tts = _install_fake_module( + monkeypatch, + "vllm_omni.model_executor.models.qwen3_tts", + prompt_embeds_builder=fake_prompt_builder, + qwen3_tts_talker=fake_talker, + ) + # Mark the package fakes as real packages so the patch's lazy submodule + # imports (e.g. ``qwen3_tts_code2wav``) resolve inside the harness. + pkg_qwen3_tts.__path__ = [] + _install_fake_module( + monkeypatch, + "vllm_omni.platforms", + current_omni_platform=SimpleNamespace(is_npu=lambda: False), + ) + _install_fake_module(monkeypatch, "vllm_omni.platforms.npu", is_a5=lambda: a5) + _install_fake_module(monkeypatch, "vllm") + _install_fake_module(monkeypatch, "vllm.config", VllmConfig=object) + _install_fake_module( + monkeypatch, + "vllm.logger", + init_logger=lambda _name: SimpleNamespace(info=lambda *_: None, debug=lambda *_: None), + ) + _install_fake_module(monkeypatch, "vllm_ascend") + _install_fake_module(monkeypatch, "vllm_ascend.utils", maybe_trans_nz=lambda weight: weight) + _install_fake_module(monkeypatch, "torch_npu", npu_format_cast=lambda weight, _fmt: weight) + + path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "models" / "qwen3_tts.py" + module = _load_source_module("vllm_omni_test_qwen3_tts_code2wav_patch", path) + return module, fake_prompt_builder, fake_talker + + +def test_is_a5_detects_a5_device(monkeypatch: pytest.MonkeyPatch) -> None: + module = _load_npu_init(monkeypatch) + _install_fake_module(monkeypatch, "vllm_ascend") + _install_fake_module(monkeypatch, "vllm_ascend.utils", is_950=lambda: True) + + assert module.is_a5() is True + + class NpuDevice: + type = "npu" + + assert module.is_a5(NpuDevice()) is True + assert module.is_a5(torch.device("cpu")) is False + + +def test_is_a5_falls_back_false_without_is_950(monkeypatch: pytest.MonkeyPatch) -> None: + module = _load_npu_init(monkeypatch) + # Older vllm-ascend releases have no `is_950`; detection must fall back. + _install_fake_module(monkeypatch, "vllm_ascend") + _install_fake_module(monkeypatch, "vllm_ascend.utils") + + assert module.is_a5(torch.device("cuda")) is False + assert module.is_a5() is False + + +def test_qwen3_tts_patch_swaps_prompt_builder_on_a5(monkeypatch: pytest.MonkeyPatch) -> None: + module, fake_prompt_builder, fake_talker = _load_qwen3_tts_code2wav_patch( + monkeypatch, + a5=True, + ) + original = fake_prompt_builder.Qwen3TTSPromptEmbedsBuilder + + module.apply_qwen3_tts_patches() + + assert fake_prompt_builder.Qwen3TTSPromptEmbedsBuilder is not original + assert fake_prompt_builder.Qwen3TTSPromptEmbedsBuilder is fake_talker.Qwen3TTSPromptEmbedsBuilder + + module.apply_qwen3_tts_patches() + assert fake_prompt_builder.Qwen3TTSPromptEmbedsBuilder is fake_talker.Qwen3TTSPromptEmbedsBuilder + + +def test_qwen3_tts_patch_keeps_prompt_builder_off_a5(monkeypatch: pytest.MonkeyPatch) -> None: + module, fake_prompt_builder, fake_talker = _load_qwen3_tts_code2wav_patch( + monkeypatch, + a5=False, + ) + original = fake_prompt_builder.Qwen3TTSPromptEmbedsBuilder + + module.apply_qwen3_tts_patches() + + assert fake_prompt_builder.Qwen3TTSPromptEmbedsBuilder is original + assert fake_talker.Qwen3TTSPromptEmbedsBuilder is original + + +def test_a5_prompt_patch_sets_cpu_mel_front_end_flag(monkeypatch: pytest.MonkeyPatch) -> None: + module, fake_prompt_builder, fake_talker = _load_qwen3_tts_code2wav_patch( + monkeypatch, + a5=True, + ) + original = fake_prompt_builder.Qwen3TTSPromptEmbedsBuilder + + module.apply_qwen3_tts_patches() + + # The A5 builder keeps the base implementation and only opts into the CPU + # mel front-end, because the NPU has no torch.stft. + builder_cls = fake_prompt_builder.Qwen3TTSPromptEmbedsBuilder + assert builder_cls is not original + assert builder_cls._mel_spectrogram_on_cpu is True + assert fake_talker.Qwen3TTSPromptEmbedsBuilder is builder_cls + + +def test_qwen3_tts_code2wav_patch_skips_fractal_z_conv_cast_on_a5( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module, _, _ = _load_qwen3_tts_code2wav_patch(monkeypatch, a5=True) + linear_weights = [] + conv_weights = [] + + monkeypatch.setattr( + module, + "maybe_trans_nz", + lambda weight: linear_weights.append(weight) or weight, + ) + monkeypatch.setattr( + module.torch_npu, + "npu_format_cast", + lambda weight, fmt: conv_weights.append((weight, fmt)) or weight, + ) + + class FakeDecoder(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(4, 4) + self.conv = torch.nn.Conv1d(4, 4, 3) + self.deconv = torch.nn.ConvTranspose1d(4, 4, 4) + self.grouped_conv = torch.nn.Conv1d(4, 4, 3, groups=2) + self.cache_precompute_calls = 0 + + def precompute_snake_caches(self): + self.cache_precompute_calls += 1 + + class FakeCode2Wav: + def __init__(self, *, vllm_config, prefix=""): + self.vllm_config = vllm_config + self.prefix = prefix + self.decoder = FakeDecoder() + + def _npu_decoder_runtime_dtype(self, _device): + return torch.float16 + + def load_weights(self, weights): + assert list(weights) == [] + return {"loaded"} + + target = _install_fake_module( + monkeypatch, + "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav", + Qwen3TTSCode2Wav=FakeCode2Wav, + ) + module.apply_qwen3_tts_code2wav_patch() + + model = target.Qwen3TTSCode2Wav( + vllm_config=SimpleNamespace(device_config=SimpleNamespace(device=torch.device("cpu"))), + prefix="stage1", + ) + assert model.load_weights(iter(())) == {"loaded"} + + assert model.prefix == "stage1" + assert model.decoder.linear.weight.dtype is torch.float16 + assert [weight.data_ptr() for weight in linear_weights] == [model.decoder.linear.weight.data_ptr()] + # On A5 convs keep the contiguous ND layout: no FRACTAL_Z cast is issued. + assert conv_weights == [] + assert model.decoder.cache_precompute_calls == 1 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 ddb55cd5171..609117ce3b6 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 @@ -643,6 +643,10 @@ def put_ref_audio_artifacts( # -------------------- speaker encoder / codec encoder -------------------- + # Some NPUs (310P, Ascend 950) do not support torch.stft; builders on + # such devices set this flag to compute the mel spectrogram on CPU. + _mel_spectrogram_on_cpu: bool = False + def extract_speaker_embedding(self, wav: np.ndarray, sr: int) -> torch.Tensor: # vLLM workers do not automatically move arbitrary torch.nn.Modules to # CUDA. Ensure the speaker encoder is on the same device/dtype as the @@ -666,7 +670,9 @@ def extract_speaker_embedding(self, wav: np.ndarray, sr: int) -> torch.Tensor: # Follow official implementation: mel_spectrogram expects 24kHz. Move # the waveform first so STFT/mel computation stays on the model device # instead of materializing a CPU mel tensor and copying it per request. - wav_tensor = torch.from_numpy(wav).to(device=dev, dtype=torch.float32).unsqueeze(0) + # Devices without torch.stft set _mel_spectrogram_on_cpu instead. + mel_device = torch.device("cpu") if self._mel_spectrogram_on_cpu else dev + wav_tensor = torch.from_numpy(wav).to(device=mel_device, dtype=torch.float32).unsqueeze(0) mels = mel_spectrogram( wav_tensor, n_fft=1024, @@ -677,7 +683,7 @@ def extract_speaker_embedding(self, wav: np.ndarray, sr: int) -> torch.Tensor: fmin=0, fmax=12000, ).transpose(1, 2) - spk = self._speaker_encoder(mels.to(dtype=dtype))[0] + spk = self._speaker_encoder(mels.to(device=dev, dtype=dtype))[0] return spk.to(dtype=dtype) def encode_ref_audio_batch( diff --git a/vllm_omni/platforms/npu/_310p/__init__.py b/vllm_omni/platforms/npu/_310p/__init__.py index 129cf8e2041..605a1aa53f5 100644 --- a/vllm_omni/platforms/npu/_310p/__init__.py +++ b/vllm_omni/platforms/npu/_310p/__init__.py @@ -6,12 +6,10 @@ from __future__ import annotations import torch -from vllm.logger import init_logger - -logger = init_logger(__name__) def is_310p(device: torch.device | None = None) -> bool: + """Return True on Ascend 310P devices, False otherwise.""" if device is not None and device.type != "npu": return False try: @@ -20,9 +18,6 @@ def is_310p(device: torch.device | None = None) -> bool: return bool(ascend_is_310p()) except (ImportError, AttributeError): return False - except Exception: - logger.debug("Failed to detect Ascend 310P device.", exc_info=True) - return False def disable_jit_compile() -> None: diff --git a/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py b/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py index dd03ba7d26b..4ec4b6b72d0 100644 --- a/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py +++ b/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py @@ -81,35 +81,9 @@ def _encode_ref_audio_batch( class _Qwen3TTSPromptEmbedsBuilder310P(prompt_embeds_builder.Qwen3TTSPromptEmbedsBuilder): - def extract_speaker_embedding(self, wav: np.ndarray, sr: int) -> torch.Tensor: - dev = self._device() - dtype = self._embedding_dtype - try: - spk_param = next(self._speaker_encoder.parameters()) - if spk_param.device != dev or spk_param.dtype != dtype: - self._speaker_encoder.to(device=dev, dtype=dtype) - except StopIteration: - pass - - target_sr = int(getattr(self._config.speaker_encoder_config, "sample_rate", 24000)) - if sr != target_sr: - resampler = self._get_resampler(int(sr), target_sr) - wav = resampler.resample(wav.astype(np.float32), orig_sr=int(sr)) - - # 310P does not support torch.stft on NPU. - wav_tensor = torch.from_numpy(wav).to(device=_CPU_DEVICE, dtype=torch.float32).unsqueeze(0) - mels = prompt_embeds_builder.mel_spectrogram( - wav_tensor, - n_fft=1024, - num_mels=128, - sampling_rate=24000, - hop_size=256, - win_size=1024, - fmin=0, - fmax=12000, - ).transpose(1, 2) - spk = self._speaker_encoder(mels.to(device=dev, dtype=dtype))[0] - return spk.to(dtype=dtype) + """Qwen3-TTS prompt-embeds builder specialized for the 310P NPU path.""" + + _mel_spectrogram_on_cpu = True # =================================================================== @@ -192,10 +166,15 @@ def forward( cos, sin = position_embeddings cos = cos.unsqueeze(1) sin = sin.unsqueeze(1) - # Use the fused Ascend RoPE op instead of expanding RoPE into - # elementwise mul/add/rotate-half kernels. - q = torch_npu.npu_rotary_mul(q, cos, sin) - k = torch_npu.npu_rotary_mul(k, cos, sin) + # Use the fused aclnn rope op, which is NPUGraph-capturable; it + # requires the 4D BSND layout with head_dim 128 or 64. + q = q.transpose(1, 2).contiguous() + k = k.transpose(1, 2).contiguous() + cos = cos.transpose(1, 2).contiguous() + sin = sin.transpose(1, 2).contiguous() + q, k = torch_npu.npu_apply_rotary_pos_emb(q, k, cos, sin, rotary_mode="half") + q = q.transpose(1, 2) + k = k.transpose(1, 2) real_tokens = int(bsz) * int(seq_len) output_dtype = q.dtype diff --git a/vllm_omni/platforms/npu/__init__.py b/vllm_omni/platforms/npu/__init__.py index 74dfc419ce6..3df9b5ca360 100644 --- a/vllm_omni/platforms/npu/__init__.py +++ b/vllm_omni/platforms/npu/__init__.py @@ -1,14 +1,30 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + from typing import TYPE_CHECKING +import torch + if TYPE_CHECKING: from vllm_omni.platforms.npu.platform import NPUOmniPlatform __all__ = ["NPUOmniPlatform"] +def is_a5(device: torch.device | None = None) -> bool: + """Return True on Ascend 950 (A5) devices, False otherwise.""" + if device is not None and device.type != "npu": + return False + try: + from vllm_ascend.utils import is_950 + + return bool(is_950()) + except (ImportError, AttributeError): + return False + + def __getattr__(name: str): if name != "NPUOmniPlatform": raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm_omni/platforms/npu/models/qwen3_tts_code2wav.py b/vllm_omni/platforms/npu/models/qwen3_tts.py similarity index 54% rename from vllm_omni/platforms/npu/models/qwen3_tts_code2wav.py rename to vllm_omni/platforms/npu/models/qwen3_tts.py index bd74e9dc6ad..a91d2cf80cd 100644 --- a/vllm_omni/platforms/npu/models/qwen3_tts_code2wav.py +++ b/vllm_omni/platforms/npu/models/qwen3_tts.py @@ -1,11 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Patch Qwen3-TTS Code2Wav NPU runtime setup and weight preparation.""" +"""Patch Qwen3-TTS NPU runtime setup and weights. -from __future__ import annotations +The code2wav runtime/weight fixes apply on every NPU. On A5 (Ascend 950 PR) +``torch.stft`` is unavailable, so the speaker-encoder mel front-end of the +prompt-embeds builder runs on CPU and the result is moved back to NPU +(same trade-off as the 310P path). +""" -from typing import TYPE_CHECKING +from __future__ import annotations import torch import torch.nn as nn @@ -14,12 +18,12 @@ from vllm.logger import init_logger from vllm_ascend.utils import maybe_trans_nz -if TYPE_CHECKING: - pass +from vllm_omni.platforms.npu import is_a5 logger = init_logger(__name__) _PATCHED = False +_A5_PATCHED = False _original_init = None _original_load_weights = None _ACL_FORMAT_FRACTAL_Z = 4 @@ -49,7 +53,12 @@ def _prepare_npu_decoder_weights(decoder: nn.Module) -> None: module.weight.data = maybe_trans_nz(module.weight.data) linear_count += 1 elif isinstance(module, (nn.Conv1d, nn.ConvTranspose1d)) and module.groups == 1: - module.weight.data = torch_npu.npu_format_cast(module.weight.data.contiguous(), _ACL_FORMAT_FRACTAL_Z) + # A5 (Ascend 950 PR) does not support the FRACTAL_Z conv + # weight layout, keep the contiguous ND layout there. + if not is_a5(): + module.weight.data = torch_npu.npu_format_cast( + module.weight.data.contiguous(), _ACL_FORMAT_FRACTAL_Z + ) conv_count += 1 logger.info("Prepared NPU Code2Wav weights: linear=%d conv=%d", linear_count, conv_count) @@ -67,6 +76,44 @@ def _patched_load_weights(self, weights): return loaded +def apply_qwen3_tts_patches() -> None: + """Install all Qwen3-TTS NPU patches. + + The code2wav runtime/weight patch applies on every NPU; the A5 + prompt-embeds builder swap is gated on ``is_a5()`` because it changes the + speaker-encoder mel front-end (``torch.stft`` is unsupported there). + """ + apply_qwen3_tts_code2wav_patch() + if not is_a5(): + return + _apply_a5_prompt_embeds_builder_patch() + + +def _apply_a5_prompt_embeds_builder_patch() -> None: + global _A5_PATCHED + if _A5_PATCHED: + return + + # Import and patch lazily so the Qwen3-TTS model modules are only loaded + # on A5 devices (same gate-then-import convention as the 310P patch). + from vllm_omni.model_executor.models.qwen3_tts import ( + prompt_embeds_builder, + qwen3_tts_talker, + ) + + class _Qwen3TTSPromptEmbedsBuilderA5(prompt_embeds_builder.Qwen3TTSPromptEmbedsBuilder): + """Qwen3-TTS prompt-embeds builder specialized for the A5 NPU path.""" + + _mel_spectrogram_on_cpu = True + + # The talker imports the builder by name at module scope, so its reference + # must be swapped as well (same as the 310P patch). + prompt_embeds_builder.Qwen3TTSPromptEmbedsBuilder = _Qwen3TTSPromptEmbedsBuilderA5 + qwen3_tts_talker.Qwen3TTSPromptEmbedsBuilder = _Qwen3TTSPromptEmbedsBuilderA5 + _A5_PATCHED = True + logger.debug("Applied A5 prompt-embeds builder patch for Qwen3-TTS") + + def apply_qwen3_tts_code2wav_patch() -> None: global _PATCHED, _original_init, _original_load_weights if _PATCHED: diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index cf65da414d4..2c8d42f6093 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -62,8 +62,8 @@ def __init__(self) -> None: from vllm_omni.platforms.npu.models.minicpmo_4_5_code2wav import ( apply_minicpmo_4_5_code2wav_patch, ) - from vllm_omni.platforms.npu.models.qwen3_tts_code2wav import ( - apply_qwen3_tts_code2wav_patch, + from vllm_omni.platforms.npu.models.qwen3_tts import ( + apply_qwen3_tts_patches, ) from vllm_omni.platforms.npu.models.qwen3_tts_tokenizer_v2 import ( apply_qwen3_tts_tokenizer_v2_patch, @@ -71,7 +71,7 @@ def __init__(self) -> None: adapt_patch(is_global_patch=True) apply_minicpmo_4_5_code2wav_patch() - apply_qwen3_tts_code2wav_patch() + apply_qwen3_tts_patches() apply_qwen3_tts_tokenizer_v2_patch() apply_310p_patches() From 7c1b1e4189d0dbd26a6ecdb05dc20426a6078a24 Mon Sep 17 00:00:00 2001 From: Zheng Wengang Date: Wed, 9 Sep 2026 20:25:41 +0800 Subject: [PATCH 31/93] [Bugfix][Engine] Group overlapping device stages into one sequential init component (#7328) Signed-off-by: ZhengWG Co-authored-by: amy-why-3459 --- .../test_async_omni_engine_stage_init.py | 5 +- tests/engine/test_parallel_stage_init.py | 136 +++++++++++++----- vllm_omni/engine/stage_init_utils.py | 35 +++++ vllm_omni/engine/stage_runtime.py | 98 ++++++------- 4 files changed, 189 insertions(+), 85 deletions(-) diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index b1771da2850..e6f7f8a8ba9 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -926,10 +926,7 @@ def test_remote_replicas_use_distinct_init_group_keys(): replica.launch_mode = "remote" replica.metadata.runtime_cfg = None - assert [runtime._replica_init_group_key(replica) for replica in plan.replicas] == [ - "remote:1:0", - "remote:1:1", - ] + assert runtime._init_group_keys(plan.replicas) == ["remote:1:0", "remote:1:1"] def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_replica_failure(monkeypatch): diff --git a/tests/engine/test_parallel_stage_init.py b/tests/engine/test_parallel_stage_init.py index 5952d710e85..d88d49825af 100644 --- a/tests/engine/test_parallel_stage_init.py +++ b/tests/engine/test_parallel_stage_init.py @@ -1,7 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Unit tests for parallel stage initialization (CPU-only). Covers the pure/lock logic added for ``parallel_stage_init``: - * init-group device keying (resolved physical set; parallel unique keys), + * init-group device keying (overlap components; parallel unique keys), * admission control arithmetic + graph reserve, * SH/EX device phase locks (real flock on a temp dir), * the phase-locked executor wrapper ordering, @@ -18,6 +21,7 @@ import os import subprocess import threading +import time import types from pathlib import Path @@ -37,7 +41,9 @@ LogicalStageInitPlan, ReplicaInitPlan, device_init_lock_path, + device_overlap_group_keys, open_device_lock_file, + parse_physical_device_ids, ) from vllm_omni.engine.stage_phase_lock import ( DeviceLockTimeoutError, @@ -46,6 +52,7 @@ wrap_executor_with_phase_locks, ) from vllm_omni.engine.stage_runtime import StageRuntime +from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -104,49 +111,114 @@ def _fake_vllm_config(util: float, *, model="m", dtype="bf16", cudagraph_mode="F # --------------------------------------------------------------------------- # # A1: init-group device key # --------------------------------------------------------------------------- # -def test_init_group_key_resolves_physical_devices_serial(): - runtime = _runtime(parallel_stage_init=False) - runtime._init_visible_devices_baseline = None # no CUDA_VISIBLE_DEVICES mapping - r0 = _llm_replica(0, 0, "0") - r1 = _llm_replica(1, 0, "1") - r0b = _llm_replica(2, 0, "0") +@pytest.mark.parametrize( + ("devices", "expected"), + [ + ("0", {0}), + ("1,0,1", {0, 1}), + (" 2, 3 ", {2, 3}), + (None, None), + ("", None), + ("GPU-uuid", None), + ], +) +def test_parse_physical_device_ids(devices, expected): + assert parse_physical_device_ids(devices) == expected + + +@pytest.mark.parametrize( + ("device_sets", "expected"), + [ + # equal / disjoint + ([{0}, {1}, {0}], ["device-group:0", "device-group:1", "device-group:0"]), + # unequal overlap, canonical order + ([{0, 1}, {0}, {2}], ["device-group:0,1", "device-group:0,1", "device-group:2"]), + # transitive: {0,1} ~ {1,2} ~ {2} + ([{0, 1}, {1, 2}, {2}], ["device-group:0,1,2"] * 3), + # unresolved may touch any GPU -> everything serial + ([{0}, None, {1}], ["device-group:*"] * 3), + ([], []), + ], +) +def test_device_overlap_group_keys(device_sets, expected): + sets = [None if s is None else frozenset(s) for s in device_sets] + assert device_overlap_group_keys(sets) == expected - k0 = runtime._replica_init_group_key(r0) - k1 = runtime._replica_init_group_key(r1) - k0b = runtime._replica_init_group_key(r0b) - assert k0 == "device:0" - assert k1 == "device:1" - # Different physical devices -> different groups (parallel). - assert k0 != k1 - # Same physical device -> same group (serialized under LOCK_EX). - assert k0 == k0b +def test_init_group_keys_serial_groups_by_device_overlap(): + runtime = _runtime(parallel_stage_init=False) + keys = runtime._init_group_keys( + [_llm_replica(0, 0, "0,1"), _llm_replica(1, 0, "0"), _llm_replica(2, 0, "2"), _llm_replica(3, 0, "1,0")] + ) -def test_init_group_key_parallel_is_unique_per_replica(): - runtime = _runtime(parallel_stage_init=True) - runtime._init_visible_devices_baseline = None - # Both on the SAME device, but parallel mode must give distinct keys so they - # run concurrently, coordinated by the child SH/EX locks + admission. - keys = [ - runtime._replica_init_group_key(_llm_replica(0, 0, "0")), - runtime._replica_init_group_key(_llm_replica(0, 1, "0")), - runtime._replica_init_group_key(_llm_replica(1, 0, "0")), - ] - assert keys == ["parallel:0:0", "parallel:0:1", "parallel:1:0"] - assert len(set(keys)) == 3 + assert keys == ["device-group:0,1", "device-group:0,1", "device-group:2", "device-group:0,1"] -def test_init_group_key_remote_and_diffusion_unchanged(): +def test_init_group_keys_fixed_for_parallel_remote_and_diffusion(): runtime = _runtime(parallel_stage_init=True) remote = _llm_replica(1, 0, "0") remote.launch_mode = "remote" remote.metadata.runtime_cfg = None - assert runtime._replica_init_group_key(remote) == "remote:1:0" - diffusion = _llm_replica(2, 0, "0") diffusion.metadata.stage_type = "diffusion" - assert runtime._replica_init_group_key(diffusion) == "inline:diffusion" + + # Same device everywhere; parallel mode still keys per replica (child SH/EX locks coordinate). + keys = runtime._init_group_keys([_llm_replica(0, 0, "0"), _llm_replica(0, 1, "0"), remote, diffusion]) + + assert keys == ["parallel:0:0", "parallel:0:1", "remote:1:0", "inline:diffusion"] + + +def test_unresolved_devices_collapse_serial_groups_and_warn(caplog): + runtime = _runtime(parallel_stage_init=False) + with caplog.at_level("WARNING"): + keys = runtime._init_group_keys([_llm_replica(0, 0, "0"), _llm_replica(1, 0, "GPU-uuid")]) + + assert keys == ["device-group:*", "device-group:*"] + assert "Stage-1 replica 0" in caplog.text + assert "GPU-uuid" in caplog.text + + +def _stage_plans(*devices: str) -> list[LogicalStageInitPlan]: + return [ + LogicalStageInitPlan(stage_idx=i, stage_id=i, replicas=[_llm_replica(i, 0, dev)]) + for i, dev in enumerate(devices) + ] + + +def test_overlapping_device_stages_initialize_sequentially(monkeypatch): + monkeypatch.delenv(current_omni_platform.device_control_env_var, raising=False) + runtime = _runtime(parallel_stage_init=False) + in_flight, peak = 0, 0 + guard = threading.Lock() + + def _initialize_replica(plan, _timeout): + nonlocal in_flight, peak + with guard: + in_flight += 1 + peak = max(peak, in_flight) + time.sleep(0.05) + with guard: + in_flight -= 1 + return object() + + monkeypatch.setattr(runtime, "_initialize_replica", _initialize_replica) + runtime._initialize_stage_replicas(_stage_plans("0,1", "0"), stage_init_timeout=5) + + assert peak == 1 + + +def test_disjoint_device_stages_initialize_in_parallel(monkeypatch): + monkeypatch.delenv(current_omni_platform.device_control_env_var, raising=False) + runtime = _runtime(parallel_stage_init=False) + barrier = threading.Barrier(2, timeout=2) # breaks unless both inits are in flight together + + def _initialize_replica(_plan, _timeout): + barrier.wait() + return object() + + monkeypatch.setattr(runtime, "_initialize_replica", _initialize_replica) + runtime._initialize_stage_replicas(_stage_plans("0", "1"), stage_init_timeout=5) # --------------------------------------------------------------------------- # diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 1e57036a0d1..0562d2b085a 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -1668,6 +1668,41 @@ def record_lock_holder_pid(fd: int, writable: bool) -> None: pass +def parse_physical_device_ids(devices: str | None) -> frozenset[int] | None: + """Parse ``"0,1"`` into ``{0, 1}``; ``None`` if missing, empty or non-integer (UUID/MIG).""" + tokens = [tok.strip() for tok in str(devices or "").split(",") if tok.strip()] + if not tokens or not all(tok.isdigit() for tok in tokens): + return None + return frozenset(int(tok) for tok in tokens) + + +def device_overlap_group_keys(device_sets: Sequence[frozenset[int] | None]) -> list[str]: + """Key each device set by the connected component of sets it shares a GPU with. + + Transitively overlapping sets get one key, the sorted union of the component + (``{0,1}`` and ``{0}`` -> ``device-group:0,1``); disjoint sets get distinct + keys. ``None`` (unresolved: may touch any GPU) overlaps everything, so one + ``None`` collapses all sets into a single group. + """ + resolved = [devices for devices in device_sets if devices is not None] + if len(resolved) != len(device_sets): + return ["device-group:*"] * len(device_sets) + + components: list[set[int]] = [] + for devices in resolved: + merged = set(devices) + disjoint = [] + for comp in components: + if comp & merged: + merged |= comp + else: + disjoint.append(comp) + components = [*disjoint, merged] + + key_of = {device: "device-group:" + ",".join(map(str, sorted(comp))) for comp in components for device in comp} + return [key_of[next(iter(devices))] for devices in resolved] + + def acquire_device_locks( stage_id: int, engine_args_dict: dict[str, Any], diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index 93e482b4255..5d7d07653af 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -52,11 +52,13 @@ build_llm_stage_output_processor, build_vllm_config, compute_replica_layout, + device_overlap_group_keys, extract_legacy_stage_metadata, get_stage_connector_spec, inject_kv_stage_info, inject_omni_kv_connector_config, load_omni_transfer_config_for_model, + parse_physical_device_ids, prepare_engine_environment, release_device_locks, stage_runtime_env, @@ -446,9 +448,10 @@ def _initialize_stage_replicas( ) -> dict[int, list[StagePoolClient | None]]: """Initialize all stage replicas. - Stages sharing the same GPU are initialized sequentially to avoid - memory profiling interference. Stages on different GPUs are - initialized in parallel. + Stages that share any physical GPU — including overlapping-but-unequal + sets such as ``{0,1}`` and ``{0}`` — initialize sequentially to avoid + memory profiling interference and same-process ``flock`` self-contention. + Stages on disjoint GPUs initialize in parallel. """ initialized_clients_by_stage: dict[int, list[StagePoolClient | None]] = { plan.stage_idx: [None] * len(plan.replicas) for plan in stage_plans @@ -464,10 +467,7 @@ def _initialize_stage_replicas( self._reject_unguardable_executors(stage_plans) self._run_stage_admission(stage_plans) - init_groups: dict[str, list[tuple[int, ReplicaInitPlan]]] = {} - for plan in stage_plans: - for replica in plan.replicas: - init_groups.setdefault(self._replica_init_group_key(replica), []).append((plan.stage_idx, replica)) + init_groups = self._build_init_groups(stage_plans) def _init_group(group: list[tuple[int, ReplicaInitPlan]]) -> None: """Initialize replicas in one scheduling group sequentially.""" @@ -589,19 +589,8 @@ def _resolve(replica: ReplicaInitPlan) -> list[int] | AdmissionExempt | None: # fails admission (fail-closed in check_admission). return ADMISSION_EXEMPT resolved = self._resolve_replica_physical_devices(replica.metadata.stage_id, replica.metadata.runtime_cfg) - if not resolved: - return None - ids: list[int] = [] - for tok in str(resolved).split(","): - tok = tok.strip() - if not tok: - continue - try: - ids.append(int(tok)) - except ValueError: - # Non-integer visibility (UUID / MIG) can't be reasoned about. - return None - return ids or None + parsed = parse_physical_device_ids(resolved) + return sorted(parsed) if parsed else None def _visible_ordinal(physical_id: int) -> int: """Translate a physical device id to this process's visible ordinal. @@ -656,47 +645,58 @@ def _total_memory(device_id: int) -> int: device_total_memory=_total_memory, ) - def _replica_init_group_key(self, replica: ReplicaInitPlan) -> str: - """Return the scheduling group used during replica initialization. - - Replicas sharing a group initialize sequentially; different groups - initialize in parallel threads. Local LLM replicas are keyed by their - **resolved canonical physical device set** so stages on different - physical GPUs land in different groups (parallel) while stages sharing a - device stay in one group (serialized, then further guarded by the - per-device ``LOCK_EX`` file lock). Keying on the raw ``runtime.devices`` - config value instead would fail to parallelize logically-distinct stages - that map to different physical GPUs. + def _build_init_groups( + self, + stage_plans: Sequence[LogicalStageInitPlan], + ) -> dict[str, list[tuple[int, ReplicaInitPlan]]]: + items = [(plan.stage_idx, replica) for plan in stage_plans for replica in plan.replicas] + groups: dict[str, list[tuple[int, ReplicaInitPlan]]] = {} + for key, item in zip(self._init_group_keys([replica for _, replica in items]), items, strict=True): + groups.setdefault(key, []).append(item) + return groups + + def _init_group_keys(self, replicas: Sequence[ReplicaInitPlan]) -> list[str]: + """One init-group key per replica: same key -> sequential, distinct -> parallel. + + Diffusion, remote and parallel-stage-init replicas keep fixed per-replica + keys. Serial local LLM replicas are grouped by connected components of + their resolved physical-device overlap, so ``{0,1}`` and ``{0}`` share a + group and never re-``flock`` a device this process already holds. """ + keys = [self._init_group_key_override(replica) for replica in replicas] + pending = [i for i, key in enumerate(keys) if key is None] + device_sets: list[frozenset[int] | None] = [] + for i in pending: + replica = replicas[i] + physical = self._resolve_replica_physical_devices(replica.metadata.stage_id, replica.metadata.runtime_cfg) + parsed = parse_physical_device_ids(physical) + if parsed is None: + logger.warning( + "[stage_init] Stage-%s replica %s physical devices %r are not integer GPU ids; " + "all serial LLM replicas will share one init group", + replica.metadata.stage_id, + replica.replica_id, + physical, + ) + device_sets.append(parsed) + for i, key in zip(pending, device_overlap_group_keys(device_sets), strict=True): + keys[i] = key + return cast(list[str], keys) + + def _init_group_key_override(self, replica: ReplicaInitPlan) -> str | None: + """Key that bypasses device-overlap grouping, or ``None`` to use it.""" if replica.launch_mode == "local" and replica.metadata.stage_type == "diffusion": # Local diffusion process spawning must stay on the orchestrator # thread. Keep all local diffusion replicas in one sequential group. return "inline:diffusion" if replica.launch_mode == "remote": return f"remote:{replica.metadata.stage_id}:{replica.replica_id}" - - physical_devices = self._resolve_replica_physical_devices( - replica.metadata.stage_id, - replica.metadata.runtime_cfg, - ) if self._parallel_stage_init: # Same-device concurrency is coordinated by the engine-core SH/EX # device locks + pre-launch admission, so give every replica its own # group to let them all initialize in parallel. return f"parallel:{replica.metadata.stage_id}:{replica.replica_id}" - - runtime_cfg = replica.metadata.runtime_cfg or {} - raw_devices = ( - runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None) - ) - if str(raw_devices) != str(physical_devices): - logger.debug( - "[stage_init] Stage-%s init-group key: raw devices=%s -> resolved physical=%s", - replica.metadata.stage_id, - raw_devices, - physical_devices, - ) - return f"device:{physical_devices}" + return None def _initialize_replica( self, From 2f1845fa4e74d78609a2c301d3b843628cbf6fb8 Mon Sep 17 00:00:00 2001 From: kunkun <72174834+kunkunblueberry@users.noreply.github.com> Date: Wed, 9 Sep 2026 20:35:16 +0800 Subject: [PATCH 32/93] fix: reserve Qwen3-Omni NVFP4 backend fix (#7200) Signed-off-by: kunkunblueberry <1833921874@qq.com> --- .../test_qwen3_omni_modelopt_nvfp4_w4a4.py | 7 +++- tests/engine/test_stage_engine_args.py | 40 +++++++++++++++++++ vllm_omni/config/omni_config.py | 2 +- vllm_omni/engine/stage_init_utils.py | 18 ++++++--- 4 files changed, 59 insertions(+), 8 deletions(-) diff --git a/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py b/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py index 9502ed5e631..6537b1cac40 100644 --- a/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py +++ b/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py @@ -65,7 +65,12 @@ def _stage_config(): # gains enough VRAM headroom for a no-eager run, add a companion test. return modify_stage_config( _CI_DEPLOY, - updates={"stages": {0: {"enforce_eager": True}, 1: {"enforce_eager": True}}}, + updates={ + "stages": { + 0: {"enforce_eager": True, "moe_backend": "auto"}, + 1: {"enforce_eager": True}, + } + }, ) diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index e2196dfd1ca..250aff88464 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/tests/engine/test_stage_engine_args.py @@ -11,6 +11,7 @@ import pytest from pydantic.fields import FieldInfo +from transformers import Qwen3OmniMoeConfig from vllm.config import CacheConfig as VllmCacheConfig from vllm.config import CompilationConfig as VllmCompilationConfig from vllm.config import LoadConfig as VllmLoadConfig @@ -19,6 +20,7 @@ from vllm.config import SchedulerConfig as VllmSchedulerConfig from vllm.engine.arg_utils import EngineArgs +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config from vllm_omni.config.omni_config import ( _LLM_STAGE_ENGINE_FIELDS, OmniStageCacheConfig, @@ -155,6 +157,44 @@ def test_qwen3_omni_preserves_explicit_moe_backend(): assert engine_args["moe_backend"] == "flashinfer" +def test_qwen3_omni_preserves_explicit_auto_moe_backend(): + engine_args = { + "model_arch": "Qwen3OmniMoeForConditionalGeneration", + "moe_backend": "auto", + } + + stage_init_utils._maybe_set_qwen3_omni_moe_backend(engine_args) + + assert engine_args["moe_backend"] == "auto" + + +def test_qwen3_omni_nvfp4_ci_config_preserves_auto_moe_backend(tmp_path): + """The NVFP4 CI deploy opts stage 0 into vLLM's auto backend selection.""" + hf_config = Qwen3OmniMoeConfig(enable_audio_output=True) + pipeline = resolve_pipeline_config("qwen3_omni_moe", hf_config) + assert pipeline is not None + nvfp4_deploy = modify_stage_config( + get_deploy_config_path("ci/qwen3_omni_moe.yaml"), + updates={"stages": {0: {"moe_backend": "auto"}}}, + ) + deploy = load_deploy_config(nvfp4_deploy) + legacy_stages = merge_pipeline_deploy(pipeline, copy.deepcopy(deploy)) + omni_config = VllmOmniConfig.from_pipeline_config( + pipeline, + user_deploy_config=copy.deepcopy(deploy), + cli_overrides={"model": str(tmp_path)}, + ) + + for stage_id, expected_backend in ((0, "auto"), (1, "triton")): + legacy_stage = next(stage for stage in legacy_stages if stage.stage_id == stage_id) + typed_stage = omni_config.stage_by_id(stage_id) + for engine_args in ( + build_engine_args_dict(legacy_stage.to_omegaconf(), str(tmp_path)), + build_engine_args_dict_from_omni_stage_config(typed_stage, str(tmp_path)), + ): + assert engine_args["moe_backend"] == expected_backend + + def _engine_arg_inputs(tmp_path: Path) -> tuple[PipelineConfig, DeployConfig, str]: parent_model = tmp_path / "parent-model" stage_model = tmp_path / "stage-model" diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index aaa9c35122d..81eb76f0b41 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -400,7 +400,7 @@ def _get_deploy_config( @config -class OmniStageModelConfig: +class OmniStageModelConfig(_TrackExplicitConfigFields): """Per-stage model behavior and resolved model-engine inputs.""" model: str | None = None diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 0562d2b085a..765549bcb37 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -764,12 +764,15 @@ def prepare_engine_environment() -> None: pass -def _maybe_set_qwen3_omni_moe_backend(engine_args_dict: dict[str, Any]) -> None: +def _maybe_set_qwen3_omni_moe_backend( + engine_args_dict: dict[str, Any], + *, + moe_backend_is_explicit: bool | None = None, +) -> None: """Choose the stable MoE backend when Qwen3-Omni has no explicit choice.""" - if ( - engine_args_dict.get("model_arch") == "Qwen3OmniMoeForConditionalGeneration" - and engine_args_dict.get("moe_backend", "auto") == "auto" - ): + if moe_backend_is_explicit is None: + moe_backend_is_explicit = "moe_backend" in engine_args_dict + if engine_args_dict.get("model_arch") == "Qwen3OmniMoeForConditionalGeneration" and not moe_backend_is_explicit: engine_args_dict["moe_backend"] = "triton" logger.info("[stage_init] Set moe_backend=triton for Qwen3-Omni stage") @@ -1121,6 +1124,10 @@ def _project_omni_stage_engine_args( # The legacy builder always emits this key, including for pipelines such # as Audex that intentionally defer architecture discovery to HF config. engine_args["model_arch"] = copy.deepcopy(stage_config.model_config.model_arch) + _maybe_set_qwen3_omni_moe_backend( + engine_args, + moe_backend_is_explicit="moe_backend" in getattr(stage_config.model_config, "_omni_explicit_fields", ()), + ) topology = stage_config.stage_pipeline_config topology_engine_args = { @@ -1283,7 +1290,6 @@ def _finalize_engine_args_dict( engine_args_dict["has_sampling_extra_args"] = has_sampling_extra_args engine_args_dict["sampling_extra_args_keys"] = sampling_extra_args_keys - # Select the typed backend option during stage argument finalization. _maybe_set_qwen3_omni_moe_backend(engine_args_dict) return engine_args_dict From 0e602509aa9460cedca2de45094e6ea1d2be47ba Mon Sep 17 00:00:00 2001 From: Shaun Walsh <153730091+Shaun-Walsh@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:38:42 +0100 Subject: [PATCH 33/93] [BugFix] Add field validators for /v1/audio/generate request (#4741) Signed-off-by: Shaun Walsh Co-authored-by: Claude Opus 4.6 Co-authored-by: Nick Cao --- .../test_invalid_audio_diffusion.py | 33 +++++++++---------- .../entrypoints/openai/protocol/audio.py | 13 +++++++- 2 files changed, 27 insertions(+), 19 deletions(-) diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_diffusion.py b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_diffusion.py index f14dcf8c11b..720d8615172 100644 --- a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_diffusion.py +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_diffusion.py @@ -13,9 +13,6 @@ pytestmark = [pytest.mark.slow, pytest.mark.diffusion] -# https://github.com/vllm-project/vllm-omni/issues/3649 — invalid inputs still HTTP 200 / timeout / missing validation. -_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649") - _PARAMS = [ pytest.param( OmniServerParams( @@ -36,10 +33,8 @@ @pytest.mark.parametrize( "overrides, err_message", [ - pytest.param({"input": ""}, ("cannot be empty", "input"), id="input_empty", marks=_SKIP_ISSUE_3649), - pytest.param( - {"input": " "}, ("cannot be empty", "input"), id="input_whitespace_only", marks=_SKIP_ISSUE_3649 - ), + pytest.param({"input": ""}, ("cannot be empty", "input"), id="input_empty"), + pytest.param({"input": " "}, ("cannot be empty", "input"), id="input_whitespace_only"), pytest.param({"input": 123}, ("input", "a valid string"), id="input_wrong_type"), pytest.param( {"response_format": "mpeg"}, @@ -54,10 +49,8 @@ ("sse", "stream_format", "value_error", "not a supported stream_format"), id="stream_format_sse_blocked", ), - pytest.param( - {"audio_length": -1.0}, ("audio_end_in_s", "higher than", "audio_start_in_s"), id="audio_length_negative" - ), - pytest.param({"audio_length": 0}, "audio_length", id="audio_length_zero", marks=_SKIP_ISSUE_3649), + pytest.param({"audio_length": -1.0}, ("audio_length", "greater_than", "0"), id="audio_length_negative"), + pytest.param({"audio_length": 0}, ("audio_length", "greater_than", "0"), id="audio_length_zero"), # ``OmniOpenAIServingAudioGenerate`` only forwards ``audio_start`` when ``audio_length`` is set. pytest.param( {"audio_start": -0.5, "audio_length": 5.0}, @@ -67,26 +60,30 @@ pytest.param( {"negative_prompt": ["noise"]}, ("negative_prompt", "a valid string"), id="negative_prompt_wrong_type" ), - pytest.param({"guidance_scale": -1.0}, "guidance_scale", id="guidance_scale_negative", marks=_SKIP_ISSUE_3649), - pytest.param({"guidance_scale": 0}, "guidance_scale", id="guidance_scale_zero", marks=_SKIP_ISSUE_3649), + pytest.param( + {"guidance_scale": -1.0}, + ("guidance_scale", "greater_than_equal", "0"), + id="guidance_scale_negative", + ), pytest.param( {"num_inference_steps": 0}, - ("num_inference_steps", "greater than or equal to 1"), + ("num_inference_steps", "greater_than_equal", "1"), id="num_inference_steps_zero", ), pytest.param( {"num_inference_steps": -1}, - ("num_inference_steps", "greater than or equal to 1"), + ("num_inference_steps", "greater_than_equal", "1"), id="num_inference_steps_negative", ), pytest.param( {"num_inference_steps": 6000}, - "num_inference_steps", + ("num_inference_steps", "less_than_equal", "1000"), id="num_inference_steps_above_max", - marks=_SKIP_ISSUE_3649, ), pytest.param( - {"guidance_scale": 1001.0}, "guidance_scale", id="guidance_scale_above_max", marks=_SKIP_ISSUE_3649 + {"guidance_scale": 1001.0}, + ("guidance_scale", "less_than_equal", "1000"), + id="guidance_scale_above_max", ), pytest.param( {"audio_length": 86401.0}, ("Requested audio length", "exceeds maximum"), id="audio_length_above_max" diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index 2446114df17..459f9d4ae39 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -353,6 +353,14 @@ class OpenAICreateAudioGenerateRequest(BaseModel): input: str = Field( description="Text prompt describing the audio to generate", ) + + @field_validator("input") + @classmethod + def validate_input(cls, v): + if not v or not v.strip(): + raise ValueError("input cannot be empty") + return v + model: str | None = None response_format: Literal["wav", "pcm", "flac", "mp3", "opus"] = DEFAULT_AUDIO_FORMAT speed: float | None = Field( @@ -363,6 +371,7 @@ class OpenAICreateAudioGenerateRequest(BaseModel): stream_format: Literal["sse", "audio"] | None = "audio" audio_length: float | None = Field( default=None, + gt=0, description="Audio length in seconds", ) audio_start: float | None = Field( @@ -375,12 +384,14 @@ class OpenAICreateAudioGenerateRequest(BaseModel): ) guidance_scale: float | None = Field( default=None, + ge=0, + le=1000, description="Guidance scale for diffusion models", ) num_inference_steps: int | None = Field( default=None, ge=1, - le=_INT64_MAX, + le=1000, description="Number of inference steps", ) seed: int | None = Field( From c94cf4ddde3cddf28fb02fbe27029824ae5ad6c8 Mon Sep 17 00:00:00 2001 From: andyluo7 <43718156+andyluo7@users.noreply.github.com> Date: Wed, 9 Sep 2026 06:47:15 -0700 Subject: [PATCH 34/93] [CI][ROCm] Match CUDA/NPU L2/L3 label routing (#6966) Signed-off-by: andyluo7 Co-authored-by: Hongsheng Liu --- .buildkite/amd/scripts/bootstrap-amd-omni.sh | 166 +++++++++---------- .buildkite/amd/scripts/select_test_suites.py | 93 +++++++++++ docs/contributing/ci/ci_settings.md | 4 +- tests/buildkite/test_amd_bootstrap.py | 156 +++++++++++++++++ tests/buildkite/test_amd_suite_selection.py | 61 +++++++ 5 files changed, 391 insertions(+), 89 deletions(-) create mode 100644 .buildkite/amd/scripts/select_test_suites.py create mode 100644 tests/buildkite/test_amd_bootstrap.py create mode 100644 tests/buildkite/test_amd_suite_selection.py diff --git a/.buildkite/amd/scripts/bootstrap-amd-omni.sh b/.buildkite/amd/scripts/bootstrap-amd-omni.sh index 4eb7f6eb47a..9eff6b6c7a4 100755 --- a/.buildkite/amd/scripts/bootstrap-amd-omni.sh +++ b/.buildkite/amd/scripts/bootstrap-amd-omni.sh @@ -14,6 +14,8 @@ set -euo pipefail # shellcheck disable=SC1091 source .buildkite/common/scripts/resolve_skip_ci.sh +AMD_SUITE_SELECTOR=".buildkite/amd/scripts/select_test_suites.py" + if [[ -z "${RUN_ALL:-}" ]]; then RUN_ALL=0 fi @@ -30,31 +32,20 @@ if [[ -z "${AMD_MIRROR_HW:-}" ]]; then AMD_MIRROR_HW="amdproduction" fi -fail_fast() { - DISABLE_LABEL="ci-no-fail-fast" - # If BUILDKITE_PULL_REQUEST != "false", then we check the PR labels using curl and jq - if [ "$BUILDKITE_PULL_REQUEST" != "false" ]; then - PR_LABELS=$(curl -s "https://api.github.com/repos/vllm-project/vllm-omni/pulls/$BUILDKITE_PULL_REQUEST" | jq -r '.labels[].name') - if [[ $PR_LABELS == *"$DISABLE_LABEL"* ]]; then - echo false - else - echo true - fi - else - echo false # not a PR or BUILDKITE_PULL_REQUEST not set +PR_LABELS="" +# Explicit debug suites do not depend on PR labels. +if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" && -z "${DEBUG_TEST_YAML:-}" ]]; then + if ! PR_LABELS=$(curl -fsSL \ + "https://api.github.com/repos/vllm-project/vllm-omni/pulls/$BUILDKITE_PULL_REQUEST" \ + | jq -r '.labels[].name'); then + echo "ERROR: Could not read PR labels; refusing to select an incorrect AMD test tier." >&2 + exit 1 fi -} +fi -check_run_all_label() { - RUN_ALL_LABEL="ready-run-all-tests" - # If BUILDKITE_PULL_REQUEST != "false", then we check the PR labels using curl and jq - if [ "$BUILDKITE_PULL_REQUEST" != "false" ]; then - PR_LABELS=$(curl -s "https://api.github.com/repos/vllm-project/vllm-omni/pulls/$BUILDKITE_PULL_REQUEST" | jq -r '.labels[].name') - if [[ $PR_LABELS == *"$RUN_ALL_LABEL"* ]]; then - echo true - else - echo false - fi +fail_fast() { + if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then + echo true else echo false # not a PR or BUILDKITE_PULL_REQUEST not set fi @@ -64,6 +55,57 @@ if [[ -z "${COV_ENABLED:-}" ]]; then COV_ENABLED=0 fi +resolve_test_specs() { + local selected_specs + if ! selected_specs=$(python3 "$AMD_SUITE_SELECTOR" \ + --branch "$BUILDKITE_BRANCH" \ + --labels "$PR_LABELS" \ + --debug-test-yaml "${DEBUG_TEST_YAML:-}"); then + exit 1 + fi + + TEST_SPECS=() + while IFS= read -r suite_spec; do + [[ -z "$suite_spec" ]] && continue + TEST_SPECS+=("$suite_spec") + done <<< "$selected_specs" +} + +filter_test_specs_by_skip_ci() { + local suite_spec level decision + local skip_all=0 + local -a runnable_specs=() + + for suite_spec in "${TEST_SPECS[@]}"; do + case "$suite_spec" in + READY_TESTS:*) level="l2" ;; + MERGE_TESTS:*) level="l3" ;; + *) + echo "ERROR: unknown AMD test suite spec '$suite_spec'" >&2 + exit 1 + ;; + esac + + if decision=$(python3 "$SKIP_CI_PY" gate amd "$level"); then + echo "Skipping AMD $level suite due to $decision decision." + [[ "$decision" == "skip-all" ]] && skip_all=1 + else + runnable_specs+=("$suite_spec") + fi + done + + TEST_SPECS=("${runnable_specs[@]}") + if [[ ${#TEST_SPECS[@]} -eq 0 ]]; then + if [[ $skip_all -eq 1 ]]; then + buildkite-agent annotate \ + ":memo: CI skipped — docs or pytest skip-mark changes only" \ + --style "info" 2>/dev/null || true + fi + echo "No AMD L2/L3 suites remain after skip-ci filtering." + exit 0 + fi +} + upload_pipeline() { echo "Uploading pipeline..." # Install minijinja @@ -84,7 +126,7 @@ upload_pipeline() { # (WIP) Use pipeline generator instead of jinja template if [ -e ".buildkite/amd/pipeline_generator/pipeline_generator.py" ]; then python -m pip install click pydantic - python .buildkite/amd/pipeline_generator/pipeline_generator.py --run_all=$RUN_ALL --list_file_diff="$LIST_FILE_DIFF" --nightly="$NIGHTLY" --mirror_hw="$AMD_MIRROR_HW" + python .buildkite/amd/pipeline_generator/pipeline_generator.py --run_all="$RUN_ALL" --list_file_diff="$LIST_FILE_DIFF" --nightly="$NIGHTLY" --mirror_hw="$AMD_MIRROR_HW" buildkite-agent pipeline upload .buildkite/amd/pipeline.yaml exit 0 fi @@ -97,49 +139,12 @@ upload_pipeline() { cd .buildkite/amd - # Select test definition file: merge suite for main, ready suite for PRs. - # For debugging, DEBUG_TEST_YAML accepts a comma-separated list containing - # "merge" and/or "ready" (case-insensitive). Multiple suites are combined - # before rendering so they share one amd-build step. - if [[ -n "${DEBUG_TEST_YAML:-}" ]]; then - declare -a DEBUG_TEST_SPECS=() - declare -A SEEN_DEBUG_TESTS=() - IFS=',' read -ra REQUESTED_DEBUG_TESTS <<< "${DEBUG_TEST_YAML,,}" - - for requested_test in "${REQUESTED_DEBUG_TESTS[@]}"; do - # Permit readable values such as "ready, merge," and ignore the - # empty item produced by a trailing comma. - requested_test="${requested_test//[[:space:]]/}" - [[ -z "$requested_test" ]] && continue - - if [[ -n "${SEEN_DEBUG_TESTS[$requested_test]:-}" ]]; then - echo "ERROR: duplicate DEBUG_TEST_YAML suite '$requested_test'" >&2 - exit 1 - fi - SEEN_DEBUG_TESTS[$requested_test]=1 - - case "$requested_test" in - ready) - DEBUG_TEST_SPECS+=("READY_TESTS:test-amd-ready.yml") - ;; - merge) - DEBUG_TEST_SPECS+=("MERGE_TESTS:test-amd-merge.yml") - ;; - *) - echo "ERROR: DEBUG_TEST_YAML entries must be 'merge' or 'ready', got '$requested_test'" >&2 - exit 1 - ;; - esac - done - - if [[ ${#DEBUG_TEST_SPECS[@]} -eq 0 ]]; then - echo "ERROR: DEBUG_TEST_YAML did not contain a test suite" >&2 - exit 1 - elif [[ ${#DEBUG_TEST_SPECS[@]} -eq 1 ]]; then - TEST_YAML="${DEBUG_TEST_SPECS[0]#*:}" - else - TEST_YAML=$(mktemp "${TMPDIR:-/tmp}/amd-debug-tests.XXXXXX.yml") - python - "$TEST_YAML" "${DEBUG_TEST_SPECS[@]}" <<'PY' + # Multiple label-selected or debug-selected suites share one image build. + if [[ ${#TEST_SPECS[@]} -eq 1 ]]; then + TEST_YAML="${TEST_SPECS[0]#*:}" + else + TEST_YAML=$(mktemp "${TMPDIR:-/tmp}/amd-selected-tests.XXXXXX.yml") + python - "$TEST_YAML" "${TEST_SPECS[@]}" <<'PY' import sys import yaml @@ -172,13 +177,8 @@ for suite_spec in suite_specs: with open(output_path, "w", encoding="utf-8") as output_file: yaml.safe_dump(combined, output_file, sort_keys=False) PY - fi - echo "DEBUG_TEST_YAML override: using ${DEBUG_TEST_SPECS[*]}" - elif [[ $BUILDKITE_BRANCH == "main" ]]; then - TEST_YAML="test-amd-merge.yml" - else - TEST_YAML="test-amd-ready.yml" fi + echo "AMD test suites: ${TEST_SPECS[*]}" ( set -x @@ -198,7 +198,7 @@ PY > pipeline.yaml ) cat pipeline.yaml - if [[ "$TEST_YAML" == "${TMPDIR:-/tmp}/amd-debug-tests."*.yml ]]; then + if [[ "$TEST_YAML" == "${TMPDIR:-/tmp}/amd-selected-tests."*.yml ]]; then rm -f -- "$TEST_YAML" fi buildkite-agent artifact upload pipeline.yaml @@ -219,12 +219,10 @@ if [[ $BUILDKITE_BRANCH == "main" ]]; then file_diff=$(get_diff_main) fi -# Early exit: unified skip-ci (docs / skip marks) and CI-yaml-only level targeting. -if [[ $BUILDKITE_BRANCH == "main" ]]; then - gate_bootstrap_ci amd l3 -else - gate_bootstrap_ci amd l2 -fi +# Resolve PR tier labels before skip-ci so ready and merge suites can be +# filtered independently for CI-YAML-only changes. +resolve_test_specs +filter_test_specs_by_skip_ci patterns=( "docker/Dockerfile" @@ -274,14 +272,6 @@ for file in $file_diff; do fi done -# Check for ready-run-all-tests label -LABEL_RUN_ALL=$(check_run_all_label) -if [[ $LABEL_RUN_ALL == true ]]; then - RUN_ALL=1 - NIGHTLY=1 - echo "Found 'ready-run-all-tests' label. Running all tests including optional tests." -fi - # Decide whether to use precompiled wheels # Relies on existing patterns array as a basis. if [[ -n "${VLLM_USE_PRECOMPILED:-}" ]]; then diff --git a/.buildkite/amd/scripts/select_test_suites.py b/.buildkite/amd/scripts/select_test_suites.py new file mode 100644 index 00000000000..b7c72f6e472 --- /dev/null +++ b/.buildkite/amd/scripts/select_test_suites.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Select AMD L2/L3 suites after Buildkite has created an AMD build. + +This repository-side selector does not decide which GitHub label events start +the external ``vllm-omni-amd-ci`` pipeline. Its Buildkite PR build condition +must admit the relevant label event before this code can inspect the labels. +""" + +from __future__ import annotations + +import argparse +import sys +from collections.abc import Iterable + +SUITE_SPECS = { + "ready": "READY_TESTS:test-amd-ready.yml", + "merge": "MERGE_TESTS:test-amd-merge.yml", +} + + +def _parse_debug_suites(value: str) -> tuple[str, ...]: + suites = tuple(item.strip().lower() for item in value.split(",") if item.strip()) + if not suites: + raise ValueError("DEBUG_TEST_YAML did not contain a test suite") + + seen: set[str] = set() + for suite in suites: + if suite not in SUITE_SPECS: + raise ValueError(f"DEBUG_TEST_YAML entries must be 'merge' or 'ready', got {suite!r}") + if suite in seen: + raise ValueError(f"duplicate DEBUG_TEST_YAML suite {suite!r}") + seen.add(suite) + return suites + + +def select_amd_test_suites( + *, + branch: str, + labels: Iterable[str], + debug_test_yaml: str = "", +) -> tuple[str, ...]: + """Return ordered suite names while preserving legacy unlabeled PRs.""" + + if debug_test_yaml: + return _parse_debug_suites(debug_test_yaml) + if branch == "main": + return ("merge",) + + label_set = {label.strip() for label in labels if label.strip()} + selected = tuple(suite for label, suite in (("ready", "ready"), ("merge-test", "merge")) if label in label_set) + # AMD historically uploaded L2 for every PR build. Keep that behavior for + # triggers such as amd-test while adding tier selection for ready/merge-test. + return selected or ("ready",) + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--branch", required=True) + parser.add_argument("--labels", default="") + parser.add_argument("--debug-test-yaml", default="") + args = parser.parse_args() + + labels = tuple(args.labels.splitlines()) + try: + suites = select_amd_test_suites( + branch=args.branch, + labels=labels, + debug_test_yaml=args.debug_test_yaml, + ) + except ValueError as exc: + parser.error(str(exc)) + + label_set = {label.strip() for label in labels if label.strip()} + if not args.debug_test_yaml and args.branch != "main" and not ({"ready", "merge-test"} & label_set): + print( + "No AMD L2/L3 tier label found; preserving the legacy ready-suite fallback.", + file=sys.stderr, + ) + if "nightly-test" in label_set: + print( + "AMD nightly-test has no suite yet and does not affect L2/L3 selection.", + file=sys.stderr, + ) + + for suite in suites: + print(SUITE_SPECS[suite]) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/docs/contributing/ci/ci_settings.md b/docs/contributing/ci/ci_settings.md index 7348ad1832a..d578b4b02df 100644 --- a/docs/contributing/ci/ci_settings.md +++ b/docs/contributing/ci/ci_settings.md @@ -161,7 +161,9 @@ Canonical layout (prefer these paths for new changes): **Bootstrap:** [`amd/scripts/bootstrap-amd-omni.sh`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/amd/scripts/bootstrap-amd-omni.sh)—skip-ci, diff filtering, Jinja render, then `buildkite-agent pipeline upload`. - **Test YAML (data):** `amd/test-amd-ready.yml` (L2 / PR), `test-amd-merge.yml` (L3 / main). + **Test YAML (data):** `amd/test-amd-ready.yml` (L2 / `ready`), `test-amd-merge.yml` (L3 / `merge-test` and `main`). On a PR carrying both labels, AMD combines both suites behind one image build. PR builds without either tier label retain the legacy L2 fallback for compatibility with the `amd-test` trigger. `DEBUG_TEST_YAML` remains an explicit override. AMD has no L4 file yet, so `nightly-test` does not select an additional suite. + + **Trigger boundary:** AMD label handling has two layers. First, the external `vllm-omni-amd-ci` Buildkite pipeline condition decides whether a GitHub label event creates a build. Only after that build starts does this repository's bootstrap inspect all current PR labels and select L2, L3, or both. Repository-side selection therefore cannot make a `merge-test` event start AMD CI by itself. The external condition must admit both `ready` and `merge-test`, while preserving the legacy `amd-test` trigger and non-PR `main` / scheduled builds. For PR builds without a `DEBUG_TEST_YAML` override, an unreadable label set fails the bootstrap instead of silently selecting the wrong tier. Debug overrides skip label lookup and are validated directly; invalid values, including whitespace-only values, fail the bootstrap. Keep `nightly-test` out of the AMD condition until an AMD L4 suite exists; otherwise it would start a build without providing nightly coverage. **Rendering:** [`test-template-amd-omni.j2`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/amd/test-template-amd-omni.j2) wraps data steps with `amd-build` image build and `amd_` queues. Do **not** hand-edit generated `pipeline.yaml`. diff --git a/tests/buildkite/test_amd_bootstrap.py b/tests/buildkite/test_amd_bootstrap.py new file mode 100644 index 00000000000..593ba8b2602 --- /dev/null +++ b/tests/buildkite/test_amd_bootstrap.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +REPO_ROOT = Path(__file__).resolve().parents[2] +BOOTSTRAP_PATH = REPO_ROOT / ".buildkite/amd/scripts/bootstrap-amd-omni.sh" + + +def _run_bootstrap( + tmp_path: Path, + *, + debug_test_yaml: str | None = None, + curl_status: int = 0, + jq_status: int = 0, + labels: tuple[str, ...] = (), + pull_request: bool = True, +) -> subprocess.CompletedProcess[str]: + for relative_path in ( + ".buildkite/common/scripts/resolve_skip_ci.sh", + ".buildkite/amd/scripts/select_test_suites.py", + ): + destination = tmp_path / relative_path + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(REPO_ROOT / relative_path, destination) + + # Stop after real suite selection, before rendering or uploading a pipeline. + (tmp_path / ".buildkite/common/scripts/skip_ci.py").write_text("print('test-skip')\n", encoding="utf-8") + env = { + "PATH": os.environ["PATH"], + "BUILDKITE_PULL_REQUEST": "123" if pull_request else "false", + "BUILDKITE_BRANCH": "feature" if pull_request else "main", + "TEST_PYTHON": sys.executable, + "TEST_CURL_STATUS": str(curl_status), + "TEST_JQ_STATUS": str(jq_status), + "TEST_PR_LABELS": "\n".join(labels), + } + if debug_test_yaml is not None: + env["DEBUG_TEST_YAML"] = debug_test_yaml + + # Source the unmodified bootstrap with subprocess-local command doubles. + script = r""" +curl() { + echo 'curl called' >&2 + return "$TEST_CURL_STATUS" +} +jq() { + echo 'jq called' >&2 + printf '%s\n' "$TEST_PR_LABELS" + return "$TEST_JQ_STATUS" +} +git() { :; } +python3() { "$TEST_PYTHON" "$@"; } +source "$1" +""" + return subprocess.run( + ["bash", "-c", script, "test-amd-bootstrap", str(BOOTSTRAP_PATH)], + cwd=tmp_path, + env=env, + capture_output=True, + text=True, + timeout=10, + check=False, + ) + + +@pytest.mark.parametrize(("curl_status", "jq_status"), [(22, 0), (0, 4)], ids=["request-failure", "parse-failure"]) +@pytest.mark.parametrize( + ("debug_test_yaml", "expected_levels"), + [ + ("ready", ["l2"]), + ("merge", ["l3"]), + ("ready,merge", ["l2", "l3"]), + (" MERGE, ready,", ["l3", "l2"]), + ], +) +def test_debug_override_skips_label_lookup(tmp_path, debug_test_yaml, expected_levels, curl_status, jq_status): + result = _run_bootstrap( + tmp_path, + debug_test_yaml=debug_test_yaml, + curl_status=curl_status, + jq_status=jq_status, + ) + + assert result.returncode == 0, result.stderr + assert "curl called" not in result.stderr + assert "jq called" not in result.stderr + assert "legacy ready-suite fallback" not in result.stderr + assert [ + line.split()[2] for line in result.stdout.splitlines() if line.startswith("Skipping AMD ") + ] == expected_levels + + +@pytest.mark.parametrize(("curl_status", "jq_status"), [(22, 0), (0, 4)], ids=["request-failure", "parse-failure"]) +@pytest.mark.parametrize("debug_test_yaml", [None, ""], ids=["unset", "empty"]) +def test_label_lookup_failure_without_override(tmp_path, debug_test_yaml, curl_status, jq_status): + result = _run_bootstrap( + tmp_path, + debug_test_yaml=debug_test_yaml, + curl_status=curl_status, + jq_status=jq_status, + ) + + assert result.returncode == 1 + assert "Could not read PR labels" in result.stderr + assert "Skipping AMD " not in result.stdout + + +@pytest.mark.parametrize( + ("labels", "expected_levels"), + [ + (("ready",), ["l2"]), + (("merge-test",), ["l3"]), + (("ready", "merge-test"), ["l2", "l3"]), + ((), ["l2"]), + ], +) +def test_label_lookup_selects_tiers(tmp_path, labels, expected_levels): + result = _run_bootstrap(tmp_path, labels=labels) + + assert result.returncode == 0, result.stderr + assert "curl called" in result.stderr + assert "jq called" in result.stderr + assert [ + line.split()[2] for line in result.stdout.splitlines() if line.startswith("Skipping AMD ") + ] == expected_levels + + +@pytest.mark.parametrize("debug_test_yaml", ["ready,ready", "nightly", ", ,", " \t"]) +def test_invalid_debug_override_fails_without_label_lookup(tmp_path, debug_test_yaml): + result = _run_bootstrap(tmp_path, debug_test_yaml=debug_test_yaml, curl_status=22) + + assert result.returncode == 1 + assert "DEBUG_TEST_YAML" in result.stderr + assert "curl called" not in result.stderr + assert "jq called" not in result.stderr + assert "Skipping AMD " not in result.stdout + + +def test_main_does_not_lookup_labels(tmp_path): + result = _run_bootstrap(tmp_path, pull_request=False, curl_status=22, jq_status=4) + + assert result.returncode == 0, result.stderr + assert "curl called" not in result.stderr + assert "jq called" not in result.stderr + assert "Skipping AMD l3 suite" in result.stdout diff --git a/tests/buildkite/test_amd_suite_selection.py b/tests/buildkite/test_amd_suite_selection.py new file mode 100644 index 00000000000..b1b25aa2651 --- /dev/null +++ b/tests/buildkite/test_amd_suite_selection.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +SELECTOR_PATH = Path(".buildkite/amd/scripts/select_test_suites.py") +SPEC = importlib.util.spec_from_file_location("select_test_suites", SELECTOR_PATH) +assert SPEC is not None and SPEC.loader is not None +SELECTOR = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(SELECTOR) + + +@pytest.mark.parametrize( + ("branch", "labels", "expected"), + [ + ("main", (), ("merge",)), + ("feature", ("ready",), ("ready",)), + ("feature", ("merge-test",), ("merge",)), + ("feature", ("merge-test", "ready"), ("ready", "merge")), + ("feature", ("amd-test",), ("ready",)), + ("feature", ("amd-test", "merge-test"), ("merge",)), + ("feature", ("amd-test", "ready", "merge-test"), ("ready", "merge")), + ("feature", ("nightly-test",), ("ready",)), + ("feature", ("not-ready", "merge-test-extra"), ("ready",)), + ], +) +def test_label_suite_selection(branch, labels, expected): + assert SELECTOR.select_amd_test_suites(branch=branch, labels=labels) == expected + + +def test_debug_override_takes_precedence_and_normalizes_input(): + assert SELECTOR.select_amd_test_suites( + branch="main", + labels=("ready",), + debug_test_yaml=" MERGE, ready,", + ) == ("merge", "ready") + + +def test_empty_debug_override_uses_normal_selection(): + assert SELECTOR.select_amd_test_suites( + branch="feature", + labels=("merge-test",), + debug_test_yaml="", + ) == ("merge",) + + +@pytest.mark.parametrize("value", ["ready,ready", "nightly", ", ,", " \t"]) +def test_invalid_debug_override(value): + with pytest.raises(ValueError): + SELECTOR.select_amd_test_suites( + branch="feature", + labels=(), + debug_test_yaml=value, + ) From 06f67695f1819f795956b8d51f5c982aa082937d Mon Sep 17 00:00:00 2001 From: shiyichuan <93317314+CarrotSwordsman@users.noreply.github.com> Date: Wed, 9 Sep 2026 22:05:44 +0800 Subject: [PATCH 35/93] [CI/Build] Avoid duplicate stage CLI deploy config (#7007) Signed-off-by: mershi Co-authored-by: mershi --- tests/helpers/runtime.py | 1 - tests/helpers/tests/test_runtime_stage_cli.py | 38 ++++++++++++++++++- 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index db4cba893f3..66e0cefab75 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -968,7 +968,6 @@ def iter_omni_server( raise ValueError("omni_server with use_stage_cli=True requires use_omni=True") if stage_config_path is None: raise ValueError("omni_server with use_stage_cli=True requires a stage_config_path") - server_args += ["--deploy-config", stage_config_path] with OmniServerStageCli( model, diff --git a/tests/helpers/tests/test_runtime_stage_cli.py b/tests/helpers/tests/test_runtime_stage_cli.py index 564f12cf154..712a87e0936 100644 --- a/tests/helpers/tests/test_runtime_stage_cli.py +++ b/tests/helpers/tests/test_runtime_stage_cli.py @@ -5,12 +5,15 @@ from __future__ import annotations +import threading +from contextlib import nullcontext from pathlib import Path from types import SimpleNamespace import pytest -from tests.helpers.runtime import OmniServerStageCli +from tests.helpers import runtime +from tests.helpers.runtime import OmniServerParams, OmniServerStageCli pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -49,6 +52,39 @@ def test_stage_cli_loads_stage_ids_and_replica_counts(tmp_path: Path) -> None: assert server.stage_replica_counts == {0: 1, 1: 3} +def test_stage_cli_fixture_leaves_deploy_config_to_server(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + stage_config_path = _write_stage_config(tmp_path) + captured_args: list[str] = [] + + def fake_stage_cli(model, config_path, serve_args, **kwargs): + del kwargs + assert config_path == stage_config_path + captured_args.extend(serve_args) + return nullcontext(SimpleNamespace(model=model)) + + monkeypatch.setattr(runtime, "OmniServerStageCli", fake_stage_cli) + monkeypatch.setattr(runtime, "_whisper_device_free_around", nullcontext) + monkeypatch.setattr( + "tests.helpers.stage_config.stage_config_path_for_run_level", + lambda config_path, _run_level: config_path, + ) + request = SimpleNamespace( + param=OmniServerParams( + model="fake-model", + stage_config_path=stage_config_path, + use_stage_cli=True, + ), + node=SimpleNamespace(get_closest_marker=lambda _name: None), + ) + + server_fixture = runtime.iter_omni_server(request, "core_model", threading.Lock()) + next(server_fixture) + with pytest.raises(StopIteration): + next(server_fixture) + + assert "--deploy-config" not in captured_args + + def test_stage_cli_enter_rolls_back_launched_stages_on_startup_failure( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: From b9932aa17b0ea5c5e8ec49681eef656403285344 Mon Sep 17 00:00:00 2001 From: andyluo7 <43718156+andyluo7@users.noreply.github.com> Date: Wed, 9 Sep 2026 07:32:38 -0700 Subject: [PATCH 36/93] [CI/Build][ROCm] Normalize SenseNova paged-decode hardware markers (#6935) Signed-off-by: andyluo7 --- .../test_sensenova_u1_paged_decode.py | 34 +++++++++---------- 1 file changed, 16 insertions(+), 18 deletions(-) diff --git a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py index 368852a1d27..ac517945d90 100644 --- a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py +++ b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_paged_decode.py @@ -17,6 +17,7 @@ from vllm.lora.layers import BaseLayerWithLoRA import vllm_omni.diffusion.models.sensenova_u1.paged_decode as paged_decode +from tests.helpers.mark import hardware_test from vllm_omni.diffusion.models.sensenova_u1.paged_decode import ( BLOCK_SIZE, BUCKETS, @@ -169,7 +170,8 @@ def test_length_lives_in_device_tensors(): # --------------------------------------------------------------------------- -# The capture contract. These need CUDA and the bundled flash-attention kernel. +# The capture contract. Graph tests need a torch.cuda-compatible accelerator; +# paged-attention equivalence also needs the bundled kernel's paged API. # --------------------------------------------------------------------------- cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -181,8 +183,7 @@ def test_length_lives_in_device_tensors(): @cuda_only @vllm_flash_attn_only -@pytest.mark.cuda -@pytest.mark.L4 +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) def test_paged_attention_matches_sdpa_over_the_same_buffer(): import torch.nn.functional as F @@ -223,8 +224,7 @@ def test_paged_attention_matches_sdpa_over_the_same_buffer(): @cuda_only @vllm_flash_attn_only -@pytest.mark.cuda -@pytest.mark.L4 +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) def test_a_captured_graph_follows_a_later_length(): """One capture has to serve every length in the bucket. @@ -285,8 +285,7 @@ def test_a_captured_graph_follows_a_later_length(): @cuda_only @vllm_flash_attn_only -@pytest.mark.cuda -@pytest.mark.L4 +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) def test_a_captured_graph_writes_each_step_to_its_own_slot(): """The write slot has to be a device tensor, not a Python index. @@ -347,8 +346,7 @@ def test_a_captured_graph_writes_each_step_to_its_own_slot(): @cuda_only @vllm_flash_attn_only -@pytest.mark.cuda -@pytest.mark.L4 +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) def test_a_whole_decode_loop_matches_the_unpaged_path(): """Per-step kernel equivalence is not loop equivalence. @@ -413,8 +411,7 @@ def forward(self, input_ids, indexes, past_key_values=None, use_cache=False, pag @cuda_only -@pytest.mark.cuda -@pytest.mark.L4 +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) def test_the_decode_graph_captures_into_the_shared_platform_pool(monkeypatch): """``_decode_context`` builds a runner per request, so a private capture pool hands every request its own arena and never gives it back. Measured on a @@ -556,10 +553,12 @@ def current_kernel( def test_only_q_k_and_v_reach_the_kernel_positionally(monkeypatch): """A positional standard argument binds to whatever a future wheel puts in that slot, and the probe above only checks that the names exist.""" - seen = {} + seen_args: tuple[object, ...] = () + seen_kwargs: dict[str, object] = {} - def recorder(*args, **kwargs): - seen["args"], seen["kwargs"] = args, kwargs + def recorder(*args: object, **kwargs: object): + nonlocal seen_args, seen_kwargs + seen_args, seen_kwargs = args, kwargs return torch.zeros(1, N_HEADS, HEAD_DIM) monkeypatch.setattr(paged_decode, "_flash_varlen", lambda: recorder) @@ -572,8 +571,8 @@ def recorder(*args, **kwargs): torch.zeros(1, KV_HEADS, 1, HEAD_DIM), 1.0, ) - assert len(seen["args"]) == 3, "only q, k and v may be positional" - assert set(seen["kwargs"]) == { + assert len(seen_args) == 3, "only q, k and v may be positional" + assert set(seen_kwargs) == { "max_seqlen_q", "cu_seqlens_q", "max_seqlen_k", @@ -694,8 +693,7 @@ def test_no_capture_crosses_a_request_once_lora_wrappers_are_installed(monkeypat @cuda_only -@pytest.mark.cuda -@pytest.mark.L4 +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) def test_a_second_request_replays_the_first_capture(monkeypatch): """Reuse is only worth anything if the capture survives it.""" dev = torch.device("cuda") From 6fb7b0a05e3758e41e3871804598cea18a0b42fb Mon Sep 17 00:00:00 2001 From: hyw <109567717+yuweih205@users.noreply.github.com> Date: Wed, 9 Sep 2026 22:34:46 +0800 Subject: [PATCH 37/93] [Model] Skip unused frame packing in Wan2.2 S2V (#7155) Signed-off-by: hyw --- .../diffusion/models/wan2_2/wan2_2_s2v_transformer.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.py index d1b74d020ab..cce5e43be8b 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.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 """ Wan2.2 Speech-to-Video (S2V) Transformer using vllm-omni ops. @@ -1359,11 +1359,9 @@ def process_motion(self, motion_latents, drop_motion_frames=False): return flatten_mot, mot_remb def process_motion_frame_pack(self, motion_latents, drop_motion_frames=False, add_last_motion=2): - flatten_mot, mot_remb = self.frame_packer(motion_latents, add_last_motion) if drop_motion_frames: - return [m[:, :0] for m in flatten_mot], [m[:, :0] for m in mot_remb] - else: - return flatten_mot, mot_remb + return [], [] + return self.frame_packer(motion_latents, add_last_motion) def process_motion_transformer_motioner(self, motion_latents, drop_motion_frames=False, add_last_motion=True): batch_size = motion_latents.shape[0] From 504e5a6a6e7d3d00105a1c04ba5bd84fd62d5d15 Mon Sep 17 00:00:00 2001 From: bojiang-li <327132355+bojiang-li@users.noreply.github.com> Date: Wed, 9 Sep 2026 11:48:41 -0700 Subject: [PATCH 38/93] [Doc] Add dual DGX Spark MiniMax-H3 results (#7343) Signed-off-by: bojiang-li <327132355+bojiang-li@users.noreply.github.com> --- recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md | 54 ++++++++++++++++++---- recipes/README.md | 1 + 2 files changed, 47 insertions(+), 8 deletions(-) diff --git a/recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md b/recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md index b9d32947b8f..efd88c06849 100644 --- a/recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md +++ b/recipes/MiniMaxAI/MiniMax-H3-Spark-GB10.md @@ -7,13 +7,14 @@ recipes it uses **no offload of any kind** — see the capacity note below. Validated on: - Host: DGX Spark (GB10), aarch64 -- GPUs: 1 (unified memory) -- Driver: +- GPUs: 1, or 2 hosts with one unified-memory GPU each +- Driver: 580.173.02 for the two-host qualification - vLLM: 0.26.0 - vLLM-Omni: `main` at `e1aa6eae75c460cd1893bc320546e81e66973831` -- Workloads: T2VA and Ref2VA, 960x576, `duration=8.0`, 50 steps, `flow_shift=12`, - `seed=1101` - +- One-host workloads: T2VA and Ref2VA, 960x576, `duration=8.0`, 50 steps +- Two-host workload: T2VA, 1344x768, `duration=5.0` and `duration=10.0`, + 50 steps, text-encoder TP2, and DiT USP2 +- Common sampling settings: `flow_shift=12`, `seed=1101` ## Capacity requirements @@ -24,6 +25,7 @@ Validated on: | Checkpoint storage | 135 GiB per partition | | Measured peak allocator high-water, T2VA (FP8) | 97.7 GiB | | Measured peak allocator high-water, Ref2VA (FP8) | 102.8 GiB | +| Measured peak per rank, two-host T2VA 5 s / 10 s (FP8) | 75.45 / 87.17 GiB | `FL2VA` and `Ref2VA` are separate 135 GiB checkpoint partitions. Start one server at a time. @@ -164,6 +166,41 @@ curl --fail-with-body -sS -X POST http://127.0.0.1:8000/v1/videos/sync \ The response is `200 OK` with an MP4 body: H.264 at the requested geometry plus a 32 kHz stereo AAC track. `ffprobe` reports 8.032 s for `duration=8.0`. +## Two GB10 hosts: T2VA at 1344x768 + +The two-host qualification uses one GB10 GPU per DGX Spark, connected through a +single ConnectX-7 RoCE interface. The FL2VA checkpoint and software revision are +identical on both hosts. The distributed configuration is: + +- world size 2; +- text-encoder tensor parallel size 2; +- DiT Ulysses sequence parallel size 2 and ring degree 1; +- VAE patch parallel size 1; +- online FP8, eager execution, cuDNN attention, and tiled VAE decode. + +The current diffusion multiprocess executor starts all ranks on one host and sets +`MASTER_ADDR=localhost`. These measurements therefore used an external `env://` +launcher with one rank on each host and a benchmark-only adaptation that maps +both hosts' local worker to `cuda:0`. This qualifies the two-rank MiniMax-H3 +pipeline and its output, but it is not yet a supported cross-host `vllm serve` +command. Do not infer that the single-host launch command above works unchanged +across two machines. + +One 5-second, 2-step request warmed both workers before the measured runs. The +reported pipeline elapsed time covers distributed model execution through GPU +completion; MP4 encoding happened afterward and is not included. + +| Duration | Steps | Text encode | Denoise | Per step | VAE decode | Pipeline elapsed | Max peak per rank | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| 5 s | 50 | 0.21 s | 1570.19 s | 31.40 s | 79.45 s | 1649.87 s (27 min 30 s) | 75.45 GiB | +| 10 s | 50 | 0.15 s | 4448.12 s | 88.96 s | 160.71 s | 4609.00 s (76 min 49 s) | 87.17 GiB | + +Both ranks reported stage times and allocator peaks within measurement noise. +The generated outputs were H.264 at 1344x768 and 24 FPS with 32 kHz stereo AAC: +124 frames (5.167 s) and 243 frames (10.125 s). Both files decoded successfully, +and neither run encountered an OOM or NCCL error. These are single qualification +runs (`n=1`), not throughput guarantees. + ## Measured latency Single GB10, FP8, eager, cuDNN attention, tiled VAE, one request at a time: @@ -238,9 +275,11 @@ shorter than a 50-step Ref2VA run. ## Known limitations -- One GPU means no Ulysses, ring, TP, or VAE patch parallelism. Pass +- With one GPU there is no Ulysses, ring, TP, or VAE patch parallelism. Pass `--usp 1 --ring 1 --vae-patch-parallel-size 1` explicitly if a profile or script would otherwise supply higher degrees. +- The two-host results require the external launch adaptation described above; + cross-host diffusion worker launch is not exposed by the documented server CLI. - `--enforce-eager` is used because regional compile adds startup time without a parallelism win on a single rank. - `t2va` rejects requests without an explicit `aspect_ratio`. When `width` and @@ -255,8 +294,7 @@ shorter than a 50-step Ref2VA run. adding references or raising the output shape. - Ref2VA requires a vLLM-Omni build newer than `v0.26.0`. See the version note under **Validated on**: FP8 weight loading and image-only Ref2VA are both - broken on `release/v0.26.0`.suggestion is `v0.26.1`. + broken on `release/v0.26.0`. The suggested version is `v0.26.1`. - Online FP8 is incompatible with layerwise offload — the offload path produces a weight stride the Cutlass FP8 kernel rejects. This is not a practical restriction here since offload is unusable on GB10 anyway. -``` diff --git a/recipes/README.md b/recipes/README.md index 6b80f8a0343..ea45762014c 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -55,6 +55,7 @@ recipes/ | [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU | | [`IndexTeam/IndexTTS-2_5.md`](./IndexTeam/IndexTTS-2_5.md) | Online + offline multilingual voice-cloned TTS with native speed and emotion control | 1x NVIDIA H20 96GB | | [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 | +| [`MiniMaxAI/MiniMax-H3-Spark-GB10.md`](./MiniMaxAI/MiniMax-H3-Spark-GB10.md) | MiniMax H3 T2VA and Ref2VA on DGX Spark | 1x or 2x DGX Spark (GB10) | | [`MiniMaxAI/MiniMax-H3-MUSA.md`](./MiniMaxAI/MiniMax-H3-MUSA.md) | MiniMax H3 T2VA, FL2VA, and Ref2VA on MUSA | 4x MTT S5000 (TP4 Ref2VA profile) | | [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB | | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU | From b02862beee2f086e7f46df5aff3ca06026041045 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Wed, 9 Sep 2026 13:02:12 -0700 Subject: [PATCH 39/93] [Model] Optimize MOSS-TTS Local batched execution and streaming codec (#7202) Signed-off-by: Sy03 <1370724210@qq.com> --- .buildkite/cuda/test-ready.yml | 2 +- .../online_serving/text_to_speech/README.md | 53 +++++++- recipes/OpenMOSS/MOSS-TTS.md | 12 +- .../test_seed_tts_dataset_variants.py | 20 ++- .../openai_api/test_tts_adapter.py | 50 ++++++++ tests/entrypoints/test_omni_entrypoints.py | 26 +++- .../moss_tts/test_local_depth_rope_cache.py | 15 ++- .../moss_tts/test_streaming_attention.py | 58 +++++++++ .../moss_tts/test_streaming_terminal_batch.py | 89 +++++++++++++ tests/worker/test_batched_omni_output.py | 76 +++++++++++ .../benchmarks/data_modules/seed_tts_eval.py | 25 ++-- vllm_omni/deploy/moss_tts_local_h200.yaml | 64 ++++++++++ vllm_omni/entrypoints/omni_base.py | 9 +- .../openai/tts_adapters/moss_tts.py | 28 +++-- vllm_omni/entrypoints/utils.py | 4 +- vllm_omni/metrics/stats.py | 15 ++- .../models/moss_tts/audio_tokenizer_v2.py | 14 ++- .../moss_tts/modeling_moss_tts_codec.py | 42 +++++-- .../moss_tts/modeling_moss_tts_local_depth.py | 80 ++++++++---- .../moss_tts/modeling_moss_tts_talker.py | 28 +++-- .../models/moss_tts/pipeline.py | 9 ++ .../models/moss_tts/streaming_attention.py | 119 ++++++++++++++++++ vllm_omni/outputs/__init__.py | 5 +- vllm_omni/worker/gpu_ar_model_runner.py | 61 +++++++-- vllm_omni/worker/gpu_model_runner.py | 53 +++++--- 25 files changed, 843 insertions(+), 114 deletions(-) create mode 100644 tests/model_executor/models/moss_tts/test_streaming_attention.py create mode 100644 tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py create mode 100644 tests/worker/test_batched_omni_output.py create mode 100644 vllm_omni/deploy/moss_tts_local_h200.yaml create mode 100644 vllm_omni/model_executor/models/moss_tts/streaming_attention.py diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index d7e14ed64c1..f54a58fa48e 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.buildkite/cuda/test-ready.yml @@ -93,7 +93,7 @@ steps: commands: - | timeout 15m bash -c " - pytest -sv tests/engine/ tests/model_executor/ -m 'core_model and cuda' --run-level 'core_model' + pytest -sv tests/engine/ tests/model_executor/ tests/worker/test_batched_omni_output.py -m 'core_model and cuda' --run-level 'core_model' " mirror_hardwares: l4_1 diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index 5dbe583146b..5141a37c2bd 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -259,21 +259,25 @@ python examples/online_serving/text_to_speech/indextts2/speech_client.py \ 0.6B DualAR TTS at 44.1 kHz, 11 languages, zero-shot voice cloning. ### Prerequisites + None beyond the base install: the neural audio codec is implemented in tree and its weights (`codec.pth`) ship with the checkpoint. ### Launch + ```bash vllm serve Audio8/Audio8-TTS-Preview-0.6b --omni --port 8092 # or: ./audio8_tts/run_server.sh ``` + The deploy config auto-loads from `vllm_omni/deploy/audio8_tts.yaml` (HF `model_type` is `arktts`). Do **not** pass `--trust-remote-code`: vllm-omni registers its own `arktts` config, and transformers would otherwise prefer the checkpoint's remote code and bypass it. ### Text-only synthesis + ```bash curl -X POST http://localhost:8092/v1/audio/speech \ -H "Content-Type: application/json" \ @@ -284,6 +288,7 @@ curl -X POST http://localhost:8092/v1/audio/speech \ ``` ### Voice cloning + ```bash curl -X POST http://localhost:8092/v1/audio/speech \ -H "Content-Type: application/json" \ @@ -293,18 +298,21 @@ curl -X POST http://localhost:8092/v1/audio/speech \ "ref_text": "The exact transcript of the reference recording." }' --output cloned.wav ``` + `ref_audio` also accepts a `data:audio/wav;base64,...` URL. `ref_text` is mandatory whenever `ref_audio` is present and must match the recording. Uploading a voice via `POST /v1/audio/voices` lets the server reuse the encoded reference codes across requests (`voice: ""`). ### Streaming + ```bash python audio8_tts/speech_client.py --text "Welcome to Audio8 TTS." --stream --output out.pcm ffplay -f s16le -ar 44100 -ac 1 out.pcm ``` ### Gradio demo + ```bash ./audio8_tts/run_gradio_demo.sh # server + demo python audio8_tts/gradio_demo.py --api-base http://localhost:8092 # demo only @@ -328,7 +336,7 @@ these as a lower bound. 20 prompts from `seed_tts_smoke/en`, mean audio 4.5 s, deploy defaults: | | HF reference (batch=1) | vllm-omni c=1 | c=4 | c=8 | -|---|---|---|---|---| +| --- | --- | --- | --- | --- | | RTF (mean) | 1.008 | **0.19** | 0.30 | 0.54 | | Throughput (req/s) | 0.23 | 1.18 | 2.84 | **2.93** | | E2E latency (mean, ms) | 4375 | 850 | 1349 | 2361 | @@ -347,7 +355,7 @@ job holding ~80% SM, no repeats), so read them as directional, not as measured speedups. Same setup, varying stage 0's `max_num_seqs` (stage 1 stays at 1): | stage-0 `max_num_seqs` | client concurrency | req/s | RTF | TTFP mean / p99 (ms) | underrun p99 (s) | -|---|---|---|---|---|---| +| --- | --- | --- | --- | --- | --- | | 4 (default) | 1 | 1.18 | 0.19 | 63 / 67 | 0.00 | | 4 | 4 | 2.84 | 0.30 | 105 / 148 | 0.00 | | 4 | 8 | 2.93 | 0.54 | 1124 / 1642 | 0.00 | @@ -376,6 +384,7 @@ underrun. Raise it to 8 if first-packet latency under load matters more than a rare sub-300 ms gap. ### Notes + - Output: 44.1 kHz mono; ~21.5 codec frames per second. - No built-in speaker presets. Omit `voice` for a random timbre, or clone one. - `max_new_tokens` caps generated codec frames (1 frame ~= 46 ms). @@ -553,6 +562,46 @@ python examples/online_serving/text_to_speech/ming_flash_omni_tts/speech_client. --- +## MOSS-TTS Local Transformer v1.5 + +For a single H200, the optional +[`moss_tts_local_h200.yaml`](../../../vllm_omni/deploy/moss_tts_local_h200.yaml) +deployment places both the talker and codec on logical GPU 0: + +```bash +CUDA_VISIBLE_DEVICES=0 VLLM_OMNI_EVENT_DRIVEN_ORCH=1 \ + vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 \ + --omni --trust-remote-code \ + --deploy-config vllm_omni/deploy/moss_tts_local_h200.yaml \ + --disable-log-stats +``` + +Run from the repository root and select an available physical GPU with +`CUDA_VISIBLE_DEVICES`. The command enables event-driven orchestration and +disables detailed per-request statistics logging to reduce CPU overhead at +high concurrency. Remove `--disable-log-stats` when those statistics are needed; +keep these settings identical when comparing performance. +This preset configures 256 request slots per stage, +a 32 GiB talker KV cache, talker CUDA Graph buckets through 512 scheduled tokens, +and codec batch buckets through 256. It requires more than 80 GiB of GPU memory; +request capacity also depends on input and generated lengths. The default +[`moss_tts_local.yaml`](../../../vllm_omni/deploy/moss_tts_local.yaml) remains +available for smaller deployments. + +The preset enables the optional codec backend with +`hf_overrides.codec_attention_backend: triton` on stage 1. It preserves the +streaming ring-cache mask and uses BF16 attention with 64-dimensional heads; +other attention shapes use PyTorch SDPA. Set the backend to `sdpa` to use the +default attention implementation. Codec terminal tails share execution only +when they already map to the same padded CUDA Graph; returned audio retains +each request's actual length. + +Voice cloning requests use `ref_audio` and `ref_text`. For streaming output, +set `stream: true`, `stream_format: "audio"`, and `response_format: "pcm"`; +the native PCM format is 48 kHz, stereo, signed 16-bit little-endian. Compute +audio throughput as `PCM bytes / (48000 * 2 * 2) / elapsed seconds`, and compare +the same requests, concurrency, and reference-cache state. + ## MOSS-TTS-Nano Single-stage 0.1B AR LM + MOSS-Audio-Tokenizer-Nano codec at 48 kHz mono. Every request must include `ref_audio`; there are no built-in speaker presets. diff --git a/recipes/OpenMOSS/MOSS-TTS.md b/recipes/OpenMOSS/MOSS-TTS.md index e9cb798fe27..8d37039d80c 100644 --- a/recipes/OpenMOSS/MOSS-TTS.md +++ b/recipes/OpenMOSS/MOSS-TTS.md @@ -16,7 +16,7 @@ Use this recipe for 24 kHz multilingual TTS with voice cloning (20 languages inc Chinese and English). Choose a variant based on your latency and quality requirements: | Model | Params | Use case | -|---|---|---| +| --- | --- | --- | | MOSS-TTS | 8B | General TTS, highest quality | | MOSS-TTS-v1.5 | 8B | General TTS upgrade of 1.0: 31 languages, steadier cloning, `[pause Xs]` markers (set `language` for best results); same `MossTTSDelay` API | | MOSS-TTS-Realtime | 1.7B | Lowest latency (TTFB ~180 ms), streaming-first | @@ -24,9 +24,17 @@ Chinese and English). Choose a variant based on your latency and quality require | MOSS-SoundEffect | 8B | Sound effect synthesis from text description | | MOSS-VoiceGenerator | 1.7B | Zero-shot voice design | -All variants share the same codec (`OpenMOSS-Team/MOSS-Audio-Tokenizer`, ~7 GB) and +The variants above share the same codec (`OpenMOSS-Team/MOSS-Audio-Tokenizer`, ~7 GB) and output 24 kHz mono audio. +MOSS-TTS-Local-Transformer-v1.5 uses MOSS-Audio-Tokenizer-v2 and outputs 48 kHz +stereo audio. For Local voice cloning through `/v1/audio/speech`, provide an +accurate `ref_text` transcript alongside `ref_audio`. The adapter uses the +reference transcript and audio as a continuation prefix, then generates only +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. + ## References - Offline inference example: [`examples/offline_inference/text_to_speech/moss_tts/`](../../examples/offline_inference/text_to_speech/moss_tts/) diff --git a/tests/benchmarks/test_seed_tts_dataset_variants.py b/tests/benchmarks/test_seed_tts_dataset_variants.py index 31f66bd7d73..17ee6df4c14 100644 --- a/tests/benchmarks/test_seed_tts_dataset_variants.py +++ b/tests/benchmarks/test_seed_tts_dataset_variants.py @@ -26,6 +26,7 @@ if _MODULE_NAME not in sys.modules: _spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH) + assert _spec is not None and _spec.loader is not None _mod = importlib.util.module_from_spec(_spec) sys.modules[_MODULE_NAME] = _mod _spec.loader.exec_module(_mod) @@ -237,7 +238,8 @@ def test_seed_tts_design_dataset_rejects_missing_description(seed_tts_design_roo assert req.seed_tts_utterance_id == "ok" -def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch): +@pytest.mark.parametrize("audio_seconds", [0.1, 61]) +def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch, audio_seconds): from vllm_omni.benchmarks.data_modules import seed_tts_eval calls = {} @@ -252,8 +254,9 @@ def to(self, device): return self class FakeProcessor: - def __call__(self, wav, *, sampling_rate, return_tensors, return_attention_mask=False): + def __call__(self, wav, *, sampling_rate, return_tensors, return_attention_mask=False, **kwargs): calls["return_attention_mask"] = return_attention_mask + calls["processor_kwargs"] = kwargs assert sampling_rate == 16000 assert return_tensors == "pt" assert len(wav) > 0 @@ -262,9 +265,10 @@ def __call__(self, wav, *, sampling_rate, return_tensors, return_attention_mask= attention_mask=FakeTensor("mask") if return_attention_mask else None, ) - def get_decoder_prompt_ids(self, *, language, task): + def get_decoder_prompt_ids(self, *, language, task, no_timestamps=True): assert language == "english" assert task == "transcribe" + calls["no_timestamps"] = no_timestamps return [(1, 2)] def batch_decode(self, predicted_ids, *, skip_special_tokens): @@ -283,10 +287,18 @@ def generate(self, input_features, **kwargs): monkeypatch.setattr(seed_tts_eval, "_en_model", FakeModel()) monkeypatch.setattr(seed_tts_eval, "_device", "cuda:1") - text = seed_tts_eval._transcribe_en_f32_16k(np.ones(1600, dtype=np.float32)) + text = seed_tts_eval._transcribe_en_f32_16k(np.ones(int(audio_seconds * 16000), dtype=np.float32)) assert text == "hello" assert calls["return_attention_mask"] is True assert calls["input_features"].device == "cuda:1" assert calls["generate_kwargs"]["attention_mask"].device == "cuda:1" assert calls["generate_kwargs"]["forced_decoder_ids"] == [(1, 2)] + if audio_seconds > 30: + assert calls["processor_kwargs"] == {"truncation": False, "padding": "longest"} + assert calls["generate_kwargs"]["return_timestamps"] is True + assert calls["no_timestamps"] is False + else: + assert calls["processor_kwargs"] == {} + assert "return_timestamps" not in calls["generate_kwargs"] + assert calls["no_timestamps"] is True diff --git a/tests/entrypoints/openai_api/test_tts_adapter.py b/tests/entrypoints/openai_api/test_tts_adapter.py index a2c0ace0be7..a17a6ccbdb3 100644 --- a/tests/entrypoints/openai_api/test_tts_adapter.py +++ b/tests/entrypoints/openai_api/test_tts_adapter.py @@ -166,6 +166,56 @@ def test_moss_tts_seed_falls_back_to_stage_default(adapter_cls, mocker): assert prepared.tts_params["seed"] == [42] +@pytest.mark.parametrize( + "variant,ref_text,mode", + [ + ("local", " Reference. ", "continuation"), + ("local", None, "generation"), + ("local", " ", "generation"), + ("tts", "Reference.", "generation"), + ], +) +def test_moss_reference_transcript_mode(variant, ref_text, mode, mocker): + import torch + + reference = [torch.ones((3, 12), dtype=torch.int64)] + unified = torch.arange(52, dtype=torch.int64).reshape(1, 4, 13) + processor = mocker.Mock(return_value={"input_ids": unified}) + server = mocker.Mock(_moss_variant="local", uploaded_speakers={"speaker": {}}) + server._voice_created_at.return_value = 123 + request = OpenAICreateSpeechRequest( + input="Target.", ref_text=ref_text, language="English", voice="Speaker", seed=0, max_new_tokens=2048 + ) + adapter = MossTTSAdapter(SpeechServingContext(server=server)) + + adapter._moss_variant = variant + mocker.patch.object(adapter, "_get_moss_processor", return_value=processor) + mocker.patch.object( + adapter, "_encode_moss_references", new=mocker.AsyncMock(return_value=(reference, {0: "reference-key"})) + ) + + prepared = asyncio.run(adapter.build(request, [], has_inline_ref_audio=False)) + + if mode == "continuation": + processor.build_user_message.assert_called_once_with(text="Reference. Target.", language="English") + processor.build_assistant_message.assert_called_once_with(audio_codes_list=reference) + conversation = [processor.build_user_message.return_value, processor.build_assistant_message.return_value] + else: + processor.build_user_message.assert_called_once_with(text="Target.", language="English", reference=reference) + processor.build_assistant_message.assert_not_called() + conversation = [processor.build_user_message.return_value] + processor.assert_called_once_with(conversations=[conversation], mode=mode) + assert prepared.prompt["prompt_token_ids"] == unified[0, :, 0].tolist() + assert torch.equal(prepared.tts_params["codes"]["ref"], unified[0, :, 1:]) + assert prepared.tts_params["max_new_frames"] == [2048] + assert prepared.tts_params["ref_audio_cache_key"] == "reference-key" + assert prepared.tts_params["voice_name"] == ["speaker"] + assert prepared.tts_params["voice_created_at"] == [123] + assert prepared.tts_params["seed"] == [0] + assert prepared.prompt["cache_salt"] + assert request.input == "Target." + + def test_qwen3_tts_metadata(): assert Qwen3TTSAdapter.backend == "ar" assert issubclass(Qwen3TTSAdapter, ARTTSAdapter) diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 4f9fc592b4c..6fe04fdb7c6 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.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 queue @@ -97,9 +100,11 @@ def __init__( ) -> None: self.model = model self.config_path = None - self.stage_configs: list[Any] = [] self.stage_metadata = stage_metadata or [THREE_STAGE_META[-1]] self.num_stages = len(self.stage_metadata) + self.stage_configs = [ + StageConfig(stage_id=i, model_stage="dummy-model").to_omegaconf() for i in range(self.num_stages) + ] self.default_sampling_params_list = default_sampling_params_list or [ SamplingParams(max_tokens=8) for _ in range(self.num_stages) ] @@ -222,6 +227,25 @@ def test_resolve_sampling_params_list_preserves_stage_constraints(): assert caller_params.stop_token_ids == [7] +@pytest.mark.parametrize("use_defaults", [False, True]) +def test_moss_local_output_policy_preserves_codec_streaming(use_defaults): + from vllm_omni.model_executor.models.moss_tts.pipeline import MOSS_TTS_LOCAL_PIPELINE + + base = _make_base() + base.engine.num_stages = 2 + base.sampling_constraints_list = [stage.sampling_constraints for stage in MOSS_TTS_LOCAL_PIPELINE.stages] + base.default_sampling_params_list = [ + base._apply_sampling_constraints(SamplingParams(), constraints) + for constraints in base.sampling_constraints_list + ] + caller = [SamplingParams(output_kind=RequestOutputKind.DELTA) for _ in range(2)] + result = base.resolve_sampling_params_list(None if use_defaults else caller, allow_delta_coercion=True) + + assert [params.output_kind for params in result] == [RequestOutputKind.FINAL_ONLY, RequestOutputKind.DELTA] + assert all(params.output_kind == RequestOutputKind.DELTA for params in caller) + assert base.default_sampling_params_list[0].output_kind == RequestOutputKind.FINAL_ONLY + + def _stage_spec( stage_id: int, *, diff --git a/tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py b/tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py index 90c386ee08a..566ba4595b2 100644 --- a/tests/model_executor/models/moss_tts/test_local_depth_rope_cache.py +++ b/tests/model_executor/models/moss_tts/test_local_depth_rope_cache.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 """CPU parity tests for MOSS-TTS Local Depth RoPE caching.""" import pytest @@ -42,3 +42,16 @@ def test_rope_cache_matches_reference_and_reuses_storage() -> None: # Smaller slices retain the fixed allocation used by CUDA Graph. attn.prepare_rope_cache(6, device, torch.float32) assert attn._rope_cos_cache.data_ptr() == cache_ptr + + +def test_frame_local_kv_matches_full_prefix_and_reuses_slots() -> None: + model = MossTTSLocalDepthTransformer(_test_config()).eval() + model.h[0].attn.prepare_rope_cache(12, torch.device("cpu"), torch.float32) + cache = (torch.empty(2, 4, 12, 8), torch.empty(2, 4, 12, 8)) + with torch.inference_mode(): + for _ in range(2): + inputs = torch.randn(2, 12, 32) + for position in range(12): + expected = model._forward_prefix(inputs[:, : position + 1])[:, -1:] + actual = model._forward_prefix(inputs[:, position : position + 1], cache, position) + torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-5) diff --git a/tests/model_executor/models/moss_tts/test_streaming_attention.py b/tests/model_executor/models/moss_tts/test_streaming_attention.py new file mode 100644 index 00000000000..3134575ca5f --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_streaming_attention.py @@ -0,0 +1,58 @@ +# 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, + StreamingExecutionContext, +) +from vllm_omni.model_executor.models.moss_tts.streaming_attention import masked_attention + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + + +@pytest.mark.parametrize( + "batch,heads,queries,keys", [(2, 20, 1, 125), (4, 20, 15, 125), (1, 12, 480, 400), (4, 12, 32, 250)] +) +def test_masked_attention_matches_sdpa_with_strides_and_invalid_rows(batch, heads, queries, keys): + torch.manual_seed(7) + q = torch.randn(batch, queries, 3, heads, 64, device="cuda", dtype=torch.bfloat16)[:, :, 0].permute(0, 2, 1, 3) + k = torch.randn(batch, heads, keys, 64, device="cuda", dtype=torch.bfloat16) + v = torch.randn_like(k) + positions = torch.arange(keys, device="cuda").roll(17) + query_positions = torch.arange(queries, device="cuda") + keys - queries + mask = (positions[None, :] <= query_positions[:, None])[None, None].expand(batch, 1, -1, -1).clone() + mask[0, :, 0] = False # Cover an all-masked row alongside valid rows. + actual = masked_attention(q, k, v, mask) + expected = F.scaled_dot_product_attention(q, k, v, mask) + assert bool(actual.isfinite().all()) + torch.testing.assert_close(actual, expected, atol=0.008, rtol=0.01) + assert actual[0, :, 0].count_nonzero() == 0 + + +def test_streaming_attention_preserves_ring_wrap_reordering_and_slot_reset(): + torch.manual_seed(9) + reference = MossAudioTokenizerMultiheadAttention( + 128, 2, causal=True, context=9, device="cuda", dtype=torch.bfloat16 + ) + candidate = copy.deepcopy(reference) + candidate._streaming_attention = masked_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") + context = StreamingExecutionContext(state_slot_ids=slots, valid_rows=valid) + x = torch.randn(3, 3, 128, device="cuda", dtype=torch.bfloat16) + with torch.inference_mode(): + expected = reference(x, x, x, execution_context=context) + actual = candidate(x, x, x, execution_context=context) + torch.testing.assert_close(actual, expected, atol=0.01, rtol=0.02) + torch.testing.assert_close(candidate._streaming_state.offset, reference._streaming_state.offset, 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_terminal_batch.py b/tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py new file mode 100644 index 00000000000..fff97cd96bb --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_streaming_terminal_batch.py @@ -0,0 +1,89 @@ +# 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.moss_tts.modeling_moss_tts_codec import ( + MossTTSCodecDecoder, + _MossCodecStreamSession, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class CausalCodec(nn.Module): + downsample_rate = 1 + + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.zeros(())) + + def initialize_decoder_state_pool(self, capacity, scratch): + self.state = torch.zeros(capacity + scratch) + + 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] + return SimpleNamespace(audio=audio[:, None], audio_lengths=lengths) + + +def session(): + return _MossCodecStreamSession(CausalCodec(), state_capacity=8, n_vq=2, vllm_config=None) + + +def test_terminal_padding_preserves_audio_live_state_and_reused_slots(): + s = session() + assert [s.acquire(), s.acquire()] == [0, 1] + plan = {0: torch.ones(2, 3), 1: torch.ones(2, 15)} + with pytest.raises(ValueError, match="Only terminal"): + s.step(plan) + assert s._codec.state.count_nonzero() == 0 + s.step({0: torch.ones(2, 1), 1: torch.full((2, 1), 2)}) + actual = s.step(plan, terminal_slots={0}) + torch.testing.assert_close(actual[0], torch.tensor([[4.0, 6.0, 8.0]])) + torch.testing.assert_close(actual[1], torch.arange(6.0, 36.0, 2.0)[None]) + s.release(0, state_already_reset=True) + assert s.acquire() == 0 + resumed = s.step({0: torch.ones(2, 1), 1: torch.ones(2, 1)}) + 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): + s = session() + s._cudagraph_wrapper = ( + None + if graph_capacity is None + else SimpleNamespace( + batch_sizes=[1, graph_capacity], + _select_frame_size=lambda size, allow_padding: next((b for b in [1, 8, 15] if b >= size), None), + decode=lambda *args, **kwargs: None, + ) + ) + step = mocker.spy(s, "step") + 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 + lengths = [3, 7, 8, 1, 15] + finished = [True, True, False, True, False] + result = decoder._decode_streaming_batch( + [ + (i, str(i), torch.ones(2, length), done) + for i, (length, done) in enumerate(zip(lengths, finished, strict=True)) + ] + ) + assert step.call_count == expected_calls + 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. diff --git a/tests/worker/test_batched_omni_output.py b/tests/worker/test_batched_omni_output.py new file mode 100644 index 00000000000..d0b18b3c15f --- /dev/null +++ b/tests/worker/test_batched_omni_output.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import pytest +import torch + +from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import MossTTSLocalTalkerForGeneration +from vllm_omni.worker.gpu_ar_model_runner import _snapshot_tensor_payload_to_cpu_async + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + + +def test_async_packed_snapshot_survives_source_reuse(): + values = [ + torch.arange(12, device="cuda").view(1, 12), + torch.empty((0, 12), device="cuda", dtype=torch.long), + torch.arange(24, device="cuda").view(2, 12), + ] + expected = [v.cpu() for v in values] + snapshot = _snapshot_tensor_payload_to_cpu_async( + {"list": values, "tuple": tuple(values)}, copy_stream=torch.cuda.Stream(), pin_memory=True + ) + for v in values: + v.fill_(-99) + snapshot.wait() + torch.testing.assert_close(snapshot.payload, {"list": expected, "tuple": tuple(expected)}, rtol=0, atol=0) + + +@pytest.mark.parametrize("offsets", [[0, 1], [1, 4]]) +def test_mtp_preserves_rows_and_snapshot_without_cuda_sync(monkeypatch, offsets): + import vllm_omni.worker.gpu_model_runner as mod + from tests.worker.test_omni_gpu_model_runner import _make_runner, _noop_forward_context + + monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context) + runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4) + runner.model.gpu_resident_buffer_keys = {("codes", "audio")} + for name in ("talker_mtp_input_ids", "talker_mtp_inputs_embeds", "last_talker_hidden", "text_step"): + buffer = getattr(runner, name) + buffer.gpu = buffer.gpu.to("cuda") + embeds = torch.ones((2, 4), device="cuda") + codes = torch.arange(2, device="cuda").view(2, 1) + runner.talker_mtp = lambda *args, **kwargs: (embeds, codes) + output = torch.full((5, 4), -10.0, device="cuda") + previous = torch.cuda.get_sync_debug_mode() + try: + torch.cuda.set_sync_debug_mode("error") + runner._talker_mtp_forward(["r1", "r2"], output, offsets) + finally: + torch.cuda.set_sync_debug_mode(previous) + expected = torch.full((5, 4), -10.0) + expected[offsets] = 1.0 + torch.testing.assert_close(output.cpu(), expected, rtol=0, atol=0) + codes.fill_(999) + assert runner.model_intermediate_buffer["r1"]["codes"]["audio"].item() == 0 + assert runner.model_intermediate_buffer["r2"]["codes"]["audio"].item() == 1 + + +def test_mixed_stop_flags_do_not_synchronize_cuda(): + model = object.__new__(MossTTSLocalTalkerForGeneration) + torch.nn.Module.__init__(model) + model.n_vq = 12 + model.audio_pad_token_id = 1024 + valid = torch.arange(12, device="cuda").view(1, 12) + stopped = torch.full_like(valid, 1024) + infos = [{"audio_codes": {"current": valid}}, {}, {"audio_codes": {"current": stopped}}] + hidden = torch.ones((3, 4), device="cuda") + previous = torch.cuda.get_sync_debug_mode() + try: + torch.cuda.set_sync_debug_mode("error") + output = model.make_omni_output(hidden, runtime_additional_information=infos) + finally: + torch.cuda.set_sync_debug_mode(previous) + assert model._batch_should_continue.tolist() == [True, True, False] + torch.testing.assert_close(output.multimodal_outputs["codes"]["audio"], [valid, valid[:0], stopped]) + assert all("current" not in info.get("audio_codes", {}) for info in infos) + prefill = model.make_omni_output(hidden[:1], runtime_additional_information=[{}]) + assert model._batch_should_continue is None and prefill.multimodal_outputs == {} diff --git a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py index 9dc99548aeb..1009338700d 100644 --- a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py +++ b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py @@ -464,25 +464,33 @@ def _transcribe_en_f32_16k(wav_f32: np.ndarray) -> str: return "" with _lock: assert _en_processor is not None and _en_model is not None and _device is not None + # Whisper's default feature extraction truncates at 30 seconds. Keep + # abnormal tails so WER evaluates the complete generated response. + long_audio = len(wav_f32) > 30 * 16000 + processor_kwargs = {"truncation": False, "padding": "longest"} if long_audio else {} try: inputs = _en_processor( wav_f32, sampling_rate=16000, return_tensors="pt", return_attention_mask=True, + **processor_kwargs, ) except TypeError: - inputs = _en_processor(wav_f32, sampling_rate=16000, return_tensors="pt") + inputs = _en_processor(wav_f32, sampling_rate=16000, return_tensors="pt", **processor_kwargs) input_features = inputs.input_features.to(_device) attention_mask = getattr(inputs, "attention_mask", None) if attention_mask is None and isinstance(inputs, dict): attention_mask = inputs.get("attention_mask") generate_kwargs: dict[str, Any] = {} + if long_audio: + generate_kwargs["return_timestamps"] = True if attention_mask is not None: generate_kwargs["attention_mask"] = attention_mask.to(_device) with torch.no_grad(): try: - forced = _en_processor.get_decoder_prompt_ids(language="english", task="transcribe") + prompt_kwargs = {"no_timestamps": False} if long_audio else {} + forced = _en_processor.get_decoder_prompt_ids(language="english", task="transcribe", **prompt_kwargs) predicted_ids = _en_model.generate(input_features, forced_decoder_ids=forced, **generate_kwargs) except Exception: predicted_ids = _en_model.generate( @@ -548,14 +556,13 @@ def _expand_seed_tts_turn_outputs( turn_pcm = getattr(output, "tts_turn_pcm_bytes", None) session_pcm = getattr(output, "tts_output_pcm_bytes", None) for turn_index, turn in enumerate(turns): - expanded_requests.append( - replace( - request, - prompt=turn.target_text, - seed_tts_utterance_id=turn.utterance_id, - seed_tts_turns=(), - ) + turn_request = replace( + request, + seed_tts_utterance_id=turn.utterance_id, + seed_tts_turns=(), ) + turn_request.prompt = turn.target_text + expanded_requests.append(turn_request) turn_output = copy(output) if isinstance(turn_pcm, list) and turn_index < len(turn_pcm): turn_output.tts_output_pcm_bytes = turn_pcm[turn_index] diff --git a/vllm_omni/deploy/moss_tts_local_h200.yaml b/vllm_omni/deploy/moss_tts_local_h200.yaml new file mode 100644 index 00000000000..6fb5e010d4d --- /dev/null +++ b/vllm_omni/deploy/moss_tts_local_h200.yaml @@ -0,0 +1,64 @@ +# MOSS-TTS Local Transformer v1.5: both stages on one H200. +# Select one free physical GPU with CUDA_VISIBLE_DEVICES; devices below are logical. +# Capacity and graph capture require more than 80 GiB of GPU memory. +async_chunk: true + +connectors: + shm: + name: SharedMemoryConnector + extra: + connector_get_sleep_s: 0.005 + connector_get_max_wait_first_chunk: 1000 + connector_get_max_wait: 300 + initial_codec_chunk_frames: 1 + codec_chunk_frames: 15 + codec_left_context_frames: 0 + +stages: + - stage_id: 0 + max_num_seqs: 256 + gpu_memory_utilization: 0.60 + kv_cache_memory_bytes: 34359738368 # 32 GiB; leave room for codec state and graphs. + enable_prefix_caching: false + async_scheduling: true + max_num_batched_tokens: 65536 + max_model_len: 4096 + devices: "0" + compilation_config: + # Capture mixed prefill/decode batches by total scheduled tokens. + cudagraph_capture_sizes: [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 160, 192, 224, 256, 320, 384, 448, 512] + max_cudagraph_capture_size: 512 + output_connectors: + to_stage_1: shm + default_sampling_params: + temperature: 1.7 + top_p: 0.8 + top_k: 25 + max_tokens: 4096 + seed: 42 + repetition_penalty: 1.0 + + - stage_id: 1 + max_num_seqs: 256 + gpu_memory_utilization: 0.30 + enforce_eager: false + enable_prefix_caching: false + async_scheduling: true + compilation_config: + # Codec execution-batch buckets; frame counts come from the connector. + cudagraph_capture_sizes: [1, 2, 4, 8, 16, 32, 64, 128, 160, 192, 256] + cudagraph_num_of_warmups: 1 + max_num_batched_tokens: 65536 + max_model_len: 65536 + devices: "0" + hf_overrides: + codec_attention_backend: triton + input_connectors: + from_stage_0: shm + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + repetition_penalty: 1.0 diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 0aa6832d071..264b37fb527 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -339,6 +339,7 @@ def resolve_sampling_params_list( allow_delta_coercion: bool = False, ) -> Sequence[Any]: """Resolve request parameters; pipeline sampling constraints override caller values.""" + normalized: Sequence[Any] if sampling_params_list is None: normalized = self.default_sampling_params_list # Set the output kind to delta since no params were specified @@ -354,7 +355,9 @@ def resolve_sampling_params_list( if len(normalized) != self.num_stages: raise ValueError(f"Expected {self.num_stages} sampling params, got {len(normalized)}") - if sampling_params_list is not None: + # Streaming coercion may also change a constrained output kind in + # the defaults (for example, an internal TTS stage's FINAL_ONLY). + if sampling_params_list is not None or allow_delta_coercion: normalized = [ self._apply_sampling_constraints(params, constraints) for params, constraints in zip(normalized, self.sampling_constraints_list, strict=True) @@ -364,7 +367,7 @@ def resolve_sampling_params_list( @staticmethod def _get_sampling_constraints_list(stage_configs: Sequence[Any]) -> list[dict[str, Any]]: """Extract each stage's required sampling settings from runtime configs.""" - constraints_list = [] + constraints_list: list[dict[str, Any]] = [] for stage_config in stage_configs: constraints = getattr(stage_config, "sampling_constraints", {}) if not isinstance(constraints, Mapping): @@ -398,7 +401,7 @@ def _record_request_failure_once(self, request_id: str, reason: str) -> None: req_state = self.request_states.get(request_id) prom = getattr(self, "prom_metrics", None) metrics = getattr(req_state, "metrics", None) - if metrics is None or prom is None: + if req_state is None or metrics is None or prom is None: return if str(request_id) in metrics.e2e_done or getattr(req_state, "failure_recorded", False): return diff --git a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py index e264c7af564..07e4bb197e3 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py @@ -6,7 +6,7 @@ handles each); they are registered under distinct model-type names. """ -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast from vllm.inputs import tokens_input @@ -61,7 +61,7 @@ def _get_moss_ref_encoder(self): # the process-wide speaker cache. encoder = build_reference_encoder( self._get_moss_processor(), - variant=self._moss_variant, + variant=cast(str, self._moss_variant), speaker_cache=self._speaker_cache, ) self._moss_ref_encoder = encoder @@ -101,7 +101,7 @@ async def _encode_moss_references( return await encode_request_references( self._get_moss_ref_encoder(), - request.ref_audio, + cast(str, request.ref_audio), request.ref_audio_2 if two_speaker else None, resolve_ref_audio=self._resolve_ref_audio, get_artifact_key=self._get_resolved_ref_audio_artifact_key, @@ -179,15 +179,16 @@ async def _build_moss_tts_params( v = self._moss_variant + params: dict[str, Any] # ---- Legacy nano path (unchanged) ---- if v is None: # moss_tts_nano - params: dict[str, Any] = { + params = { "text": [request.input or ""], "mode": ["voice_clone"], } if request.max_new_tokens is not None: params["max_new_frames"] = [request.max_new_tokens] - wav_list, sr, cache_key = await self._resolve_ref_audio(request.ref_audio) + wav_list, sr, cache_key = await self._resolve_ref_audio(cast(str, request.ref_audio)) params["prompt_audio_array"] = [[wav_list, sr]] params["ref_audio_cache_key"] = cache_key return params @@ -201,13 +202,13 @@ async def _build_moss_tts_params( # talker reads for short prompts; full Realtime support needs a # separate processor.from_module path which we don't wire here. if v == "realtime": - params: dict[str, Any] = { + params = { "text": [request.input or ""], "mode": ["voice_clone"], } if request.max_new_tokens is not None: params["max_new_frames"] = [request.max_new_tokens] - wav_list, sr, cache_key = await self._resolve_ref_audio(request.ref_audio) + wav_list, sr, cache_key = await self._resolve_ref_audio(cast(str, request.ref_audio)) params["prompt_audio_array"] = [[wav_list, sr]] params["ref_audio_cache_key"] = cache_key return params @@ -254,13 +255,20 @@ async def _build_moss_tts_params( # 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). - user_msg = proc.build_user_message(**user_kwargs) - batch = proc(conversations=[[user_msg]], mode="generation") + 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") + 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) - params: dict[str, Any] = { + params = { "prompt_token_ids": text_ids, "codes": {"ref": audio_codes}, } diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index e1efd7ac4ab..39131fd58bd 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -278,9 +278,9 @@ def maybe_coerce_to_message_type(params: SamplingParams, is_streaming: bool): if params.output_kind == target_type: return params elif is_streaming and params.output_kind == RequestOutputKind.FINAL_ONLY: - logger.warning("Request appears to be streaming, but got request type final only!") + logger.debug("Coercing FINAL_ONLY output to DELTA for streaming") elif not is_streaming and params.output_kind == RequestOutputKind.DELTA: - logger.warning("Request appears to not be streaming, but got request type delta!") + logger.debug("Coercing DELTA output to FINAL_ONLY for non-streaming") if not params.skip_clone: params = params.clone() diff --git a/vllm_omni/metrics/stats.py b/vllm_omni/metrics/stats.py index ef86a5a3599..e38d83133ac 100644 --- a/vllm_omni/metrics/stats.py +++ b/vllm_omni/metrics/stats.py @@ -48,7 +48,7 @@ class StageRequestStats: finish_reason: str | None = None request_id: str | None = None postprocess_time_ms: float = 0.0 - diffusion_metrics: dict[str, float] = None + diffusion_metrics: dict[str, float] | None = None audio_generated_frames: int = 0 audio_sample_rate: int = 0 audio_duration_s: float = 0.0 @@ -181,11 +181,11 @@ def init_run_state(self, wall_start_ts: float) -> None: self.e2e_total_ms = 0.0 self.e2e_total_tokens = 0 self.e2e_count = 0 - self.e2e_done = set() + self.e2e_done: set[str] = set() self.wall_start_ts = float(wall_start_ts) self.last_finish_ts = float(wall_start_ts) - self.stage_first_ts = [None for _ in range(self.num_stages)] - self.stage_last_ts = [None for _ in range(self.num_stages)] + self.stage_first_ts: list[float | None] = [None for _ in range(self.num_stages)] + self.stage_last_ts: list[float | None] = [None for _ in range(self.num_stages)] self.accumulated_gen_time_ms: defaultdict[str, defaultdict[int, float]] = defaultdict( lambda: defaultdict(float) ) # {request_id: {stage_id:accumulated_gen_time_ms}} @@ -618,6 +618,7 @@ def on_stage_metrics( final_output_type: str | None = None, ) -> None: stats = self._as_stage_request_stats(stage_id, req_id, metrics, final_output_type) + assert stats.stage_id is not None self.stage_total_tokens[stats.stage_id] += int(stats.num_tokens_out) if stats.stage_id == 0: self.stage_total_tokens[stats.stage_id] += int(stats.num_tokens_in) @@ -775,10 +776,8 @@ def build_and_log_summary(self) -> dict[str, Any]: final_stage_id_map = self.final_stage_id_for_e2e stage_wall_time_ms = [ - ((self.stage_last_ts[i] - self.stage_first_ts[i]) * 1000.0) - if (self.stage_first_ts[i] is not None and self.stage_last_ts[i] is not None) - else 0.0 - for i in range(self.num_stages) + ((last - first) * 1000.0) if first is not None and last is not None else 0.0 + for first, last in zip(self.stage_first_ts, self.stage_last_ts, strict=True) ] overall_summary = { 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 380096af59a..4dbedbd19c1 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 @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -801,7 +804,16 @@ def forward( else: attn_bias = None - x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0) + 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) 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 231e3385701..c279493102c 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 @@ -145,15 +145,22 @@ def step( terminal_slots = terminal_slots or set() if not terminal_slots.issubset(slot_codes): raise ValueError("terminal_slots must be a subset of the decode batch slots.") - step_lengths = {int(codes.shape[1]) for codes in slot_codes.values()} - if len(step_lengths) != 1: - raise ValueError(f"MOSS codec streaming step needs uniform T, got {sorted(step_lengths)}") - (step_t,) = step_lengths + step_t = max(int(codes.shape[1]) for codes in slot_codes.values()) + if any(int(codes.shape[1]) != step_t and slot not in terminal_slots for slot, codes in slot_codes.items()): + raise ValueError("Only terminal codec rows may be padded to a larger step length") slots = list(slot_codes) if any(slot not in self._leased_slots for slot in slots): raise RuntimeError(f"Streaming decode references an unleased state slot: {slots}") codes_step = torch.stack( - [slot_codes[slot].to(device=self._device, dtype=torch.long) for slot in slots], + [ + torch.nn.functional.pad( + slot_codes[slot].to(device=self._device, dtype=torch.long), + (0, step_t - int(slot_codes[slot].shape[1])), + ) + if int(slot_codes[slot].shape[1]) != step_t + else slot_codes[slot].to(device=self._device, dtype=torch.long) + for slot in slots + ], dim=1, ) state_slot_ids = torch.tensor(slots, device=self._device, dtype=torch.long) @@ -197,9 +204,9 @@ def step( self._reset_slot_ids(terminal_slot_ids) audio = audio_tensor.detach().to("cpu", torch.float32) - audio_length = step_t * self._samples_per_frame out: dict[int, torch.Tensor] = {} for row, slot in enumerate(slots): + audio_length = int(slot_codes[slot].shape[1]) * self._samples_per_frame out[slot] = audio[row, ..., :audio_length].contiguous() return out @@ -518,6 +525,8 @@ def _decode_streaming_batch( outputs: 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 + coalesce_tails = graph is not None and max(graph.batch_sizes, default=0) >= self._stream_state_capacity for output_index, request_id, codes_nq_t, finished in items: slot = self._stream_req_slots.get(request_id) @@ -548,9 +557,13 @@ def _decode_streaming_batch( ) continue - grouped.setdefault(int(codes_nq_t.shape[1]), []).append( - (output_index, request_id, slot, codes_nq_t, finished) - ) + 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: + 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)) for group in grouped.values(): plan = {slot: codes_nq_t for _, _, slot, codes_nq_t, _ in group} @@ -756,6 +769,17 @@ def _remap(name: str) -> str: codec.eval() if device.type != "cpu": 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": + 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 + + 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") 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 e2d46ce2cbb..ba19c02d186 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 @@ -1,3 +1,5 @@ +# 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"). @@ -101,8 +103,13 @@ def _rotate_half(x: torch.Tensor) -> torch.Tensor: odd = x[..., 1::2] return torch.stack((-odd, even), dim=-1).reshape_as(x) - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - """``hidden_states``: ``(B, S, H)``. Re-prefills with a fresh causal mask over ``[0, S)`` every call.""" + def forward( + self, + hidden_states: torch.Tensor, + kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None, + position: int = 0, + ) -> torch.Tensor: + """Run a causal prefix, or one new token using frame-local K/V.""" batch_size, seq_len, _ = hidden_states.shape qkv = self.c_attn(hidden_states) query, key, value = qkv.split(self.embed_dim, dim=-1) @@ -110,14 +117,22 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: key = key.view(batch_size, seq_len, self.n_head, self.head_dim) value = value.view(batch_size, seq_len, self.n_head, self.head_dim) - cos, sin = self._rope_cos_sin(seq_len, hidden_states.device, hidden_states.dtype) + cos, sin = self._rope_cos_sin(position + seq_len, hidden_states.device, hidden_states.dtype) + cos, sin = cos[:, position:], sin[:, position:] query = query * cos + self._rotate_half(query) * sin key = key * cos + self._rotate_half(key) * sin query = query.transpose(1, 2) key = key.transpose(1, 2) value = value.transpose(1, 2) - attn_output = F.scaled_dot_product_attention(query, key, value, is_causal=True) + if kv_cache is not None: + assert seq_len == 1, "Cached Local Depth execution consumes one token at a time" + k_cache, v_cache = kv_cache + k_cache[:, :, position : position + 1].copy_(key) + 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) attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, self.embed_dim) return self.c_proj(attn_output) @@ -141,8 +156,13 @@ def __init__(self, hidden_size: int, n_head: int, inner_size: int, rope_base: fl self.ln_2 = nn.LayerNorm(hidden_size, eps=eps) self.mlp = _MossTTSLocalMLP(hidden_size, inner_size) - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states = hidden_states + self.attn(self.ln_1(hidden_states)) + def forward( + self, + hidden_states: torch.Tensor, + kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None, + position: int = 0, + ) -> torch.Tensor: + hidden_states = hidden_states + self.attn(self.ln_1(hidden_states), kv_cache, position) hidden_states = hidden_states + self.mlp(self.ln_2(hidden_states)) return hidden_states @@ -156,11 +176,9 @@ class MossTTSLocalDepthTransformer(nn.Module): and codebook-0's head (``audio_lm_heads[0]``) simultaneously. - codebooks 1..n_vq-1 are sampled sequentially: each sampled code is re-embedded (``audio_embeddings[c]``) and appended as the next - position, re-prefilling the block over the growing (<=n_vq) sequence - with a fresh causal mask each call -- mathematically identical to - incremental KV-cache decoding since attention is strictly causal and - only the last position is ever read, but avoids any cache plumbing - given the trivially short sequence length. + position. K/V for earlier positions are reused within the frame, so + each position's projections and MLP run once. A new frame overwrites + every cache position before reading it, including CUDA Graph replay. """ def __init__(self, gpt2_config, hidden_size: int | None = None) -> None: @@ -174,11 +192,13 @@ def __init__(self, gpt2_config, hidden_size: int | None = None) -> None: self.ln_f = nn.LayerNorm(self.hidden_size, eps=eps) self._compiled_forward_prefix = None - def _forward_prefix(self, seq_embeds: torch.Tensor) -> torch.Tensor: - hidden_states = seq_embeds - for block in self.h: - hidden_states = block(hidden_states) - return self.ln_f(hidden_states) + def _forward_prefix( + self, + seq_embeds: torch.Tensor, + kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None, + position: int = 0, + ) -> torch.Tensor: + return self.ln_f(self.h[0](seq_embeds, kv_cache, position)) def setup_compile(self) -> None: if self._compiled_forward_prefix is not None: @@ -189,14 +209,18 @@ def setup_compile(self) -> None: return self._compiled_forward_prefix = torch.compile( self._forward_prefix, - dynamic=False, + # Share batch/position shapes instead of exhausting Dynamo's + # recompilation budget while warming request-size graph buckets. + dynamic=True, options={"epilogue_fusion": False}, ) - logger.info("MOSS-TTS local depth prefix enabled with torch.compile") + logger.info("MOSS-TTS local depth frame-local KV execution enabled with torch.compile") - def _run_prefix(self, seq_embeds: torch.Tensor) -> torch.Tensor: + def _run_prefix( + self, seq_embeds: torch.Tensor, kv_cache: tuple[torch.Tensor, torch.Tensor], position: int + ) -> torch.Tensor: forward_prefix = self._compiled_forward_prefix or self._forward_prefix - return forward_prefix(seq_embeds) + return forward_prefix(seq_embeds, kv_cache, position) @torch.no_grad() def generate_frame( @@ -240,9 +264,13 @@ def generate_frame( for block in self.h: block.attn.prepare_rope_cache(n_vq, backbone_last_hidden.device, dtype) - embeds = backbone_last_hidden.new_zeros((batch_size, n_vq, self.hidden_size), dtype=dtype) - embeds[:, 0, :] = backbone_last_hidden.to(dtype) - hidden = self._run_prefix(embeds[:, :1, :]) + attn = self.h[0].attn + cache_shape = (batch_size, attn.n_head, n_vq, attn.head_dim) + kv_cache = ( + backbone_last_hidden.new_empty(cache_shape, dtype=dtype), + backbone_last_hidden.new_empty(cache_shape, dtype=dtype), + ) + hidden = self._run_prefix(backbone_last_hidden[:, None, :].to(dtype), kv_cache, 0) local_hidden = hidden[:, 0, :] binary_logits = local_text_lm_head(local_hidden).float() @@ -292,9 +320,9 @@ def generate_frame( codes[:, channel_index] = channel_token if channel_index + 1 < n_vq: - embeds[:, channel_index + 1, :] = audio_embeddings[channel_index](channel_token).to(dtype) - hidden = self._run_prefix(embeds[:, : channel_index + 2, :]) - local_hidden = hidden[:, channel_index + 1, :] + embeds = audio_embeddings[channel_index](channel_token).to(dtype)[:, None, :] + hidden = self._run_prefix(embeds, kv_cache, channel_index + 1) + local_hidden = hidden[:, 0, :] return should_continue, codes 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 0495d79d80e..9d2d2d17139 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 @@ -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"). @@ -1671,20 +1674,19 @@ def make_omni_output( self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts] self._batch_state_spans = kwargs.get("request_token_spans") - should_values: list[torch.Tensor] = [] - have_mtp_output = False + should_codes: list[torch.Tensor] = [] + should_indices: list[int] = [] per_req_codes: list[torch.Tensor] = [] have_codes = False - for info in info_dicts: + for index, info in enumerate(info_dicts): audio_codes = (info.get("audio_codes", {}) or {}) if isinstance(info, dict) else {} current = audio_codes.pop("current", None) if isinstance(current, torch.Tensor) and current.numel() > 0: codes = current.to(device=hidden.device, dtype=torch.long) if codes.dim() == 1: codes = codes.unsqueeze(0) - should_continue = codes.ne(self.audio_pad_token_id).any(dim=-1) - should_values.append(should_continue.reshape(-1)[:1]) - have_mtp_output = True + should_codes.append(codes[:1]) + should_indices.append(index) # Do not boolean-index a CUDA tensor here. Even for the # common single-row case, ``codes[should_continue]`` must # discover a dynamic output size on the host and therefore @@ -1694,9 +1696,19 @@ def make_omni_output( per_req_codes.append(codes) have_codes = True else: - should_values.append(torch.ones((1,), device=hidden.device, dtype=torch.bool)) per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long)) - self._batch_should_continue = torch.cat(should_values, dim=0) if have_mtp_output and should_values else None + self._batch_should_continue = None + if should_codes: + active = torch.cat(should_codes, dim=0).ne(self.audio_pad_token_id).any(dim=-1) + if len(should_codes) == len(info_dicts): + self._batch_should_continue = active + else: + self._batch_should_continue = torch.ones(len(info_dicts), device=hidden.device, dtype=torch.bool) + # Keep mixed prefill/decode output packing asynchronous. + indices = torch.tensor(should_indices, device="cpu", dtype=torch.long).to( + hidden.device, non_blocking=True + ) + self._batch_should_continue.index_copy_(0, indices, active) if not have_codes: return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) diff --git a/vllm_omni/model_executor/models/moss_tts/pipeline.py b/vllm_omni/model_executor/models/moss_tts/pipeline.py index f499654a463..0d6151313ff 100644 --- a/vllm_omni/model_executor/models/moss_tts/pipeline.py +++ b/vllm_omni/model_executor/models/moss_tts/pipeline.py @@ -1,8 +1,13 @@ +# 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"). """Pipeline topology for all MOSS-TTS variants (2-stage: talker → codec).""" +from vllm.sampling_params import RequestOutputKind + from vllm_omni.config.stage_config import ( PipelineConfig, StageExecutionType, @@ -100,6 +105,10 @@ sampling_constraints={ "detokenize": False, "stop_token_ids": [151645], + # The worker connector streams codes directly to the codec. + # This internal stage only needs to publish its terminal + # result; codec audio output remains incremental. + "output_kind": RequestOutputKind.FINAL_ONLY, }, ), StagePipelineConfig( diff --git a/vllm_omni/model_executor/models/moss_tts/streaming_attention.py b/vllm_omni/model_executor/models/moss_tts/streaming_attention.py new file mode 100644 index 00000000000..1acce771045 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/streaming_attention.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Masked streaming attention preserving ring-cache position semantics.""" + +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def _attention( + q_ptr, + k_ptr, + v_ptr, + mask_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, + ms0: tl.constexpr, + ms2: tl.constexpr, + ms3: tl.constexpr, + num_heads: tl.constexpr, + q_len: tl.constexpr, + kv_len: tl.constexpr, + head_dim: tl.constexpr, + block_m: tl.constexpr, + block_n: 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 + dims = tl.arange(0, head_dim) + q = tl.load( + q_ptr + batch * qs0 + head * qs1 + rows[:, None] * qs2 + dims[None, :] * qs3, + mask=rows[:, None] < q_len, + other=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) + for first in range(tl.cdiv(kv_len, block_n)): + cols = first * block_n + tl.arange(0, block_n) + k = tl.load( + k_ptr + batch * ks0 + head * ks1 + cols[None, :] * ks2 + dims[:, None] * ks3, + mask=cols[None, :] < kv_len, + other=0, + ) + allowed = tl.load( + mask_ptr + batch * ms0 + rows[:, None] * ms2 + cols[None, :] * ms3, + mask=(rows[:, None] < q_len) & (cols[None, :] < kv_len), + other=0, + ) + score = tl.dot(q, k).to(tl.float32) * (head_dim**-0.5) + score = tl.where(allowed & (cols[None, :] < kv_len), score, -float("inf")) + next_max = tl.maximum(maximum, tl.max(score, axis=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, axis=1) + v = tl.load( + v_ptr + batch * vs0 + head * vs1 + cols[:, None] * vs2 + dims[None, :] * vs3, + mask=cols[:, None] < kv_len, + other=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) * q_len + rows[:, None]) * head_dim + dims[None, :], + out, + mask=rows[:, None] < q_len, + ) + + +@torch.library.custom_op("moss_streaming::masked_attention", mutates_args=()) +def masked_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + b, h, t, d = q.shape + out = torch.empty(q.shape, device=q.device, dtype=q.dtype) + # Small query tiles expose parallelism for short streaming chunks. Larger + # tiles reuse K/V for steady chunks. Keep 64-key softmax blocks to retain + # the same accumulation granularity across both paths. + bm = 16 if t <= 32 else 64 + num_warps = 8 if k.shape[2] <= 256 else 4 + _attention[(triton.cdiv(t, bm), b * h)]( + q, + k, + v, + mask, + out, + *q.stride(), + *k.stride(), + *v.stride(), + mask.stride(0), + mask.stride(2), + mask.stride(3), + h, + t, + k.shape[2], + d, + bm, + 64, + num_warps=num_warps, + num_stages=2, + ) + return out + + +@masked_attention.register_fake +def _(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + return torch.empty(q.shape, device=q.device, dtype=q.dtype) diff --git a/vllm_omni/outputs/__init__.py b/vllm_omni/outputs/__init__.py index 90197f6e663..e74931a2011 100644 --- a/vllm_omni/outputs/__init__.py +++ b/vllm_omni/outputs/__init__.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from dataclasses import dataclass, field from typing import Any @@ -52,7 +55,7 @@ class OmniModelRunnerOutput(ModelRunnerOutput): """ multimodal_outputs: list[dict[str, object]] | None = None - inter_stage_outputs: list[dict[str, object]] | None = None + inter_stage_outputs: list[dict[str, Any] | None] | None = None # IDs of requests whose KV cache has been extracted from GPU/NPU to CPU. # The Scheduler can safely free the block tables for these requests. kv_extracted_req_ids: list[str] | None = None diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index d293f63122c..811f86f179a 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -15,7 +15,7 @@ from contextlib import nullcontext from copy import copy from dataclasses import replace -from typing import Any, NamedTuple +from typing import Any, NamedTuple, cast import numpy as np import torch @@ -87,14 +87,39 @@ def _clone_cuda_tensor_payload(value: Any, sources: list[torch.Tensor]) -> Any: return value.detach().clone() if isinstance(value, dict): return {k: _clone_cuda_tensor_payload(v, sources) for k, v in value.items()} - if isinstance(value, list): - return [_clone_cuda_tensor_payload(v, sources) for v in value] - if isinstance(value, tuple): - return tuple(_clone_cuda_tensor_payload(v, sources) for v in value) + if isinstance(value, (list, tuple)): + if ( + value + and all( + isinstance(v, torch.Tensor) + and v.is_cuda + and v.dtype == value[0].dtype + and v.device == value[0].device + and v.layout == torch.strided + and v.is_contiguous() + for v in value + ) + # Keep frame/state packing bounded; large payloads use individual snapshots. + and sum(v.numel() * v.element_size() for v in value) <= 1024 * 1024 + ): + packed = torch.cat([v.detach().reshape(-1) for v in value]) + sources.append(packed) + return _PackedTensorPayload(packed, [v.shape for v in value], isinstance(value, tuple)) + items = [_clone_cuda_tensor_payload(v, sources) for v in value] + return tuple(items) if isinstance(value, tuple) else items return value def _copy_tensor_payload_to_cpu(value: Any, pin_memory: bool) -> Any: + if isinstance(value, _PackedTensorPayload): + flat = _copy_tensor_payload_to_cpu(value.tensor, pin_memory) + items = [] + offset = 0 + for shape in value.shapes: + count = shape.numel() + items.append(flat[offset : offset + count].view(shape)) + offset += count + return tuple(items) if value.is_tuple else items if isinstance(value, torch.Tensor): if value.device.type != "cuda": return value @@ -110,6 +135,12 @@ def _copy_tensor_payload_to_cpu(value: Any, pin_memory: bool) -> Any: return value +class _PackedTensorPayload(NamedTuple): + tensor: torch.Tensor + shapes: list[torch.Size] + is_tuple: bool + + class _AsyncCPUPayloadSnapshot: def __init__( self, @@ -181,7 +212,7 @@ def __init__( if kwargs: raise TypeError(f"Unexpected OmniAsyncGPUModelRunnerOutput kwargs: {sorted(kwargs)}") - self._model_runner_output = None + self._model_runner_output: OmniModelRunnerOutput | None = None self._invalid_req_indices = invalid_req_indices self.async_copy_ready_event = torch.Event() @@ -211,7 +242,7 @@ def __init__( self._num_nans_cpu = self._num_nans.to("cpu", non_blocking=True) if self._num_nans is not None else None self.async_copy_ready_event.record() - self._model_runner_output_builder = model_runner_output_builder + self._model_runner_output_builder: Callable[[], OmniModelRunnerOutput] | None = model_runner_output_builder self._background_exception: BaseException | None = None self._background_thread: threading.Thread | None = None self._cuda_device = cuda_device @@ -226,6 +257,7 @@ def _build_model_runner_output_once(self) -> None: if self._model_runner_output is not None: return with record_function_or_nullcontext("omni_async_output:get_output/build_model_runner_output"): + assert self._model_runner_output_builder is not None self._model_runner_output = self._model_runner_output_builder() self._model_runner_output_builder = None @@ -321,6 +353,9 @@ class only overrides sample_tokens to expose hidden states + multimodal outputs per request while keeping Async output semantics. """ + execute_model_state: ExecuteModelState | None + kv_extracted_req_ids: list[str] | None + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32) @@ -436,6 +471,8 @@ def _request_final_stage_id(self, req_id: str) -> int | None: if not isinstance(info, dict): return None val = info.get("omni_final_stage_id") + if val is None: + return None try: return int(val) except (TypeError, ValueError): @@ -1485,7 +1522,7 @@ def _run_post_sample_talker_mtp( ) selected_indices.append(idx) selected_req_ids.append(req_id) - selected_req_infos.append(req_info) + selected_req_infos.append(cast(dict[str, Any], req_info)) if not selected_indices: return multimodal_outputs @@ -1552,7 +1589,7 @@ def _run_post_sample_talker_mtp( def _build_multimodal_outputs( self, - per_req_payloads: list[dict[str, object] | None] | None, + per_req_payloads: Sequence[dict[str, object] | None] | None, ) -> list[dict[str, torch.Tensor] | None] | None: """Build per-request multimodal output payloads (dedicated channel). @@ -1648,8 +1685,8 @@ def _should_defer_full_payload_d2h(self) -> bool: def _build_omni_step_outputs( self, - pooler_inter: list[dict[str, object]], - pooler_client: list[dict[str, object]], + pooler_inter: Sequence[dict[str, object] | None] | None, + pooler_client: Sequence[dict[str, object] | None] | None, *, defer_full_payload_d2h: bool, ) -> tuple[ @@ -1935,6 +1972,8 @@ def _build_omni_model_runner_output_from_snapshot( ) pooler_output.append(flatten_payload(payload)) + pooler_inter: Sequence[dict[str, object] | None] | None + pooler_client: Sequence[dict[str, object] | None] | None pooler_output = pooler_output or [] if self._async_chunk and stage_sends_async_output(self.model_config): pooler_inter, pooler_client = partition_payload_list(pooler_output) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index eca15c16e5c..80d55020f25 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -78,6 +78,8 @@ def _filter_mrope_kwargs_for_model(model: object, kwargs: dict[str, Any]) -> dic class OmniGPUModelRunner(GPUModelRunner): + intermediate_tensors: IntermediateTensors | None + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.model_intermediate_buffer: dict[str, dict[str, Any]] = {} @@ -842,7 +844,9 @@ def correct_spec_decode_token_counts(): return None @torch.inference_mode() - def extract_multimodal_outputs(self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput) -> dict: + def extract_multimodal_outputs( + self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput + ) -> tuple[Any, Any]: if ( hasattr(self.model, "have_multimodal_outputs") and self.model.have_multimodal_outputs @@ -1407,14 +1411,14 @@ def _build_model_kwargs_extra(self) -> dict: def _process_additional_information_updates( self, hidden_states: torch.Tensor, - multimodal_outputs: object, + multimodal_outputs: Any, num_scheduled_tokens_np: np.ndarray, scheduler_output: "SchedulerOutput", combined_hidden_states: dict[str, torch.Tensor] | None = None, - combined_multimodal_outputs: dict[str, object] | None = None, + combined_multimodal_outputs: dict[str, Any] | None = None, req_ids_filter: set[str] | None = None, req_ids: list[str] | None = None, - query_start_loc_cpu: object | None = None, + query_start_loc_cpu: Any = None, ) -> None: """Process model-provided per-request updates and merge into model_intermediate_buffer.""" req_ids = req_ids if req_ids is not None else self.input_batch.req_ids @@ -1476,9 +1480,9 @@ def _process_additional_information_updates( def _collect_additional_information_for_prefill( self, num_scheduled_tokens_np: np.ndarray, - ) -> dict[str, dict]: + ) -> None: """Overlay per-request prompt_embeds for the prefill portion and collect - additional_information slices for this step. Returns a map req_id -> dict.""" + additional_information slices for this step.""" for req_index, req_id in enumerate(self.input_batch.req_ids): req_state = self.requests[req_id] pe_cpu = getattr(req_state, "prompt_embeds_cpu", None) @@ -1707,9 +1711,9 @@ def _preprocess( # Overlay custom prompt_embeds per request for the prompt portion; # collect additional_information (tensor/list) for prefill portion only - decode_req_ids = [] - decode_start_offsets = [] - decode_batch_items = [] + decode_req_ids: list[str] = [] + decode_start_offsets: list[int] = [] + decode_batch_items: list[tuple[str, int, dict[str, Any]]] = [] batch_decode_preprocess = getattr(self.model, "preprocess_decode_batch", None) def flush_decode_batch() -> None: @@ -1717,6 +1721,7 @@ def flush_decode_batch() -> None: if not decode_batch_items: return + assert callable(batch_decode_preprocess) req_ids_b = [item[0] for item in decode_batch_items] start_offsets_b = [item[1] for item in decode_batch_items] req_infos_b = [item[2] for item in decode_batch_items] @@ -1761,7 +1766,7 @@ def flush_decode_batch() -> None: # mimo-audio check req_state = self.requests.get(req_id) - req_infos = self._maybe_attach_mimo_audio_req_infos(req_state, req_infos, req_id) + req_infos = cast(dict[str, Any], self._maybe_attach_mimo_audio_req_infos(req_state, req_infos, req_id)) start_offset = int(self.query_start_loc.cpu[req_index]) sched_tokens = int(num_scheduled_tokens_np[req_index]) @@ -1968,11 +1973,29 @@ def _row_generator(req_id: str) -> torch.Generator | None: if start_offsets is None: id_to_index = self.input_batch.req_id_to_index start_offsets = [int(self.query_start_loc.cpu[id_to_index[req_id]]) for req_id in decode_req_ids] - for idx, (req_id, start_offset) in enumerate(zip(decode_req_ids, start_offsets, strict=True)): - inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1] - if code_predictor_codes is not None: - update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}} - self._update_intermediate_buffer(req_id, update_dict) + if start_offsets == list(range(decode_batch_size)): + inputs_embeds[:decode_batch_size].copy_(req_embeds[:decode_batch_size]) + else: + # A device tensor constructor waits for the preceding MTP graph. + # Enqueue the small host index transfer without synchronizing it. + offsets = torch.tensor(start_offsets, device="cpu", dtype=torch.long).to( + inputs_embeds.device, non_blocking=True + ) + inputs_embeds.index_copy_(0, offsets, req_embeds[:decode_batch_size]) + if code_predictor_codes is not None: + if out_key in getattr(self.model, "gpu_resident_buffer_keys", set()): + # One owned batch snapshot protects all rows from graph reuse. + owned_codes = code_predictor_codes[:decode_batch_size].detach().clone() + for req_id, row in zip(decode_req_ids, owned_codes.split(1), strict=True): + req_state = self.requests.get(req_id) + if req_state is not None: + existing = self.model_intermediate_buffer.setdefault(req_id, {}) + existing.setdefault(out_key[0], {})[out_key[1]] = row + req_state.additional_information_cpu = existing + else: + for idx, req_id in enumerate(decode_req_ids): + update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}} + self._update_intermediate_buffer(req_id, update_dict) def _model_forward( self, From 51cf142bc1dfa372f2c61cfee44535f61afd93de Mon Sep 17 00:00:00 2001 From: Joshna-Medisetty Date: Wed, 9 Sep 2026 15:11:31 -0700 Subject: [PATCH 40/93] [Bugfix][XPU] Restore N-D output shape for W8A16 FP8 linear (#7301) Signed-off-by: Joshna Medisetty Signed-off-by: Joshna-Medisetty --- .../quantization/test_quantization_fp8.py | 3 +- vllm_omni/platforms/xpu/patch.py | 41 +++++++++++++++++++ vllm_omni/platforms/xpu/platform.py | 5 +++ 3 files changed, 48 insertions(+), 1 deletion(-) create mode 100644 vllm_omni/platforms/xpu/patch.py diff --git a/tests/diffusion/quantization/test_quantization_fp8.py b/tests/diffusion/quantization/test_quantization_fp8.py index 0da5b6dd34f..5793ad95cb6 100644 --- a/tests/diffusion/quantization/test_quantization_fp8.py +++ b/tests/diffusion/quantization/test_quantization_fp8.py @@ -300,7 +300,8 @@ def test_single_stage_zimage_fp8_uses_less_memory(): assert mem_fp8 < mem_bf16, f"FP8 ({mem_fp8:.2f} GiB) should use less memory than BF16 ({mem_bf16:.2f} GiB)" -@hardware_test(res={"cuda": "L4"}) +@hardware_test(res={"cuda": "L4", "xpu": "B60"}) +@pytest.mark.advanced_model def test_single_stage_qwen_image_fp8(): """Qwen-Image (random weights) with FP8 generates valid images.""" model = "riverclouds/qwen_image_random" diff --git a/vllm_omni/platforms/xpu/patch.py b/vllm_omni/platforms/xpu/patch.py new file mode 100644 index 00000000000..841cf25038d --- /dev/null +++ b/vllm_omni/platforms/xpu/patch.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""XPU-specific patches for upstream vLLM behavior.""" + +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +def _patch_w8a16_fp8_nd_activation() -> None: + """Restore the N-D output shape for the XPU W8A16 FP8 linear kernel. + + ``XPUW8A16FP8LinearKernel.apply_weights`` hands its activation straight to + ``torch.ops._xpu_C.fp8_gemm_w8a16``, skipping the flatten-then-reshape every + other ScaledMM kernel does. The op's registered fake always returns 2-D, so a + 3-D activation makes inductor bake in a 2-D ``assert_size_stride`` that the + rank-3 result trips, killing FLUX.1 ``--quantization fp8`` at startup warmup. + Only >2-D inputs are affected, so LLM (token-flattened, 2-D) paths are not. + """ + try: + from vllm.model_executor.kernels.linear.scaled_mm.xpu import ( + XPUW8A16FP8LinearKernel, + ) + except ImportError: + logger.debug("XPUW8A16FP8LinearKernel not available; skipping FP8 shape patch.") + return + + _original_apply_weights = XPUW8A16FP8LinearKernel.apply_weights + + def _patched_apply_weights(self, layer, x, bias=None): + output_shape = (*x.shape[:-1], layer.weight.shape[1]) + x = x.reshape(-1, x.shape[-1]) + return _original_apply_weights(self, layer, x, bias).view(*output_shape) + + XPUW8A16FP8LinearKernel.apply_weights = _patched_apply_weights + + +def apply_patches() -> None: + """Apply all XPU-specific patches.""" + _patch_w8a16_fp8_nd_activation() diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py index 3d2cd2d6221..df4e7fd55c5 100644 --- a/vllm_omni/platforms/xpu/platform.py +++ b/vllm_omni/platforms/xpu/platform.py @@ -11,6 +11,7 @@ from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum +from vllm_omni.platforms.xpu.patch import apply_patches logger = init_logger(__name__) @@ -31,6 +32,10 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform): _omni_enum = OmniPlatformEnum.XPU + def __init__(self): + super().__init__() + apply_patches() + @classmethod def get_omni_ar_worker_cls(cls) -> str: return "vllm_omni.platforms.xpu.worker.xpu_ar_worker.XPUARWorker" From faad918588e616ed073030710e653d4b1d3f3e79 Mon Sep 17 00:00:00 2001 From: Guangjian Dong <163994576+Hiro208@users.noreply.github.com> Date: Thu, 10 Sep 2026 08:00:13 +0800 Subject: [PATCH 41/93] [Doc] Document num_outputs_per_prompt for /v1/videos (#7341) Signed-off-by: Guangjian --- docs/user_guide/examples/online_serving/text_to_video.md | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/docs/user_guide/examples/online_serving/text_to_video.md b/docs/user_guide/examples/online_serving/text_to_video.md index 2f00d386b8a..894e2bf0eed 100644 --- a/docs/user_guide/examples/online_serving/text_to_video.md +++ b/docs/user_guide/examples/online_serving/text_to_video.md @@ -199,12 +199,21 @@ curl -X POST http://localhost:8091/v1/videos \ | `boundary_ratio` | float | None | Boundary split ratio for low/high DiT (Wan2.2) | | `flow_shift` | float | None | Scheduler flow shift (Wan2.2) | | `seed` | int | None | Random seed (reproducible) | +| `num_outputs_per_prompt` | int | 1 | Number of videos to generate (1-10, MiniMax H3); only the first is returned, see below | | `lora` | object | None | LoRA configuration | | `enable_frame_interpolation` | bool | false | Enable RIFE frame interpolation before MP4 encoding | | `frame_interpolation_exp` | int | 1 | Interpolation exponent; 1=2x temporal resolution, 2=4x | | `frame_interpolation_scale` | float | 1.0 | RIFE inference scale; use 0.5 for high-resolution inputs | | `frame_interpolation_model_path` | str | None | Local directory or Hugging Face repo ID with `flownet.pkl`; defaults to `elfgum/RIFE-4.22.lite` | +!!! note "Only the first output is returned" + `num_outputs_per_prompt` is forwarded to the pipeline, so the model does + generate that many videos and you pay the generation cost for all of them. + `/v1/videos` and `/v1/videos/sync` then return only the **first** one, + because an async video job stores a single file per video id. The server + logs a warning when it discards the extra outputs. Keep this at `1` + unless you have a specific reason to generate videos you will not receive. + ## Frame Interpolation Frame interpolation is an optional post-processing step for `/v1/videos` and From 06374543012ea0e10b970081a81097762acff78e Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Thu, 10 Sep 2026 08:05:47 +0800 Subject: [PATCH 42/93] [Skills] Add perf-evidence isolation, stage-attribution, and realtime-contract requirements (#6820) Signed-off-by: hsliu_ustc Co-authored-by: hsliu_ustc --- .claude/skills/review-pr/SKILL.md | 11 +- .../references/checks/perf-verification.md | 107 ++++++++++++++++++ .../references/process/general-checks.md | 13 +++ .../references/process/review-execution.md | 14 +-- 4 files changed, 133 insertions(+), 12 deletions(-) diff --git a/.claude/skills/review-pr/SKILL.md b/.claude/skills/review-pr/SKILL.md index 0b19b2b379c..5907e13f312 100644 --- a/.claude/skills/review-pr/SKILL.md +++ b/.claude/skills/review-pr/SKILL.md @@ -123,10 +123,11 @@ it changes again, report the churn and wait for a stable target. For a trusted PR head, materialize the pinned head in an isolated detached worktree. A worktree freezes identity but is not a security sandbox. Treat fork -heads as untrusted unless the user and environment policy explicitly establish -otherwise: execute them only in a disposable, secret-free sandbox with restricted -filesystem, network, and resources; without one, use static SHA-addressed reads -and CI evidence only. For a local review, freeze the committed, index, worktree, +heads as untrusted until the user and environment policy explicitly establish +trust: execute their code only after that trust is recorded for this host, +with the trusted SHA noted in the review state and secrets kept out of the +execution scope; without recorded trust, use static SHA-addressed reads and CI +evidence only. For a local review, freeze the committed, index, worktree, and NUL-safe in-scope untracked contents. Follow [review-execution.md](references/process/review-execution.md) for trust gates, state fingerprints, and byte-for-byte staleness checks. @@ -196,7 +197,7 @@ style preferences as simplification findings. Before each validation group, verify the frozen SHA plus the tracked, index, untracked, and ignored-file fingerprint, or recreate a pristine snapshot. On a -trusted head or inside the required sandbox, run an import/version preflight, +head the user explicitly trusted for this host, run an import/version preflight, then the narrowest relevant tests and low-cost static checks. Bind every result to the head SHA, snapshot fingerprint, and environment fingerprint. Never run imports, tests, builds, hooks, or repo-configurable tooling from an untrusted diff --git a/.claude/skills/review-pr/references/checks/perf-verification.md b/.claude/skills/review-pr/references/checks/perf-verification.md index bbfe5a1471d..027e91589b2 100644 --- a/.claude/skills/review-pr/references/checks/perf-verification.md +++ b/.claude/skills/review-pr/references/checks/perf-verification.md @@ -22,6 +22,19 @@ Keep exact commands and report variability rather than a cherry-picked run. Use the repository's benchmark when it covers the claim; otherwise use the smallest reproducible workload that reaches the changed production path. +## Isolation per claim + +Isolate one claim per comparison. When a PR stacks several changes (a fused or +restructured load path, a distilled step schedule, a new cache), request one +column per claim at equal steps/workload instead of one blended base-vs-head +delta, and name the unchanged stage or metric that pins equivalence. + +When the changed path is runtime-switchable (feature flag or environment +variable), run every column on the same head and toggle the switch rather than +diffing against a base commit that cannot separate the stacked changes: +base, head with the switch off, head with the switch on, then head with the +switch on plus each further claim. + ## Evidence table | Dimension | Typical evidence | @@ -34,6 +47,100 @@ smallest reproducible workload that reaches the changed production path. Every optimization needs correctness or quality evidence. A faster result on a different workload, precision, seed, or topology is not a valid comparison. +## E2e stage-attribution table + +For an end-to-end latency claim, request the stage-attribution table the +repository's perf PRs use: + +- Rows carry the benchmark's own field names: total e2e, per-stage execution, + postprocess, response encoding, and a request-to-output roll-up. Add a stage + handoff/idle row when the pipeline spans stages — connector transfer and + scheduling gaps are invisible in per-stage latency deltas. +- Sanity rows pin comparability: inference steps, output frames/pixels/samples, + and generated tokens held identical across columns. +- Report P50 and P100 from the stated warmup and measured counts. One + population per number: never mix cold-start and warm samples into one value. +- Name the baseline and candidate commit SHAs, and close with an isolation + statement naming the metric that stayed unchanged ("X remained effectively + unchanged, isolating the gain to Y"). + +When the serving path caches, captures, or regenerates at runtime (CUDA-graph +bucket growth, cache rebuilds, recompilation), medians hide the resulting tail. +Also request the event count and the maximum observed stall. + +Component-level tables pair with, not replace, the e2e table: per-op reference +vs candidate timings (with the numerical-equality result per op) plus one +complete-pipeline check under the production shapes. + +## Proportionality + +When a PR adds a model and explicitly labels its numbers as smoke observations +rather than claims, do not demand the attribution machinery. Ask only for +evidence protecting the model's stated operating contract — for a +realtime/duplex model that is sustained cadence over a realistic session +(enough ticks to leave warmup), with any mid-session capture/regeneration +event and its stall reported, since a cadence that collapses mid-session is a +functional failure, not a perf claim. + +## Realtime/duplex world models + +Request the operating-contract evidence as serving metrics, not claims: + +- Sustained per-tick latency over one realistic session: median, max, and the + argmax tick. +- Stall events with their causes. +- Peak memory at first versus last tick — session-owned state must not grow + per tick. +- Control latency from a submitted action to the tick that applies it. +- Realtime-versus-offline parity under the same controls, using the + repository's video-similarity helpers. Parity plotted over tick index is the + drift-over-horizon curve and the honest measure of session lifetime; + distribution-level suites (FVD, VBench) belong to model evaluation, not the + serving integration. + +Watch the contract boundaries the model documents: if a control input is +tick-API-only, the parity pair must use a horizon both paths can run. + +## Speech-generation (TTS) models + +The streaming contract outranks the average — a synthesis that pauses +mid-stream is broken for realtime use regardless of its averages. Request: + +- The benchmark's own speech metrics over one long utterance (enough audio to + leave warmup, tens of seconds): `audio_ttfp` (time to first packet), + `audio_rtf` (real-time factor; above 1.0 cannot keep up), `e2el`, and + `ttft`/`tpot` where tokens precede audio. +- Mid-utterance stalls as the benchmark's continuity fields, not prose: + `audio_underrun` percentiles and the `audio_continuity_ok` rate + (`vllm_omni/benchmarks/metrics/metrics.py`), computed by + `vllm_omni/benchmarks/audio_continuity.py` — the companion to `audio_ttfp` + and `audio_rtf`. A synthesis that bursts chunks fast enough on average but + underruns mid-stream fails the streaming contract even when the averages + pass. +- Quality as a paired comparison against the reference implementation under + the same text, speaker prompt, and seed; report sample rate and frame + count, and assert output completeness — truncated or dropped trailing audio + is a known failure class, not a quality nuance. +- For realtime or duplex TTS: interruption (barge-in) latency and the + per-session memory trend across turns. + +## Omni models + +Apply the e2e stage-attribution table directly to the voice-to-voice path of +omni (speech-to-speech and multimodal AR) models: + +- TTFT for text output, `audio_ttfp` for audio output, and per-stage + ITL/TPOT cadence across the perception-understanding-generation chain. +- Cross-modal parity — the omni analogue of realtime-versus-offline parity: + the same semantic request through text and audio inputs, compared on output + content or timing. +- For duplex omni sessions: turn-taking latency, interruption handling, the + concurrent session count actually validated, and a per-turn memory trend — + session-owned KV state must not grow per turn. When either turn produces + audio, continuity is judged with the same `audio_underrun` and + `audio_continuity_ok` fields as TTS, since turn-taking latency alone does + not expose mid-utterance underruns. + ## Reviewer verification ladder 1. Run base/head A/B on suitable hardware when affordable. diff --git a/.claude/skills/review-pr/references/process/general-checks.md b/.claude/skills/review-pr/references/process/general-checks.md index 88dfd405e13..da86eabc710 100644 --- a/.claude/skills/review-pr/references/process/general-checks.md +++ b/.claude/skills/review-pr/references/process/general-checks.md @@ -121,6 +121,19 @@ Search bounded callers and sibling implementations before reporting dead code, duplication, or a missed consumer. Keep new abstractions only when they have a distinct live caller, invariant, or compatibility need. +Verify absence and equality claims with whole-file evidence. A line window, a +grep for one symbol, or a commit message proves presence, never absence: to +assert that a file matches another version or lacks a change, diff or hash the +full files (for example `git diff : :` or the contents +API's size and SHA fields) before anchoring the finding. A verification +shortcut that would itself be wrong to run is not evidence. + +For a value written or finalized in more than one place, establish which site +executes first on each entry path — and whether the earlier site makes the +later ones unreachable — before naming where a fix belongs. Enumerating all +call sites correctly still misses the bug when a path the enumeration ranked +last actually finalizes first and clears the state the proposed fix checks. + ## Finding bar Anchor each finding to a changed `path:line`; name the trigger or call path, diff --git a/.claude/skills/review-pr/references/process/review-execution.md b/.claude/skills/review-pr/references/process/review-execution.md index 7c0305254f8..c78532e9d5d 100644 --- a/.claude/skills/review-pr/references/process/review-execution.md +++ b/.claude/skills/review-pr/references/process/review-execution.md @@ -39,13 +39,13 @@ isolated detached worktree. Run every source read, `rg` search, import, and test from that snapshot, not the caller's checkout. A detached worktree provides snapshot isolation, not a security boundary. -Treat a fork head as untrusted unless the user and environment policy explicitly -establish trust. Never run its imports, tests, builds, hooks, package setup, or -repo-configurable linters/plugins on the reviewer host. Execute them only in a -disposable sandbox or VM with no credentials or inherited secrets, no host agent -or service sockets, minimal read-only host mounts, disabled or explicitly -allowlisted network, resource/time limits, and destruction after the run. If -that boundary is unavailable, limit the review to the remote diff, SHA-addressed +Treat a fork head as untrusted until the user and environment policy +explicitly establish trust. Without recorded trust, never run its imports, +tests, builds, hooks, package setup, or repo-configurable linters/plugins on +the reviewer host. Once the user explicitly trusts this pinned commit for +this host, record the trusted SHA with the review state and execute with +credentials, agent sockets, and other secrets out of scope. Without recorded +trust, limit the review to the remote diff, SHA-addressed reads such as `git show :`, and existing CI evidence; report all executable validation as a gap. From 6e63ea875e9a54f640b52478a567b9393df95892 Mon Sep 17 00:00:00 2001 From: Gao Han Date: Thu, 10 Sep 2026 09:54:20 +0800 Subject: [PATCH 43/93] [Bugfix] Allow LLM replicas on different GPUs to initialize concurrently (#7292) Signed-off-by: Gao Han Co-authored-by: amy-why-3459 --- .../test_stage_runtime_launch_concurrency.py | 149 ++++++++++++++++++ vllm_omni/engine/stage_runtime.py | 61 +++---- 2 files changed, 169 insertions(+), 41 deletions(-) create mode 100644 tests/engine/test_stage_runtime_launch_concurrency.py diff --git a/tests/engine/test_stage_runtime_launch_concurrency.py b/tests/engine/test_stage_runtime_launch_concurrency.py new file mode 100644 index 00000000000..f880b9386b9 --- /dev/null +++ b/tests/engine/test_stage_runtime_launch_concurrency.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Check replica startup ordering without model-loading timing assumptions.""" + +import contextlib +import fcntl +import os +import threading +from collections.abc import Iterator +from pathlib import Path + +import pytest +from pytest_mock import MockerFixture +from vllm.config import VllmConfig +from vllm.sampling_params import SamplingParams +from vllm.v1.engine.utils import CoreEngineProcManager, EngineZmqAddresses + +from vllm_omni.config.stage_config import StageConfig +from vllm_omni.engine import stage_init_utils, stage_runtime +from vllm_omni.engine.stage_engine_startup import StageReplicaResources +from vllm_omni.engine.stage_init_utils import LogicalStageInitPlan, ReplicaInitPlan, StageMetadata + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def launch_runtime( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, mocker: MockerFixture +) -> Iterator[stage_runtime.StageRuntime]: + runtime = stage_runtime.StageRuntime( + stage_configs=[], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=5, + async_chunk=False, + ) + # These are virtual device IDs: only the real per-device file locking and + # scheduling run. No model, CUDA context, or engine process is created. + monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda _stage, cfg: cfg["devices"]) + monkeypatch.setattr( + stage_init_utils, "device_init_lock_path", lambda device: str(tmp_path / f"device-{device}.lock") + ) + mocker.patch.object( + stage_runtime.StageEngineCoreClientBase, + "make_async_mp_client", + return_value=mocker.Mock(spec=stage_runtime.StageClient), + ) + monkeypatch.setenv("VLLM_OMNI_TEST_LAUNCH_ENV", "parent") + yield runtime + if runtime._stage_init_executor is not None: + runtime._stage_init_executor.shutdown(wait=True) + + +def _plan(stage_id: int, device: str) -> LogicalStageInitPlan: + metadata = StageMetadata( + stage_id=stage_id, + stage_type="llm", + engine_output_type="text", + is_comprehension=False, + requires_multimodal_data=False, + engine_input_source=[], + final_output=False, + final_output_type=None, + default_sampling_params=SamplingParams(), + custom_process_input_func=None, + model_stage="dummy", + runtime_cfg={"devices": device, "env": {"VLLM_OMNI_TEST_LAUNCH_ENV": str(stage_id)}}, + ) + replica = ReplicaInitPlan( + replica_id=0, + num_replicas=1, + launch_mode="local", + stage_cfg=StageConfig(stage_id=stage_id, model_stage="dummy", yaml_runtime=metadata.runtime_cfg), + metadata=metadata, + stage_connector_spec={}, + omni_kv_connector=(None, None, None), + stage_vllm_config=VllmConfig(), + executor_class=object, + engine_args_dict={}, + ) + return LogicalStageInitPlan(stage_idx=stage_id, stage_id=stage_id, replicas=[replica]) + + +def _resources(mocker: MockerFixture) -> StageReplicaResources: + return StageReplicaResources( + manager=mocker.Mock(spec=CoreEngineProcManager), + addresses=EngineZmqAddresses(inputs=["unused-input"], outputs=["unused-output"]), + ) + + +def test_distinct_devices_spawn_before_either_is_ready( + launch_runtime: stage_runtime.StageRuntime, monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture +) -> None: + spawned = [threading.Event(), threading.Event()] + + @contextlib.contextmanager + def launch(*, stage_id: int, stage_visible_devices: str, **_kwargs: object) -> Iterator[StageReplicaResources]: + assert os.environ["VLLM_OMNI_TEST_LAUNCH_ENV"] == str(stage_id) + spawned[stage_id].set() + yield _resources(mocker) + # The device remains exclusively locked throughout loading/profiling, + # even though another device's replica must be able to start. + with open(stage_init_utils.device_init_lock_path(int(stage_visible_devices)), "r+") as lock: + with pytest.raises(BlockingIOError): + fcntl.flock(lock.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + assert spawned[1 - stage_id].wait(5), "A different device's spawn was blocked by the READY wait" + + monkeypatch.setattr(stage_runtime, "launch_stage_replica", launch) + clients = launch_runtime._initialize_stage_replicas([_plan(0, "0"), _plan(1, "1")], 5) + assert all(clients[stage][0] is not None for stage in (0, 1)) + assert os.environ["VLLM_OMNI_TEST_LAUNCH_ENV"] == "parent" + + +def test_same_device_waits_for_ready_before_next_spawn( + launch_runtime: stage_runtime.StageRuntime, monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture +) -> None: + events: list[tuple[int, str]] = [] + + @contextlib.contextmanager + def launch(*, stage_id: int, **_kwargs: object) -> Iterator[StageReplicaResources]: + events.append((stage_id, "spawn")) + yield _resources(mocker) + events.append((stage_id, "ready")) + + monkeypatch.setattr(stage_runtime, "launch_stage_replica", launch) + launch_runtime._initialize_stage_replicas([_plan(0, "0"), _plan(1, "0")], 5) + assert events == [(0, "spawn"), (0, "ready"), (1, "spawn"), (1, "ready")] + assert os.environ["VLLM_OMNI_TEST_LAUNCH_ENV"] == "parent" + + +def test_ready_failure_cleans_up_resources_and_releases_device( + launch_runtime: stage_runtime.StageRuntime, monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture +) -> None: + manager = mocker.Mock(spec=CoreEngineProcManager) + resources = StageReplicaResources(manager=manager, addresses=EngineZmqAddresses(inputs=[], outputs=[])) + + @contextlib.contextmanager + def launch(**_kwargs: object) -> Iterator[StageReplicaResources]: + yield resources + raise RuntimeError("READY handshake failed") + + monkeypatch.setattr(stage_runtime, "launch_stage_replica", launch) + with pytest.raises(RuntimeError, match="READY handshake failed"): + launch_runtime._initialize_local_llm_replica(_plan(0, "0").replicas[0], 5) + + manager.shutdown.assert_called_once_with() + assert os.environ["VLLM_OMNI_TEST_LAUNCH_ENV"] == "parent" + with open(stage_init_utils.device_init_lock_path(0), "r+") as lock: + fcntl.flock(lock.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index 5d7d07653af..53400df5e4b 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -11,7 +11,7 @@ import threading import time from collections.abc import Callable, Iterator, Mapping, Sequence -from contextlib import contextmanager +from contextlib import ExitStack, contextmanager from dataclasses import dataclass from typing import Any, cast @@ -149,12 +149,9 @@ def __init__( self.stage_pools: list[StagePool] = [] self._stage_init_executor: concurrent.futures.ThreadPoolExecutor | None = None self._spawn_device_lock = threading.Lock() - # Serialize all LLM replica spawning + handshake across device groups - # to prevent ZMQ port-allocation races (get_engine_zmq_addresses) and - # CUDA-context conflicts when multiple engine core subprocesses - # initialize simultaneously on different GPUs. Matches the old - # AsyncOmniEngine._initialize_llm_replica pattern which used a single - # ``llm_stage_launch_lock`` for all replicas. + # Serialize process spawning and process-global environment overlays. + # Readiness waits use per-device initialization protection so different + # device groups can initialize concurrently. self._replica_launch_lock = threading.Lock() self._init_visible_devices_baseline: str | None = None @@ -778,43 +775,25 @@ def _initialize_local_llm_replica( spawn_device_lock=self._spawn_device_lock, omni_parallel_stage_init=self._parallel_stage_init, ) - # G2 launch lock serializes engine-core *spawning* across replicas - # (ZMQ port-allocation races + simultaneous CUDA context init). - # * Default path: hold it across the whole launch context manager, - # whose __exit__ waits for READY — this serializes the full child - # init (spawn + load + profile + KV + capture). - # * Parallel path: hold it only around the spawn (__enter__); run - # the READY-wait (__exit__) outside the lock so replicas init - # concurrently, coordinated by the child SH/EX device locks. - if self._parallel_stage_init: - # The per-stage runtime.env overlay mutates os.environ - # (process-global), so it must be applied under the launch lock - # and only needs to cover the spawn (__enter__) — children - # inherit the env at spawn time; the READY-wait needs no env. - g2_start = time.perf_counter() + # Only process spawning and the runtime.env overlay need the global + # launch lock. The launch context's exit waits for READY: holding the + # lock there serializes model loading and compilation even across + # different GPUs. Default initialization keeps its per-device EX + # locks until READY; parallel_stage_init uses child phase locks. + g2_start = time.perf_counter() + with ExitStack() as launch_stack: with self._replica_launch_lock: + g2_locked = time.perf_counter() with stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg): - resources = launch_cm.__enter__() + resources = launch_stack.enter_context(launch_cm) g2_spawned = time.perf_counter() - launch_cm.__exit__(None, None, None) - logger.debug( - "[stage_init] Stage-%s G2 spawn(locked)=%.3fs, READY(unlocked)=%.3fs", - plan.metadata.stage_id, - g2_spawned - g2_start, - time.perf_counter() - g2_spawned, - ) - else: - g2_start = time.perf_counter() - with self._replica_launch_lock, stage_runtime_env(plan.metadata.stage_id, plan.metadata.runtime_cfg): - g2_locked = time.perf_counter() - with launch_cm as resources: - pass - logger.debug( - "[stage_init] Stage-%s G2 launch-lock wait=%.3fs, spawn+READY=%.3fs", - plan.metadata.stage_id, - g2_locked - g2_start, - time.perf_counter() - g2_locked, - ) + logger.debug( + "[stage_init] Stage-%s G2 launch-lock wait=%.3fs, spawn=%.3fs, READY=%.3fs", + plan.metadata.stage_id, + g2_locked - g2_start, + g2_spawned - g2_locked, + time.perf_counter() - g2_spawned, + ) logger.info("[StageRuntime] Stage %s engine startup completed", plan.metadata.stage_id) if resources is None: From 3d71a2b3adb1a556d656535795c66295c1b05855 Mon Sep 17 00:00:00 2001 From: andyluo7 <43718156+andyluo7@users.noreply.github.com> Date: Wed, 9 Sep 2026 19:09:39 -0700 Subject: [PATCH 44/93] [CI/Build] Stabilize LTX2 vocoder autocast test on ROCm (#7336) Signed-off-by: andyluo7 --- .../models/ltx2/test_ltx2_vocoder_cuda.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/tests/diffusion/models/ltx2/test_ltx2_vocoder_cuda.py b/tests/diffusion/models/ltx2/test_ltx2_vocoder_cuda.py index d3d885d574a..aa022ed61ee 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_vocoder_cuda.py +++ b/tests/diffusion/models/ltx2/test_ltx2_vocoder_cuda.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""CUDA regression tests for LTX-2 BWE vocoder precision.""" +"""GPU regression tests for LTX-2 BWE vocoder precision.""" import pytest import torch @@ -14,10 +14,13 @@ class _BWEConvVocoder(torch.nn.Module): def __init__(self): super().__init__() + # Match the production LTX vocoder's conv_pre geometry. The previous + # single-channel 1x1 probe intermittently had no MIOpen algorithm. self.conv = torch.nn.Conv1d( - 1, - 1, - kernel_size=1, + 128, + 1024, + kernel_size=7, + padding=3, bias=False, device="cuda", dtype=torch.bfloat16, @@ -33,12 +36,12 @@ def forward(self, value): return output -@hardware_test(res={"cuda": "L4"}, num_cards=1) +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) def test_ltx_bwe_vocoder_runs_real_cuda_autocast_in_fp32(): from vllm_omni.diffusion.models.ltx2.ltx2_runtime import _run_ltx_vocoder vocoder = _BWEConvVocoder() - generated_mel = torch.ones((1, 1, 4), device="cuda", dtype=torch.bfloat16) + generated_mel = torch.ones((1, 128, 64), device="cuda", dtype=torch.bfloat16) output = _run_ltx_vocoder(vocoder, generated_mel) From 35a83a08850f4184e683a82201156038be5ae725 Mon Sep 17 00:00:00 2001 From: Weiming Liao Date: Thu, 10 Sep 2026 10:16:14 +0800 Subject: [PATCH 45/93] [NPU][CI] Add A5 and 310P CI support (#6875) Signed-off-by: Weiming Liao Co-authored-by: wangyu <53896905+yenuo26@users.noreply.github.com> --- .buildkite/common/ci_mirror_hardwares.yml | 51 +++++ .buildkite/common/scripts/upload_pipeline.py | 10 +- .buildkite/npu/bootstrap-upload-steps.yml | 134 +++--------- .buildkite/npu/scripts/build-npu-image.sh | 48 +++++ .buildkite/npu/test-npu-ready.yml | 75 +++++-- docker/Dockerfile.npu | 20 +- docker/Dockerfile.npu.a3 | 36 ---- docker/Dockerfile.npu.ci | 20 +- docker/Dockerfile.npu.ci.a3 | 72 ------- docs/contributing/ci/ci_settings.md | 9 +- docs/contributing/ci/test_writing_guide.md | 4 +- .../installation/npu/npu.inc.md | 198 ++++++++++++++---- pyproject.toml | 2 + tests/diffusion/layers/test_activation.py | 16 +- 14 files changed, 387 insertions(+), 308 deletions(-) create mode 100644 .buildkite/npu/scripts/build-npu-image.sh delete mode 100644 docker/Dockerfile.npu.a3 delete mode 100644 docker/Dockerfile.npu.ci.a3 diff --git a/.buildkite/common/ci_mirror_hardwares.yml b/.buildkite/common/ci_mirror_hardwares.yml index e3c9a427ab7..d638ecb8d5e 100644 --- a/.buildkite/common/ci_mirror_hardwares.yml +++ b/.buildkite/common/ci_mirror_hardwares.yml @@ -17,6 +17,8 @@ # a2b3_npu_1 — ascend-a2b3, resource_class npu-1 (A2/B3 CI image) # a2b3_npu_4 — ascend-a2b3, resource_class npu-4 # a3_npu_2 — ascend-a3, resource_class npu-2 (A3 CI image tag …:$BUILDKITE_COMMIT-a3) +# a5_npu_2 — ascend-a5, resource_class npu-2 +# 310p_npu_1 — ascend-310p, resource_class npu-1 # # B200 (b200-k8s): matches vLLM device ``b200-k8s`` from ci-infra # (AgentQueue.B200_K8S + b200_plugin_template). Queue routes the K8s pool; @@ -646,3 +648,52 @@ mirror_hardwares: resource_class: "npu-16" image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-a3" plugins: *npu_k8s_plugins + + a5_npu_2: + agents: + queue: "ascend-a5" + resource_class: "npu-2" + image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-a5" + plugins: *npu_k8s_plugins + + a5_npu_4: + agents: + queue: "ascend-a5" + resource_class: "npu-4" + image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-a5" + plugins: *npu_k8s_plugins + + a5_npu_8: + agents: + queue: "ascend-a5" + resource_class: "npu-8" + image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-a5" + plugins: *npu_k8s_plugins + + a5_npu_16: + agents: + queue: "ascend-a5" + resource_class: "npu-16" + image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-a5" + plugins: *npu_k8s_plugins + + 310p_npu_1: + agents: + queue: "ascend-310p" + resource_class: "npu-1" + image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-310p" + plugins: *npu_k8s_plugins + + 310p_npu_2: + agents: + queue: "ascend-310p" + resource_class: "npu-2" + image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-310p" + plugins: *npu_k8s_plugins + + 310p_npu_4: + agents: + queue: "ascend-310p" + resource_class: "npu-4" + image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-310p" + plugins: *npu_k8s_plugins diff --git a/.buildkite/common/scripts/upload_pipeline.py b/.buildkite/common/scripts/upload_pipeline.py index 53da6f66065..046b136d776 100644 --- a/.buildkite/common/scripts/upload_pipeline.py +++ b/.buildkite/common/scripts/upload_pipeline.py @@ -70,7 +70,15 @@ LOG = "upload_pipeline" BOOTSTRAP_STEPS_FILENAME = "bootstrap-upload-steps.yml" -BOOTSTRAP_IMAGE_BUILD_KEYS = frozenset({"image-build", "image-build-a2", "image-build-a3"}) +BOOTSTRAP_IMAGE_BUILD_KEYS = frozenset( + { + "image-build", + "image-build-a2", + "image-build-a3", + "image-build-a5", + "image-build-310p", + } +) BOOTSTRAP_UPLOAD_IF_KEYS = { "upload-ready-pipeline": "ready", "upload-merge-pipeline": "merge", diff --git a/.buildkite/npu/bootstrap-upload-steps.yml b/.buildkite/npu/bootstrap-upload-steps.yml index 64b97f8955b..398fbc6dced 100644 --- a/.buildkite/npu/bootstrap-upload-steps.yml +++ b/.buildkite/npu/bootstrap-upload-steps.yml @@ -3,10 +3,10 @@ steps: - label: ":buildkit: Build and Push NPU Test Image - A2" key: image-build-a2 - agents: + agents: &npu-image-build-agents queue: "ascend-a2b3" resource_class: "npu-1" - plugins: + plugins: &npu-image-build-plugins - kubernetes: metadata: annotations: @@ -46,123 +46,47 @@ steps: IMAGE_NAME: "vllm-omni-ci-npu" IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry" BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234" - command: | - set -ex - - echo "--- Building and pushing NPU Test Image" - echo "Image: $${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}" - echo "buildkitd address: $${BUILDKITD_ADDR}" - - if ! command -v buildctl &> /dev/null; then - echo "Installing buildctl..." - mkdir -p /tmp/buildkit - BUILDKIT_VERSION="v0.29.0" - wget -q "https://github.com/moby/buildkit/releases/download/$${BUILDKIT_VERSION}/buildkit-$${BUILDKIT_VERSION}.linux-arm64.tar.gz" -O /tmp/buildkit.tar.gz - tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit - cp /tmp/buildkit/bin/buildctl /usr/local/bin/ - fi - - export DOCKER_CONFIG=/home/user/.docker - buildctl \ - --addr="$${BUILDKITD_ADDR}" \ - --tlscacert=/home/user/.docker/ca.pem \ - --tlscert=/home/user/.docker/cert.pem \ - --tlskey=/home/user/.docker/key.pem \ - build \ - --frontend dockerfile.v0 \ - --local context=. \ - --local dockerfile=./docker \ - --opt filename=Dockerfile.npu.ci \ - --secret id=dockerconfig,src=/home/user/.docker/config.json \ - --output type=image,name=$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG},push=true \ - --progress=plain - - echo "--- Image pushed successfully" - echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}" + command: bash .buildkite/npu/scripts/build-npu-image.sh Dockerfile.npu.ci "$${VLLM_IMAGE_TAG}" - label: ":buildkit: Build and Push NPU Test Image - A3" key: image-build-a3 - agents: - queue: "ascend-a2b3" - resource_class: "npu-1" - plugins: - - kubernetes: - metadata: - annotations: - vault.hashicorp.com/agent-init-first: "true" - vault.hashicorp.com/agent-inject: "true" - vault.hashicorp.com/agent-inject-perms-ca.pem: "0400" - vault.hashicorp.com/agent-inject-perms-cert.pem: "0400" - vault.hashicorp.com/agent-inject-perms-config.json: "0400" - vault.hashicorp.com/agent-inject-perms-key.pem: "0400" - vault.hashicorp.com/agent-inject-secret-ca.pem: internal/data/ascend/buildkitd - vault.hashicorp.com/agent-inject-secret-cert.pem: internal/data/ascend/buildkitd - vault.hashicorp.com/agent-inject-secret-config.json: internal/data/ascend/buildkitd - vault.hashicorp.com/agent-inject-secret-key.pem: internal/data/ascend/buildkitd - vault.hashicorp.com/agent-inject-template-ca.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.RootCA }}\n{{- end }}" - vault.hashicorp.com/agent-inject-template-cert.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaCert }}\n{{- end }}" - vault.hashicorp.com/agent-inject-template-config.json: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.dockerConfig }}\n{{- end }}" - vault.hashicorp.com/agent-inject-template-key.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaKey }}\n{{- end }}" - vault.hashicorp.com/agent-pre-populate-only: "true" - vault.hashicorp.com/agent-run-as-group: "1000" - vault.hashicorp.com/agent-run-as-user: "1000" - vault.hashicorp.com/agent-service-account-token-volume-name: token-vol - vault.hashicorp.com/role: ascend-gha-runners - vault.hashicorp.com/secret-volume-path: /home/user/.docker/ - vault.hashicorp.com/tls-skip-verify: "true" - podSpecPatch: - volumes: - - name: token-vol - projected: - defaultMode: 420 - sources: - - serviceAccountToken: - audience: api - expirationSeconds: 600 - path: token + agents: *npu-image-build-agents + plugins: *npu-image-build-plugins env: VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}-a3" IMAGE_NAME: "vllm-omni-ci-npu" IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry" BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234" - command: | - set -ex - - echo "--- Building and pushing NPU Test Image" - echo "Image: $${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}" - echo "buildkitd address: $${BUILDKITD_ADDR}" + command: bash .buildkite/npu/scripts/build-npu-image.sh Dockerfile.npu.ci "$${VLLM_IMAGE_TAG}" -a3 - if ! command -v buildctl &> /dev/null; then - echo "Installing buildctl..." - mkdir -p /tmp/buildkit - BUILDKIT_VERSION="v0.29.0" - wget -q "https://github.com/moby/buildkit/releases/download/$${BUILDKIT_VERSION}/buildkit-$${BUILDKIT_VERSION}.linux-arm64.tar.gz" -O /tmp/buildkit.tar.gz - tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit - cp /tmp/buildkit/bin/buildctl /usr/local/bin/ - fi - - export DOCKER_CONFIG=/home/user/.docker - buildctl \ - --addr="$${BUILDKITD_ADDR}" \ - --tlscacert=/home/user/.docker/ca.pem \ - --tlscert=/home/user/.docker/cert.pem \ - --tlskey=/home/user/.docker/key.pem \ - build \ - --frontend dockerfile.v0 \ - --local context=. \ - --local dockerfile=./docker \ - --opt filename=Dockerfile.npu.ci.a3 \ - --secret id=dockerconfig,src=/home/user/.docker/config.json \ - --output type=image,name=$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG},push=true \ - --progress=plain + - label: ":buildkit: Build and Push NPU Test Image - A5" + key: image-build-a5 + agents: *npu-image-build-agents + plugins: *npu-image-build-plugins + env: + VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}-a5" + IMAGE_NAME: "vllm-omni-ci-npu" + IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry" + BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234" + command: bash .buildkite/npu/scripts/build-npu-image.sh Dockerfile.npu.ci "$${VLLM_IMAGE_TAG}" -a5 - echo "--- Image pushed successfully" - echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}" + - label: ":buildkit: Build and Push NPU Test Image - 310P" + key: image-build-310p + agents: *npu-image-build-agents + plugins: *npu-image-build-plugins + env: + VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}-310p" + IMAGE_NAME: "vllm-omni-ci-npu" + IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry" + BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234" + command: bash .buildkite/npu/scripts/build-npu-image.sh Dockerfile.npu.ci "$${VLLM_IMAGE_TAG}" -310p - label: "Upload Ready Pipeline" depends_on: - image-build-a2 - image-build-a3 + - image-build-a5 + - image-build-310p key: upload-ready-pipeline agents: queue: "cpu_queue_premerge" @@ -174,6 +98,8 @@ steps: depends_on: - image-build-a2 - image-build-a3 + - image-build-a5 + - image-build-310p key: upload-nightly-pipeline commands: - python3 .buildkite/common/scripts/upload_pipeline.py --upload .buildkite/npu/test-npu-nightly.yml diff --git a/.buildkite/npu/scripts/build-npu-image.sh b/.buildkite/npu/scripts/build-npu-image.sh new file mode 100644 index 00000000000..415fa7d5ead --- /dev/null +++ b/.buildkite/npu/scripts/build-npu-image.sh @@ -0,0 +1,48 @@ +#!/usr/bin/env bash +set -ex + +dockerfile=${1:?Dockerfile is required} +image_tag=${2:?Image tag is required} +device_suffix=${3:-} + +image_name=${IMAGE_NAME:-vllm-omni-ci-npu} +image_registry=${IMAGE_REGISTRY:-swr.cn-southwest-2.myhuaweicloud.com/modelfoundry} +buildkitd_addr=${BUILDKITD_ADDR:-tcp://buildkitd-service.buildkitd:1234} + +echo "--- Building and pushing NPU Test Image" +echo "Image: ${image_registry}/${image_name}:${image_tag}" +echo "Dockerfile: ${dockerfile}" + +if ! command -v buildctl &> /dev/null; then + mkdir -p /tmp/buildkit + buildkit_version=v0.29.0 + wget -q \ + "https://github.com/moby/buildkit/releases/download/${buildkit_version}/buildkit-${buildkit_version}.linux-arm64.tar.gz" \ + -O /tmp/buildkit.tar.gz + tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit + cp /tmp/buildkit/bin/buildctl /usr/local/bin/ +fi + +build_args=() +if [[ -n "${device_suffix}" ]]; then + build_args+=(--opt "build-arg:VLLM_ASCEND_DEVICE_SUFFIX=${device_suffix}") +fi + +export DOCKER_CONFIG=/home/user/.docker +buildctl \ + --addr="${buildkitd_addr}" \ + --tlscacert=/home/user/.docker/ca.pem \ + --tlscert=/home/user/.docker/cert.pem \ + --tlskey=/home/user/.docker/key.pem \ + build \ + --frontend dockerfile.v0 \ + --local context=. \ + --local dockerfile=./docker \ + --opt "filename=${dockerfile}" \ + "${build_args[@]}" \ + --secret id=dockerconfig,src=/home/user/.docker/config.json \ + --output "type=image,name=${image_registry}/${image_name}:${image_tag},push=true" \ + --progress=plain + +echo "--- Image pushed successfully" +echo "${image_registry}/${image_name}:${image_tag}" diff --git a/.buildkite/npu/test-npu-ready.yml b/.buildkite/npu/test-npu-ready.yml index 01fb328790a..d475aad5bfc 100644 --- a/.buildkite/npu/test-npu-ready.yml +++ b/.buildkite/npu/test-npu-ready.yml @@ -1,28 +1,59 @@ steps: - - label: "Platforms Test" - key: ready-platforms-test + - group: ":card_index_dividers: Simple Test" depends_on: upload-ready-pipeline - mirror_hardwares: a3_npu_2 - commands: - - >- - pytest -s -v - tests/platforms/npu/test_minicpmo_code2wav_npugraph.py - -m 'core_model and npu and A3 and omni' - --run-level 'core_model' + steps: + - label: "Platforms Test" + key: ready-platforms-test + mirror_hardwares: a3_npu_2 + commands: + - >- + pytest -s -v + tests/platforms/npu/test_minicpmo_code2wav_npugraph.py + -m 'core_model and npu and A3 and omni' + --run-level 'core_model' - - label: "Diffusion NPU layer tests" - key: ready-diffusion-npu-layer-tests - depends_on: upload-ready-pipeline - mirror_hardwares: a3_npu_2 - commands: - - >- - pytest -s -v - tests/diffusion/layers/test_activation.py - tests/diffusion/layers/test_fused_qk_norm_rope_npu.py - tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py - tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py - -m 'core_model and npu and A3' - --run-level 'core_model' + - label: "Diffusion NPU layer tests - A2" + key: ready-diffusion-npu-layer-tests-a2 + mirror_hardwares: a2b3_npu_1 + commands: + - >- + pytest -s -v + tests/diffusion/layers/test_activation.py + -m 'core_model and npu and A2' + --run-level 'core_model' + + - label: "Diffusion NPU layer tests - A3" + key: ready-diffusion-npu-layer-tests + mirror_hardwares: a3_npu_2 + commands: + - >- + pytest -s -v + tests/diffusion/layers/test_activation.py + tests/diffusion/layers/test_fused_qk_norm_rope_npu.py + tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py + tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py + -m 'core_model and npu and A3' + --run-level 'core_model' + + - label: "Diffusion NPU layer tests - A5" + key: ready-diffusion-npu-layer-tests-a5 + mirror_hardwares: a5_npu_2 + commands: + - >- + pytest -s -v + tests/diffusion/layers/test_activation.py + -m 'core_model and npu and A5' + --run-level 'core_model' + + - label: "Diffusion NPU layer tests - 310P" + key: ready-diffusion-npu-layer-tests-310p + mirror_hardwares: 310p_npu_1 + commands: + - >- + pytest -s -v + tests/diffusion/layers/test_activation.py + -m 'core_model and npu and 310P' + --run-level 'core_model' - group: ":card_index_dividers: Omni Model Test" key: ready-omni-group diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index 3ebacdc9841..0d91c0b3273 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -1,20 +1,7 @@ ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend -ARG VLLM_ASCEND_TAG=v0.28.0 -FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} - -# WORKDIR /vllm-workspace/vllm -# RUN git fetch origin --tags && git checkout v0.18.0 - -# WORKDIR /vllm-workspace/vllm-ascend -# RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a -# # Install vllm-ascend -# # Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH -# RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \ -# source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ -# source /usr/local/Ascend/nnal/atb/set_env.sh && \ -# export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \ -# python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \ -# python3 -m pip cache purge +ARG VLLM_ASCEND_VERSION=v0.28.0 +ARG VLLM_ASCEND_DEVICE_SUFFIX= +FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_VERSION}${VLLM_ASCEND_DEVICE_SUFFIX} ARG APP_DIR=/vllm-workspace/vllm-omni WORKDIR ${APP_DIR} @@ -24,6 +11,7 @@ RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ source /usr/local/Ascend/nnal/atb/set_env.sh && \ export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \ + python3 -m pip install -v "mindiesd==3.1.0" && \ python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation RUN python3 -m pip install -v --no-deps \ diff --git a/docker/Dockerfile.npu.a3 b/docker/Dockerfile.npu.a3 deleted file mode 100644 index d341d456db5..00000000000 --- a/docker/Dockerfile.npu.a3 +++ /dev/null @@ -1,36 +0,0 @@ -ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend -ARG VLLM_ASCEND_TAG=v0.28.0-a3 -FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} - -# WORKDIR /vllm-workspace/vllm -# RUN git fetch origin --tags && git checkout v0.18.0 - -# WORKDIR /vllm-workspace/vllm-ascend -# RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a -# # Install vllm-ascend -# # Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH -# RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \ -# source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ -# source /usr/local/Ascend/nnal/atb/set_env.sh && \ -# export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \ -# python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \ -# python3 -m pip cache purge - -ARG APP_DIR=/vllm-workspace/vllm-omni -WORKDIR ${APP_DIR} -COPY . . - -RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \ - source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ - source /usr/local/Ascend/nnal/atb/set_env.sh && \ - export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \ - python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation - -RUN python3 -m pip install -v --no-deps \ - "stepaudio2-minicpmo>=0.1.1" \ - "step-audio2>=1.0.0" - -ENV VLLM_TARGET_DEVICE=npu -ENV VLLM_WORKER_MULTIPROC_METHOD=spawn - -ENTRYPOINT [] diff --git a/docker/Dockerfile.npu.ci b/docker/Dockerfile.npu.ci index 029a905fdc1..0a150fdd814 100644 --- a/docker/Dockerfile.npu.ci +++ b/docker/Dockerfile.npu.ci @@ -1,10 +1,10 @@ ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend -ARG VLLM_ASCEND_TAG=v0.28.0 -FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} +ARG VLLM_ASCEND_VERSION=v0.28.0 +ARG VLLM_ASCEND_DEVICE_SUFFIX= +FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_VERSION}${VLLM_ASCEND_DEVICE_SUFFIX} ARG APP_DIR=/vllm-workspace/vllm-omni WORKDIR ${APP_DIR} -COPY . . RUN sed -i "s@^\(deb.*\)https\?://[a-z0-9.-]*\.ubuntu\.com@\1http://cache-service.nginx-pypi-cache.svc.cluster.local:8081@g" /etc/apt/sources.list @@ -22,11 +22,9 @@ RUN apt-get update && \ rm -rf /var/cache/apt/* && \ rm -rf /var/lib/apt/lists/* -RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ - source /usr/local/Ascend/nnal/atb/set_env.sh && \ - export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \ - python3 -m pip install -U pip && \ +RUN python3 -m pip install -U pip && \ python3 -m pip install -v \ + "mindiesd==3.1.0" \ "pytest>=7.0.0" \ "pytest-asyncio>=0.21.0" \ "pytest-cov>=4.0.0" \ @@ -48,7 +46,13 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ "datasets>=3.0.0" \ --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \ --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \ - --default-timeout=60 && \ + --default-timeout=60 + +COPY . . + +RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ + source /usr/local/Ascend/nnal/atb/set_env.sh && \ + export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \ python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation \ --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \ --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \ diff --git a/docker/Dockerfile.npu.ci.a3 b/docker/Dockerfile.npu.ci.a3 deleted file mode 100644 index fe3741d0419..00000000000 --- a/docker/Dockerfile.npu.ci.a3 +++ /dev/null @@ -1,72 +0,0 @@ -ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend -ARG VLLM_ASCEND_TAG=v0.28.0-a3 -FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} - -ARG APP_DIR=/vllm-workspace/vllm-omni -WORKDIR ${APP_DIR} -COPY . . - -RUN sed -i "s@^\(deb.*\)https\?://[a-z0-9.-]*\.ubuntu\.com@\1http://cache-service.nginx-pypi-cache.svc.cluster.local:8081@g" /etc/apt/sources.list - -RUN apt-get update && \ - apt-get install -y --no-install-recommends \ - ffmpeg \ - libgl1 \ - libglib2.0-0 \ - libgomp1 \ - libsm6 \ - libxext6 \ - libxrender1 \ - libxcb1 \ - espeak-ng && \ - rm -rf /var/cache/apt/* && \ - rm -rf /var/lib/apt/lists/* - -RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ - source /usr/local/Ascend/nnal/atb/set_env.sh && \ - export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \ - python3 -m pip install -U pip && \ - python3 -m pip install -v \ - "pytest>=7.0.0" \ - "pytest-asyncio>=0.21.0" \ - "pytest-cov>=4.0.0" \ - "pytest-mock>=3.10.0" \ - "openpyxl>=3.0.0" \ - "psutil>=7.2.0" \ - "soundfile>=0.13.1" \ - "pyttsx3>=2.99" \ - "opencc==1.4.1" \ - "opencv-python>=4.12.0.88" \ - "imageio[ffmpeg]>=0.6.0" \ - "jiwer>=3.0.0" \ - "zhon>=2.0.0" \ - "zhconv>=1.4.2" \ - "funasr>=1.0.0" \ - "hyperpyyaml>=1.2.0" \ - "onnx>=1.16.0" \ - "librosa>=0.10.0" \ - "torchmetrics>=1.9.0" \ - "datasets>=3.0.0" \ - --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \ - --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \ - --default-timeout=60 && \ - python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation \ - --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \ - --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \ - --default-timeout=60 - -# MiniCPM-o 4.5 Stage-2 vocoder assets: s3tokenizer plus flashcosyvoice, which -# is shipped by step-audio2. --no-deps keeps pip from replacing the Ascend torch -# build and from shadowing onnxruntime-cann with the vanilla onnxruntime wheel; -# the deps these two actually need are installed above. -RUN python3 -m pip install -v --no-deps \ - "s3tokenizer>=0.3.0" \ - "step-audio2>=1.0.0" \ - --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \ - --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \ - --default-timeout=60 - -ENV VLLM_TARGET_DEVICE=npu -ENV VLLM_WORKER_MULTIPROC_METHOD=spawn - -ENTRYPOINT [] diff --git a/docs/contributing/ci/ci_settings.md b/docs/contributing/ci/ci_settings.md index d578b4b02df..885ebd5280c 100644 --- a/docs/contributing/ci/ci_settings.md +++ b/docs/contributing/ci/ci_settings.md @@ -129,7 +129,7 @@ Canonical layout (prefer these paths for new changes): === "NPU" - **Bootstrap:** [`npu/pipeline-npu.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/npu/pipeline-npu.yml) + [`npu/bootstrap-upload-steps.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/npu/bootstrap-upload-steps.yml)—same split as CUDA; builds A2/B3 and A3 CI images, then uploads child test pipelines. + **Bootstrap:** [`npu/pipeline-npu.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/npu/pipeline-npu.yml) + [`npu/bootstrap-upload-steps.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/npu/bootstrap-upload-steps.yml)—same split as CUDA; builds A2/B3, A3, A5, and 310P CI images, then uploads child test pipelines. **Test YAML:** `npu/test-npu-ready.yml` (L2), `test-npu-nightly.yml` (L4). @@ -137,7 +137,12 @@ Canonical layout (prefer these paths for new changes): **Upload:** same as CUDA—`upload_pipeline.py --upload`. - **Hardware in YAML:** `mirror_hardwares` preset (string), expanded to `agents`, top-level `image`, and `plugins`. Presets: `a2b3_npu_1`, `a2b3_npu_4`, `a3_npu_2` in `common/ci_mirror_hardwares.yml`. + **Hardware in YAML:** `mirror_hardwares` preset (string), expanded to `agents`, top-level `image`, and `plugins`. NPU presets in `common/ci_mirror_hardwares.yml` are: + + - A2: `a2b3_npu_1`, `a2b3_npu_4`, `a2b3_npu_8` + - A3: `a3_npu_2`, `a3_npu_4`, `a3_npu_8`, `a3_npu_16` + - A5: `a5_npu_2`, `a5_npu_4`, `a5_npu_8`, `a5_npu_16` + - 310P: `310p_npu_1`, `310p_npu_2`, `310p_npu_4` **Conventions** diff --git a/docs/contributing/ci/test_writing_guide.md b/docs/contributing/ci/test_writing_guide.md index 469e69cacda..762863bb3e9 100644 --- a/docs/contributing/ci/test_writing_guide.md +++ b/docs/contributing/ci/test_writing_guide.md @@ -48,6 +48,8 @@ Defined in `pyproject.toml`: | `MI325` | Tests that require MI325 GPU (AMD/ROCm) * | | `A2` | Tests that require A2 NPU * | | `A3` | Tests that require A3 NPU * | +| `A5` | Tests that require A5 NPU * | +| `310P` | Tests that require 310P NPU * | | `cards_{n}` | Tests that require *n* accelerator cards (`cards_1` … `cards_8`); auto-added from `num_cards` * | | `slow` | Slow tests (may skip in quick CI) | | `benchmark` | Benchmark tests (decorator on runner test functions; perf JSON uses `full_model` + type marker instead) | @@ -76,7 +78,7 @@ def test_video_to_audio() This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/helpers/mark.py` performs the following actions: 1. **Applies platform and resource markers** - Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `H200`, `B200`, `MI325`, `B60`, `A2`, `A3`). + Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `H200`, `B200`, `MI325`, `B60`, `A2`, `A3`, `A5`, `310P`). ```python @pytest.mark.cuda @pytest.mark.L4 diff --git a/docs/getting_started/installation/npu/npu.inc.md b/docs/getting_started/installation/npu/npu.inc.md index 0019d00bdf5..9ff477522d4 100644 --- a/docs/getting_started/installation/npu/npu.inc.md +++ b/docs/getting_started/installation/npu/npu.inc.md @@ -7,13 +7,34 @@ For detailed hardware and software requirements, please refer to the [vLLM-Ascen The recommended way to use vLLM-Omni on NPU is through the vLLM-Ascend pre-built Docker images: +=== "A2" + + ```bash + export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.28.0 + ``` + +=== "A3" + + ```bash + export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.28.0-a3 + ``` + +=== "A5" + + ```bash + export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.28.0-a5 + ``` + +=== "310P (Experimental)" + + ```bash + export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.28.0-310p + ``` + ```bash -# vLLM-Ascend image -# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.28.0 -# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.28.0-a3 docker run --rm \ --name vllm-omni-npu \ - --shm-size=1g \ + --shm-size=64g \ --device /dev/davinci0 \ --device /dev/davinci1 \ --device /dev/davinci2 \ @@ -27,8 +48,8 @@ docker run --rm \ -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \ -v /etc/ascend_install.info:/etc/ascend_install.info \ -v /root/.cache:/root/.cache \ - -p 8000:8000 \ - -it quay.io/atlas-ci/vllm-ascend:v0.28.0-a3 bash + -p 8091:8091 \ + -it "$IMAGE" bash # Inside the container, install vLLM-Omni from source cd /vllm-workspace @@ -72,53 +93,146 @@ Supported images as following. | image name | Hardware | OS | |-|-|-| -| image-tag | Atlas A2 | Ubuntu | -| image-tag-a3 | Atlas A3 | Ubuntu | +| image-tag | A2 | Ubuntu | +| image-tag-a3 | A3 | Ubuntu | +| image-tag-a5 | A5 | Ubuntu | +| image-tag-310p | 310P (experimental) | Ubuntu | Here's an example deployment command that has been verified on 4 x NPUs: -```bash -# Atlas A2: quay.io/ascend/vllm-omni:v0.28.0 -# Atlas A3: quay.io/ascend/vllm-omni:v0.28.0-a3 -docker run --rm \ - --name vllm-omni-a3 \ - --shm-size=4g \ - --device /dev/davinci0 \ - --device /dev/davinci1 \ - --device /dev/davinci2 \ - --device /dev/davinci3 \ - --device /dev/davinci_manager \ - --device /dev/devmm_svm \ - --device /dev/hisi_hdc \ - -v /usr/local/dcmi:/usr/local/dcmi \ - -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \ - -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \ - -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \ - --device /dev/davinci_manager \ - --device /dev/devmm_svm \ - --device /dev/hisi_hdc \ - -v /usr/local/dcmi:/usr/local/dcmi \ - -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \ - -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \ - -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \ - -v /etc/ascend_install.info:/etc/ascend_install.info \ - -v ~/.cache:/root/.cache \ - -p 8091:8091 \ - -it quay.io/ascend/vllm-omni:v0.28.0-a3 bash -``` +=== "A2" + + ```bash + export IMAGE=quay.io/ascend/vllm-omni:v0.28.0 + docker run --rm \ + --name vllm-omni-a2 \ + --shm-size=64g \ + --device /dev/davinci0 \ + --device /dev/davinci1 \ + --device /dev/davinci2 \ + --device /dev/davinci3 \ + --device /dev/davinci_manager \ + --device /dev/devmm_svm \ + --device /dev/hisi_hdc \ + -v /usr/local/dcmi:/usr/local/dcmi \ + -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \ + -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \ + -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \ + -v /etc/ascend_install.info:/etc/ascend_install.info \ + -v ~/.cache:/root/.cache \ + -p 8091:8091 \ + -it "$IMAGE" bash + ``` + +=== "A3" + + ```bash + export IMAGE=quay.io/ascend/vllm-omni:v0.28.0-a3 + docker run --rm \ + --name vllm-omni-a3 \ + --shm-size=64g \ + --device /dev/davinci0 \ + --device /dev/davinci1 \ + --device /dev/davinci2 \ + --device /dev/davinci3 \ + --device /dev/davinci_manager \ + --device /dev/devmm_svm \ + --device /dev/hisi_hdc \ + -v /usr/local/dcmi:/usr/local/dcmi \ + -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \ + -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \ + -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \ + -v /etc/ascend_install.info:/etc/ascend_install.info \ + -v ~/.cache:/root/.cache \ + -p 8091:8091 \ + -it "$IMAGE" bash + ``` + +=== "A5" + + ```bash + export IMAGE=quay.io/ascend/vllm-omni:v0.28.0-a5 + docker run --rm \ + --name vllm-omni-a5 \ + --shm-size=64g \ + --device /dev/davinci0 \ + --device /dev/davinci1 \ + --device /dev/davinci2 \ + --device /dev/davinci3 \ + --device /dev/davinci_manager \ + --device /dev/devmm_svm \ + --device /dev/hisi_hdc \ + -v /usr/local/dcmi:/usr/local/dcmi \ + -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \ + -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \ + -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \ + -v /etc/ascend_install.info:/etc/ascend_install.info \ + -v ~/.cache:/root/.cache \ + -p 8091:8091 \ + -it "$IMAGE" bash + ``` + +=== "310P (Experimental)" + + ```bash + export IMAGE=quay.io/ascend/vllm-omni:v0.28.0-310p + docker run --rm \ + --name vllm-omni-310p \ + --shm-size=64g \ + --device /dev/davinci0 \ + --device /dev/davinci1 \ + --device /dev/davinci2 \ + --device /dev/davinci3 \ + --device /dev/davinci_manager \ + --device /dev/devmm_svm \ + --device /dev/hisi_hdc \ + -v /usr/local/dcmi:/usr/local/dcmi \ + -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \ + -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \ + -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \ + -v /etc/ascend_install.info:/etc/ascend_install.info \ + -v ~/.cache:/root/.cache \ + -p 8091:8091 \ + -it "$IMAGE" bash + ``` !!! tip You can use this docker image to serve models the same way you would with in vLLM! To do so, make sure you overwrite the default entrypoint (`vllm serve --omni`) which works only for models supported in the vLLM-Omni project. -Or build IMAGE from **source code**: +Or build the image from **source code**: ```bash git clone https://github.com/vllm-project/vllm-omni.git cd vllm-omni -# A2 -# docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu . -# A3 -docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu.a3 . ``` +All hardware variants use `docker/Dockerfile.npu`. Choose the device suffix for your hardware: + +=== "A2" + + ```bash + docker build -t vllm-omni-dev-image:latest -f docker/Dockerfile.npu . + ``` + +=== "A3" + + ```bash + docker build -t vllm-omni-dev-image:latest -f docker/Dockerfile.npu \ + --build-arg VLLM_ASCEND_DEVICE_SUFFIX=-a3 . + ``` + +=== "A5" + + ```bash + docker build -t vllm-omni-dev-image:latest -f docker/Dockerfile.npu \ + --build-arg VLLM_ASCEND_DEVICE_SUFFIX=-a5 . + ``` + +=== "310P (Experimental)" + + ```bash + docker build -t vllm-omni-dev-image:latest -f docker/Dockerfile.npu \ + --build-arg VLLM_ASCEND_DEVICE_SUFFIX=-310p . + ``` + # --8<-- [end:pre-built-images] diff --git a/pyproject.toml b/pyproject.toml index 3ba902977a0..0393c9c6217 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -283,6 +283,8 @@ markers = [ "S5000: [hardware-resource] [musa] Tests that require S5000 GPU (Moore Threads/MUSA)", "A2: [hardware-resource] [npu] Tests that require A2 NPU", "A3: [hardware-resource] [npu] Tests that require A3 NPU", + "A5: [hardware-resource] [npu] Tests that require A5 NPU", + "310P: [hardware-resource] [npu] Tests that require 310P NPU", # Tag with [hardware-cards] for tests.helpers.mark.get_supported_card_counts(). "cards_1: [hardware-cards] Tests that require 1 accelerator card", "cards_2: [hardware-cards] Tests that require 2 accelerator cards", diff --git a/tests/diffusion/layers/test_activation.py b/tests/diffusion/layers/test_activation.py index c3205502f61..321f2ab8f2a 100644 --- a/tests/diffusion/layers/test_activation.py +++ b/tests/diffusion/layers/test_activation.py @@ -5,7 +5,7 @@ import torch import torch.nn.functional as F -from tests.helpers.mark import hardware_test +from tests.helpers.mark import hardware_marks from vllm_omni.diffusion.layers.activation import SiluAndMul pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] @@ -19,9 +19,17 @@ def test_silu_and_mul_native_matches_packed_reference() -> None: torch.testing.assert_close(SiluAndMul().forward_native(packed), F.silu(gate) * up) -@hardware_test(res={"npu": "A3"}, num_cards=1) -def test_silu_and_mul_npu_matches_packed_reference() -> None: - packed = torch.randn(257, 256, device="npu", dtype=torch.bfloat16) +@pytest.mark.parametrize( + "dtype", + [ + pytest.param(torch.bfloat16, marks=hardware_marks(res={"npu": "A2"}, num_cards=1), id="A2-bf16"), + pytest.param(torch.bfloat16, marks=hardware_marks(res={"npu": "A3"}, num_cards=1), id="A3-bf16"), + pytest.param(torch.bfloat16, marks=hardware_marks(res={"npu": "A5"}, num_cards=1), id="A5-bf16"), + pytest.param(torch.float16, marks=hardware_marks(res={"npu": "310P"}, num_cards=1), id="310P-fp16"), + ], +) +def test_silu_and_mul_npu_matches_packed_reference(dtype: torch.dtype) -> None: + packed = torch.randn(257, 256, device="npu", dtype=dtype) gate, up = packed.chunk(2, dim=-1) torch.testing.assert_close( From 4b8dcac68153cc386db1e5fb0226d52518739414 Mon Sep 17 00:00:00 2001 From: Mu GuanLin <1203789601@qq.com> Date: Thu, 10 Sep 2026 11:07:35 +0800 Subject: [PATCH 46/93] [Kernel] Enable LTX DiffVAE fusions on SM100 and SM103 (#7350) Signed-off-by: mglyn <1203789601@qq.com> --- .../models/ltx2/test_ltx2_ops_platform.py | 6 ++++-- vllm_omni/diffusion/models/ltx2/ops/__init__.py | 14 ++++++++------ 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/tests/diffusion/models/ltx2/test_ltx2_ops_platform.py b/tests/diffusion/models/ltx2/test_ltx2_ops_platform.py index 6ffff3136f4..5fca0e5e65d 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_ops_platform.py +++ b/tests/diffusion/models/ltx2/test_ltx2_ops_platform.py @@ -37,8 +37,10 @@ def clear_platform_cache(): (90, True, True, True, True), (90, True, False, True, False), (90, False, True, False, False), - (100, True, True, False, False), - (103, True, True, False, False), + (100, True, True, True, False), + (100, True, False, True, False), + (103, True, True, True, False), + (103, True, False, True, False), (110, True, True, False, False), (120, True, True, False, False), (121, True, True, False, False), diff --git a/vllm_omni/diffusion/models/ltx2/ops/__init__.py b/vllm_omni/diffusion/models/ltx2/ops/__init__.py index 93f2ff67291..7907fba564e 100644 --- a/vllm_omni/diffusion/models/ltx2/ops/__init__.py +++ b/vllm_omni/diffusion/models/ltx2/ops/__init__.py @@ -17,9 +17,10 @@ # Pointwise fusions are launch/memory-traffic reductions whose first execution # for every runtime key is checked bit-for-bit; failures use the ordinary model -# math. Both paths remain SM90-only until other architectures are validated. -# The FNA schedule is performance-tuned for SM90 and has no silent fallback. -_CUDA_COMPUTE_CAPABILITIES = frozenset({90}) +# math. The FNA schedule is performance-tuned for SM90 and has no silent +# fallback, so qualify the two paths independently. +_FUSION_CUDA_COMPUTE_CAPABILITIES = frozenset({90, 100, 103}) +_FNA_CUDA_COMPUTE_CAPABILITIES = frozenset({90}) @dataclass(frozen=True) @@ -66,9 +67,10 @@ def resolve_ltx2_vae_operators(device: torch.device) -> LTX2VAEOperatorSet | Non if device.type != "cuda" or not current_omni_platform.is_cuda() or not current_omni_platform.is_available(): return None device_index = device.index if device.index is not None else torch.accelerator.current_device_index() - if _cuda_compute_capability(int(device_index)) not in _CUDA_COMPUTE_CAPABILITIES: + capability = _cuda_compute_capability(int(device_index)) + if capability not in _FUSION_CUDA_COMPUTE_CAPABILITIES: return None - return _operator_set(has_tilelang()) + return _operator_set(capability in _FNA_CUDA_COMPUTE_CAPABILITIES and has_tilelang()) def _eligible_operators(tensor: torch.Tensor) -> LTX2VAEOperatorSet | None: @@ -78,7 +80,7 @@ def _eligible_operators(tensor: torch.Tensor) -> LTX2VAEOperatorSet | None: def is_ltx2_fusion_eligible(tensor: torch.Tensor) -> bool: - """Allow self-verifying pointwise fusions on validated SM90 devices.""" + """Allow self-verifying pointwise fusions on qualified CUDA devices.""" return _eligible_operators(tensor) is not None From 333f14e97ebbe78f208380fd82ff4523033d355b Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Wed, 9 Sep 2026 20:38:11 -0700 Subject: [PATCH 47/93] [Bugfix][MiniCPM-o] Align structured chat content with native omni rendering (#7344) Signed-off-by: Sy03 <1370724210@qq.com> --- .buildkite/cuda/test-merge.yml | 1 + .buildkite/cuda/test-ready.yml | 1 + pyproject.toml | 2 +- recipes/OpenBMB/MiniCPM-o-4_5.md | 21 +++- tests/e2e/online_serving/test_minicpmo_4_5.py | 29 ++++- .../models/minicpmo_4_5/test_chat_template.py | 45 +++++++ .../chat_templates/minicpmo45_native.jinja | 116 ++++++++++++++++++ 7 files changed, 204 insertions(+), 11 deletions(-) create mode 100644 tests/model_executor/models/minicpmo_4_5/test_chat_template.py create mode 100644 vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index 6bad98f7959..5aa49ea6cb1 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.buildkite/cuda/test-merge.yml @@ -155,6 +155,7 @@ steps: source_file_dependencies: - tests/e2e/offline_inference/test_minicpmo_4_5.py - tests/e2e/online_serving/test_minicpmo_4_5.py + - vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja - vllm_omni/model_executor/models/minicpmo_4_5/ - vllm_omni/model_executor/models/step_audio2/ - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index f54a58fa48e..6d8d1e8b98a 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.buildkite/cuda/test-ready.yml @@ -188,6 +188,7 @@ steps: - label: "Omni · MiniCPM-o 4.5 Online Test" source_file_dependencies: - tests/e2e/online_serving/test_minicpmo_4_5.py + - vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja - vllm_omni/model_executor/models/minicpmo_4_5/ - vllm_omni/model_executor/models/step_audio2/ - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ diff --git a/pyproject.toml b/pyproject.toml index 0393c9c6217..2584c8c7a24 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -161,7 +161,7 @@ where = ["."] include = ["vllm_omni*"] [tool.setuptools.package-data] -"vllm_omni" = ["_version.py", "py.typed"] +"vllm_omni" = ["_version.py", "py.typed", "transformers_utils/chat_templates/*.jinja"] [tool.setuptools_scm] # no extra settings needed, presence enables setuptools-scm diff --git a/recipes/OpenBMB/MiniCPM-o-4_5.md b/recipes/OpenBMB/MiniCPM-o-4_5.md index 86f0883f769..fb0934d4f51 100644 --- a/recipes/OpenBMB/MiniCPM-o-4_5.md +++ b/recipes/OpenBMB/MiniCPM-o-4_5.md @@ -28,13 +28,13 @@ also provided. - Default deploy configs (auto-loaded by HF `model_type=minicpmo` + `hf_config.version="4.5"`): - - Default single-GPU compatibility layout (auto-loaded): + - Default single-GPU compatibility layout (auto-loaded): [`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml) - - Recommended 2-GPU continuous-batching layout: + - Recommended 2-GPU continuous-batching layout: [`vllm_omni/deploy/minicpmo_4_5_2gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_2gpu.yaml), - - 3-GPU layout: + - 3-GPU layout: [`vllm_omni/deploy/minicpmo_4_5_3gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_3gpu.yaml) - - 8x RTX 4090 layout: + - 8x RTX 4090 layout: [`vllm_omni/deploy/minicpmo_4_5_8x4090.yaml`](../../vllm_omni/deploy/minicpmo_4_5_8x4090.yaml) - Online example + Gradio demo: [`examples/online_serving/minicpmo/`](../../examples/online_serving/minicpmo/) @@ -126,9 +126,20 @@ video. Use the two-GPU profile for production throughput. ```bash vllm serve openbmb/MiniCPM-o-4_5 --omni \ --trust-remote-code \ + --chat-template vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja \ + --chat-template-content-format openai \ --host 0.0.0.0 --port 8099 ``` +Run these commands from the repository root. The bundled chat template +matches the Hugging Face `chat(omni_mode=True)` content layout: media and text +parts are concatenated without adding separators, while whitespace explicitly +included in a text part is preserved. This matters for speech generation because +an inserted newline changes the Thinker condition passed to Talker. For image +questions, place the image before the question, as in the native image-chat +layout. Continue to request `use_tts_template: true` and +`enable_thinking: false` for spoken answers. + The deploy config is auto-loaded by the model registry — no `--deploy-config` flag needed for this default single-GPU layout. @@ -271,6 +282,8 @@ uses GPU 5. GPUs 6–7 are left free. vllm serve openbmb/MiniCPM-o-4_5 --omni \ --deploy-config vllm_omni/deploy/minicpmo_4_5_8x4090.yaml \ --trust-remote-code \ + --chat-template vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja \ + --chat-template-content-format openai \ --host 0.0.0.0 --port 8099 ``` diff --git a/tests/e2e/online_serving/test_minicpmo_4_5.py b/tests/e2e/online_serving/test_minicpmo_4_5.py index 290064fc31b..5417d2b1652 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5.py +++ b/tests/e2e/online_serving/test_minicpmo_4_5.py @@ -24,6 +24,9 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" _MODEL = "openbmb/MiniCPM-o-4_5" +_NATIVE_CHAT_TEMPLATE = ( + Path(__file__).resolve().parents[3] / "vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja" +) _CI_DEPLOY = modify_stage_config( get_deploy_config_path("minicpmo_4_5.yaml"), updates={ @@ -63,7 +66,14 @@ model=_MODEL, stage_config_path=_CI_DEPLOY, use_stage_cli=False, - server_args=["--trust-remote-code", "--async-chunk"], + server_args=[ + "--trust-remote-code", + "--async-chunk", + "--chat-template", + str(_NATIVE_CHAT_TEMPLATE), + "--chat-template-content-format", + "openai", + ], env_dict={ "VLLM_CONFIGURE_LOGGING": "1", "VLLM_LOGGING_CONFIG_PATH": str(_PROMPT_LOGGING_CONFIG_PATH), @@ -265,11 +275,18 @@ def test_image_to_text_audio_001(omni_server, openai_client) -> None: Input Setting: stream=True """ image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(24, 24)['base64']}" - messages = dummy_messages_from_mix_data( - system_prompt=get_system_prompt(), - image_data_url=image_data_url, - content_text=get_prompt("text_image"), - ) + # Match the native image-chat layout explicitly; do not rely on the + # renderer moving media placeholders ahead of text. + messages = [ + get_system_prompt(), + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": image_data_url}}, + {"type": "text", "text": get_prompt("text_image")}, + ], + }, + ] request_config = { "model": omni_server.model, diff --git a/tests/model_executor/models/minicpmo_4_5/test_chat_template.py b/tests/model_executor/models/minicpmo_4_5/test_chat_template.py new file mode 100644 index 00000000000..8f1068392c2 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_chat_template.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""MiniCPM structured-content rendering must preserve native part boundaries.""" + +from pathlib import Path + +import pytest +from jinja2 import TemplateError +from transformers.utils.chat_template_utils import _compile_jinja_template + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +TEMPLATE = Path(__file__).resolve().parents[4] / "vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja" + + +@pytest.mark.parametrize( + "parts,expected", + [ + ([{"type": "image"}, {"type": "text", "text": "Question?"}], "(./)Question?"), + ([{"type": "text", "text": "Question?"}, {"type": "image"}], "Question?(./)"), + ([{"type": "image"}, {"type": "text", "text": "\n\nQuestion?"}], "(./)\n\nQuestion?"), + ([{"type": "text", "text": " A "}, {"type": "text", "text": " B "}], " A B "), + ([{"type": "audio"}, {"type": "text", "text": "Transcribe."}], "()Transcribe."), + ([{"type": "video"}, {"type": "image"}], "()(./)"), + ("\nKeep this whitespace.\n", "\nKeep this whitespace.\n"), + ], +) +def test_native_content_concatenation(parts, expected): + rendered = _compile_jinja_template(TEMPLATE.read_text()).render( + messages=[{"role": "user", "content": parts}], + add_generation_prompt=True, + enable_thinking=False, + use_tts_template=True, + ) + assert rendered == ( + f"<|im_start|>user\n{expected}<|im_end|>\n<|im_start|>assistant\n\n\n\n\n<|tts_bos|>" + ) + + +def test_unsupported_part_is_not_silently_discarded(): + with pytest.raises(TemplateError, match="Unsupported MiniCPM-o content part"): + _compile_jinja_template(TEMPLATE.read_text()).render( + messages=[{"role": "user", "content": [{"type": "unknown"}]}], + add_generation_prompt=True, + ) diff --git a/vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja b/vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja new file mode 100644 index 00000000000..2f2bc98399a --- /dev/null +++ b/vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja @@ -0,0 +1,116 @@ +{# SPDX-License-Identifier: Apache-2.0 + Adapted from OpenBMB/MiniCPM-o-4_5 tokenizer_config.json. + Match chat(omni_mode=True): concatenate structured content parts without + injecting separators, preserving whitespace explicitly supplied as text. +#} +{%- set normalized = namespace(messages=[]) -%} +{%- for message in messages -%} + {%- if message.content is string or message.content is none -%} + {%- set normalized.messages = normalized.messages + [message] -%} + {%- else -%} + {%- set content = namespace(text='') -%} + {%- for part in message.content -%} + {%- if part.type == 'text' -%} + {%- set content.text = content.text + part.text -%} + {%- elif part.type == 'image' -%} + {%- set content.text = content.text + '(./)' -%} + {%- elif part.type == 'audio' -%} + {%- set content.text = content.text + '()' -%} + {%- elif part.type == 'video' -%} + {%- set content.text = content.text + '()' -%} + {%- else -%} + {{- raise_exception('Unsupported MiniCPM-o content part: ' + part.type) -}} + {%- endif -%} + {%- endfor -%} + {%- set normalized.messages = normalized.messages + [dict(message, content=content.text)] -%} + {%- endif -%} +{%- endfor -%} +{%- set messages = normalized.messages -%} +{%- if tools %} + {{- '<|im_start|>system\n' }} + {%- if messages[0].role == 'system' %} + {{- messages[0].content + '\n\n' }} + {%- endif %} + {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- "\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n<|im_end|>\n" }} +{%- else %} + {%- if messages[0].role == 'system' %} + {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }} + {%- endif %} +{%- endif %} +{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %} +{%- for message in messages[::-1] %} + {%- set index = (messages|length - 1) - loop.index0 %} + {%- if ns.multi_step_tool and message.role == "user" and not(message.content.startswith('') and message.content.endswith('')) %} + {%- set ns.multi_step_tool = false %} + {%- set ns.last_query_index = index %} + {%- endif %} +{%- endfor %} +{%- for message in messages %} + {%- if (message.role == "user") or (message.role == "system" and not loop.first) %} + {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {%- set content = message.content %} + {%- set reasoning_content = '' %} + {%- if message.reasoning_content is defined and message.reasoning_content is not none %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if '' in message.content %} + {%- set content = message.content.split('')[-1].lstrip('\n') %} + {%- set reasoning_content = message.content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %} + {%- endif %} + {%- endif %} + {%- if loop.index0 > ns.last_query_index %} + {%- if loop.last or (not loop.last and reasoning_content) %} + {{- '<|im_start|>' + message.role + '\n\n' + reasoning_content.strip('\n') + '\n\n\n' + content.lstrip('\n') }} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + content }} + {%- endif %} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + content }} + {%- endif %} + {%- if message.tool_calls %} + {%- for tool_call in message.tool_calls %} + {%- if (loop.first and content) or (not loop.first) %} + {{- '\n' }} + {%- endif %} + {%- if tool_call.function %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {{- '\n{"name": "' }} + {{- tool_call.name }} + {{- '", "arguments": ' }} + {%- if tool_call.arguments is string %} + {{- tool_call.arguments }} + {%- else %} + {{- tool_call.arguments | tojson }} + {%- endif %} + {{- '}\n' }} + {%- endfor %} + {%- endif %} + {{- '<|im_end|>\n' }} + {%- elif message.role == "tool" %} + {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {{- message.content }} + {{- '\n' }} + {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %} + {{- '<|im_end|>\n' }} + {%- endif %} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} + {%- if enable_thinking is defined and enable_thinking is false %} + {{- '\n\n\n\n' }} + {%- endif %} + {%- if use_tts_template is defined and use_tts_template is true %} + {{- '<|tts_bos|>' }} + {%- endif %} +{%- endif -%} From aff7d64948f6c0e81e9b3272234623044e215967 Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Thu, 10 Sep 2026 12:08:28 +0800 Subject: [PATCH 48/93] [Rebase] Rebase to vLLM 0.29.0 (#7230) Signed-off-by: tzhouam Signed-off-by: Zhou Taichang Co-authored-by: Claude Opus 5 --- .buildkite/common/ci_mirror_hardwares.yml | 27 +- .buildkite/common/scripts/upload_pipeline.py | 13 + docker/Dockerfile.ci | 29 +- docs/configuration/README.md | 2 +- docs/configuration/environment_variables.md | 2 +- docs/getting_started/installation/README.md | 2 +- docs/getting_started/installation/gpu.md | 2 + .../installation/gpu/cuda.inc.md | 22 +- .../installation/gpu/rocm.inc.md | 12 +- docs/getting_started/quickstart.md | 6 +- tests/attention/test_fish_kvcache_attn.py | 19 +- tests/buildkite/test_upload_pipeline.py | 32 + tests/config/test_omni_config.py | 10 + ...t_omni_scheduler_forces_v1_model_runner.py | 59 ++ .../diffusion_kv/test_block_tables.py | 6 +- .../diffusion_kv/test_initialization.py | 8 +- tests/diffusion/diffusion_kv/test_layout.py | 147 ++++ tests/diffusion/diffusion_kv/test_manager.py | 7 +- .../test_paged_attention_adapter.py | 4 +- .../test_paged_attention_adapter_gpu.py | 20 +- tests/diffusion/diffusion_kv/test_request.py | 6 +- .../diffusion_kv/test_spec_discovery.py | 11 +- tests/diffusion/layers/test_adalayernorm.py | 4 +- tests/diffusion/layers/test_rope_broadcast.py | 27 +- .../test_boogu_image_transformer.py | 2 +- .../models/dreamzero/test_qkv_fusion.py | 4 +- .../glm_image/test_glm_image_quantization.py | 4 +- .../test_mistral_encoder_tp.py | 11 +- .../qwen_image/test_qwen_image_rope_utils.py | 29 +- .../sana_video/test_transformer_sana_video.py | 4 +- .../test_transformer_sana_video_sp.py | 10 +- .../test_transformer_sana_video_tp.py | 4 +- .../models/t5_encoder/test_t5_encoder_tp.py | 9 +- .../quantization/test_component_routing.py | 5 +- tests/diffusion/test_diffusion_scheduler.py | 7 +- .../test_diffusers_backend_similarity.py | 1 + .../online_serving/test_dots_tts_expansion.py | 2 + .../test_async_omni_engine_stage_init.py | 79 +-- .../openai/test_batch_chat_completions.py | 5 +- .../openai_api/test_audio_format.py | 4 +- .../openai_api/test_modality_validation.py | 3 +- .../openai_api/test_serving_chat_metrics.py | 3 +- .../openai_api/test_serving_speech.py | 2 +- .../openai_api/test_video_stream_handler.py | 4 +- .../entrypoints/test_stream_finish_reason.py | 3 +- tests/helpers/kv_layout.py | 36 + tests/helpers/serving_chat.py | 3 +- tests/inputs/test_mm_processor.py | 44 ++ tests/inputs/test_preprocess.py | 178 ++++- .../models/glm_image/test_glm_image_ar.py | 37 +- .../minicpmo_4_5/test_audio_embeds_input.py | 8 +- .../test_long_audio_processing.py | 40 +- .../test_processor_prompt_tokenization.py | 102 +++ .../qwen2_5_omni/test_per_video_audio_mask.py | 14 +- .../qwen3_omni/test_qwen3_omni_moe_lora.py | 4 +- .../qwen3_tts/test_qwen3_tts_code2wav.py | 7 +- .../models/test_weight_loader.py | 38 ++ .../test_modelopt_nvfp4_nan_clamp.py | 13 +- .../test_upstream_prompt_update_api.py | 64 ++ tests/platforms/test_platform_detection.py | 38 ++ vllm_omni/benchmarks/patch/patch.py | 2 + vllm_omni/core/sched/omni_scheduler_mixin.py | 11 + .../diffusion/attention/backends/abstract.py | 7 + vllm_omni/diffusion/attention/layer.py | 11 +- .../diffusion/diffusion_kv/initialization.py | 28 +- vllm_omni/diffusion/diffusion_kv/layout.py | 174 +++++ .../diffusion_kv/model_runner_backend.py | 11 +- .../diffusion_kv/paged_attention_adapter.py | 4 +- vllm_omni/diffusion/layers/rope.py | 33 +- .../hunyuan_image3/pipeline_hunyuan_image3.py | 7 +- .../models/ltx2/ltx2_phase_adapter.py | 2 +- .../qwen_image/qwen_image_transformer.py | 37 +- vllm_omni/engine/arg_utils.py | 4 +- .../engine/stage_engine_core_proc_manager.py | 8 +- vllm_omni/engine/stage_init_utils.py | 25 +- vllm_omni/entrypoints/async_omni.py | 8 +- vllm_omni/entrypoints/cli/serve.py | 2 +- vllm_omni/entrypoints/duplex/chat_fallback.py | 2 +- vllm_omni/entrypoints/openai/api_server.py | 29 +- vllm_omni/entrypoints/openai/batch_serving.py | 4 +- .../entrypoints/openai/realtime_connection.py | 5 +- vllm_omni/entrypoints/openai/serving_chat.py | 30 +- .../entrypoints/openai/serving_speech.py | 8 +- .../entrypoints/openai/video_api_utils.py | 4 +- .../openai/video_stream_session.py | 4 +- .../ar_diffusion/kv_cache/manager.py | 5 +- vllm_omni/inputs/mm_processor.py | 39 ++ vllm_omni/inputs/preprocess.py | 276 +++----- .../models/audex/audex_omni_processing.py | 7 +- .../model_executor/models/bagel/bagel.py | 131 ++-- .../models/cosyvoice3/cosyvoice3.py | 52 +- .../models/covo_audio/covo_audio.py | 10 +- .../models/gepard/gepard_talker.py | 10 +- .../models/glm_image/glm_image_ar.py | 549 +++------------ .../model_executor/models/glm_tts/glm_tts.py | 60 +- .../higgs_audio_v3/higgs_audio_v3_talker.py | 8 +- .../hunyuan_image3/_hunyuan_v1_vendored.py | 625 ++++++++++++++++++ .../models/hunyuan_image3/hunyuan_image3.py | 39 +- .../models/mimo_audio/mimo_audio.py | 13 +- .../models/mimo_audio/mimo_audio_llm.py | 40 +- .../ming_flash_omni/ming_flash_omni_talker.py | 14 +- .../ming_flash_omni_thinker.py | 148 ++--- .../minicpmo_4_5/minicpmo_4_5_omni_llm.py | 188 ++---- .../models/minimax_h3/text_encoder.py | 84 ++- .../models/minimax_music3/talker.py | 10 +- .../models/omnivoice/omnivoice.py | 89 ++- .../qwen2_5_omni/qwen2_5_omni_talker.py | 11 +- .../qwen2_5_omni/qwen2_5_omni_thinker.py | 120 ++-- .../models/qwen2_5_omni/qwen2_old.py | 12 +- .../models/qwen3_omni/qwen3_omni_code2wav.py | 10 +- .../qwen3_omni/qwen3_omni_moe_talker.py | 14 +- .../qwen3_omni/qwen3_omni_moe_thinker.py | 184 +----- .../models/qwen3_tts/qwen3_tts_code2wav.py | 12 +- .../models/qwen3_tts/qwen3_tts_talker.py | 10 +- .../models/step_audio2/step_audio2_thinker.py | 57 +- .../voxtral_tts_audio_generation.py | 31 +- vllm_omni/platforms/__init__.py | 8 +- .../reasoning/step_audio_reasoning_parser.py | 2 +- vllm_omni/worker/gpu_ar_model_runner.py | 1 - vllm_omni/worker/sampling_utils.py | 15 +- 120 files changed, 2955 insertions(+), 1687 deletions(-) create mode 100644 tests/core/sched/test_omni_scheduler_forces_v1_model_runner.py create mode 100644 tests/diffusion/diffusion_kv/test_layout.py create mode 100644 tests/helpers/kv_layout.py create mode 100644 tests/inputs/test_mm_processor.py create mode 100644 tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py create mode 100644 tests/model_executor/models/test_weight_loader.py create mode 100644 tests/model_executor/test_upstream_prompt_update_api.py create mode 100644 tests/platforms/test_platform_detection.py create mode 100644 vllm_omni/diffusion/diffusion_kv/layout.py create mode 100644 vllm_omni/inputs/mm_processor.py create mode 100644 vllm_omni/model_executor/models/hunyuan_image3/_hunyuan_v1_vendored.py diff --git a/.buildkite/common/ci_mirror_hardwares.yml b/.buildkite/common/ci_mirror_hardwares.yml index d638ecb8d5e..f51f13f1ae0 100644 --- a/.buildkite/common/ci_mirror_hardwares.yml +++ b/.buildkite/common/ci_mirror_hardwares.yml @@ -28,6 +28,9 @@ # packed onto g6.12xlarge L4 nodes (``vllm.ci/gpu-pool=l4x4``). GPU / CPU / # memory / shm requests scale with card count (1 GPU = 10 CPU / 40Gi / 24Gi shm). # Retry covers lost agents / killed pods, not pytest or OOM (exit 1 / 137). +# The Kubernetes HF secret deliberately uses VLLM_CI_HF_TOKEN. Scheduled +# build-level HF_TOKEN values can override pod env entries with the same name; +# upload_pipeline.py restores the secret-backed value inside the job shell. x-npu-k8s-plugins: &npu_k8s_plugins - kubernetes: @@ -86,7 +89,7 @@ mirror_hardwares: env: - name: HF_HOME value: /fsx/hf_cache - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -137,7 +140,7 @@ mirror_hardwares: env: - name: HF_HOME value: /fsx/hf_cache - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -188,7 +191,7 @@ mirror_hardwares: env: - name: HF_HOME value: /fsx/hf_cache - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -239,7 +242,7 @@ mirror_hardwares: env: - name: HF_HOME value: /fsx/hf_cache - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -273,7 +276,7 @@ mirror_hardwares: env: - name: HF_HOME value: /root/.cache/huggingface - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -308,7 +311,7 @@ mirror_hardwares: env: - name: HF_HOME value: /root/.cache/huggingface - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -343,7 +346,7 @@ mirror_hardwares: env: - name: HF_HOME value: /root/.cache/huggingface - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -378,7 +381,7 @@ mirror_hardwares: env: - name: HF_HOME value: /root/.cache/huggingface - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -428,7 +431,7 @@ mirror_hardwares: value: "0" - name: HF_HOME value: /raid/hf_cache - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -479,7 +482,7 @@ mirror_hardwares: value: "0" - name: HF_HOME value: /raid/hf_cache - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -530,7 +533,7 @@ mirror_hardwares: value: "0" - name: HF_HOME value: /raid/hf_cache - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret @@ -581,7 +584,7 @@ mirror_hardwares: value: "0" - name: HF_HOME value: /raid/hf_cache - - name: HF_TOKEN + - name: VLLM_CI_HF_TOKEN valueFrom: secretKeyRef: name: hf-token-secret diff --git a/.buildkite/common/scripts/upload_pipeline.py b/.buildkite/common/scripts/upload_pipeline.py index 046b136d776..749f35f5a68 100644 --- a/.buildkite/common/scripts/upload_pipeline.py +++ b/.buildkite/common/scripts/upload_pipeline.py @@ -87,6 +87,8 @@ } E2E_GROUP_MARKER = "E2E Test" CI_MIRROR_HARDWARES_PATH = ROOT / ".buildkite/common/ci_mirror_hardwares.yml" +CUDA_HF_TOKEN_ENV = "VLLM_CI_HF_TOKEN" +CUDA_HF_TOKEN_EXPORT = f'if [ -n "$${{{CUDA_HF_TOKEN_ENV}:-}}" ]; then export HF_TOKEN="$${{{CUDA_HF_TOKEN_ENV}}}"; fi' # Bootstrap Buildkite ``if`` expressions. # ``*_MAIN_IF``: main + env schedule. ``*_LABEL_IF``: PR label (and/or composed with MAIN). @@ -509,6 +511,17 @@ def _expand_mirror_hardwares(step: dict[str, Any]) -> dict[str, Any] | None: expanded = copy.deepcopy(preset) merged = {key: value for key, value in step.items() if key != "mirror_hardwares"} | expanded + # A scheduled build's HF_TOKEN may override a pod env entry with the same + # name. The preset injects the Kubernetes secret under a collision-proof + # alias; restore the conventional name after Buildkite has merged env. + if any(preset_name == chip or preset_name.startswith(f"{chip}_") for chip in _cuda_mirror_chips()): + commands = merged.get("commands") + if isinstance(commands, list): + merged["commands"] = [CUDA_HF_TOKEN_EXPORT, *commands] + elif commands is None: + merged["commands"] = [CUDA_HF_TOKEN_EXPORT] + else: + merged["commands"] = [CUDA_HF_TOKEN_EXPORT, commands] # Preset retry (K8S_RETRY on l4_*) must not clobber a step that opted out. if "retry" in step: merged["retry"] = step["retry"] diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index b84b344a879..a8e46580b40 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -1,27 +1,10 @@ -# Target vllm version: v0.28.0 (tag v0.28.0 @ 2cf0a6915ce5). -# Keep VLLM_BASE_TAG aligned with the target vLLM release on every rebase. -# -# Release-boundary form: there is no precompiled-wheel pin here, because the -# base image already ships exactly the target vLLM. Everything the runtime -# needs — torch, flashinfer, nvJitLink, numpy — comes from that image as one -# self-consistent set. -# -# This block reappears only when the rebase targets a vLLM commit that is NOT a -# released tag. During the 0.27 rebase, targeting v0.27.0rc1 on top of the -# v0.26.0 base image required reinstalling vLLM from wheels.vllm.ai and then -# hand-repairing the fallout: pinning torch ahead of flashinfer (installing it -# after bound flashinfer's TVM FFI runtime to the old torch and aborted engine -# core init), registering libnvJitLink.so.13 with ldconfig, and re-pinning numpy -# under numba's ceiling. All of it was compensation for a base image and a vLLM -# that disagreed. At a real release boundary they agree, so none of it is -# needed. The 0.28 rebase repeated the cycle: the v0.28.0rc / final-tag wheels -# rode on the v0.27.0 base image with the full reinstall block until the -# official v0.28.0 image shipped, at which point this file dropped back to the -# form you are reading. If you pin an untagged commit again, restore the -# reinstall block — agent/tasks/20_update_wheel_url.sh fails the run rather -# than letting the Dockerfile silently ship the wrong vLLM. +# Target vLLM version: v0.29.0 (tag @ 98dff2a81d74). +# Use the matching release image and its coherent CUDA/PyTorch/FlashInfer +# runtime. Do not reinstall a commit wheel or repair dependencies on top of +# an older base image at a release boundary. For a future unreleased target, +# restore an explicit wheel pin and validate its complete runtime stack. ARG VLLM_BASE_IMAGE=vllm/vllm-openai -ARG VLLM_BASE_TAG=v0.28.0 +ARG VLLM_BASE_TAG=v0.29.0 # We use registry-based layer caching for the dependencies, # where the tag is determined based on the hash of the contents diff --git a/docs/configuration/README.md b/docs/configuration/README.md index cdfb5818ace..3d065f9a1f7 100644 --- a/docs/configuration/README.md +++ b/docs/configuration/README.md @@ -2,7 +2,7 @@ This section lists the most common options for running vLLM-Omni. -For options within a vLLM Engine, please refer to the [vLLM 0.28 configuration guide](https://docs.vllm.ai/en/v0.28.0/configuration/index.html). +For options within a vLLM Engine, please refer to the [vLLM 0.29 configuration guide](https://docs.vllm.ai/en/v0.29.0/configuration/index.html). Each model defines fixed topology in a registered `PipelineConfig` and runtime overrides in a deploy YAML. diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md index 8f944f44dfa..90b9ffae092 100644 --- a/docs/configuration/environment_variables.md +++ b/docs/configuration/environment_variables.md @@ -151,7 +151,7 @@ their keys only. ## Inherited vLLM variables vLLM-Omni also reads variables through its aligned vLLM dependency. Refer to -the [vLLM 0.28 environment-variable reference](https://docs.vllm.ai/en/v0.28.0/configuration/env_vars.html) +the [vLLM 0.29 environment-variable reference](https://docs.vllm.ai/en/v0.29.0/configuration/env_vars.html) for their definitions. This includes vLLM launch, cache, logging, plugin, ROCm, XPU, ModelScope, and FlashInfer workspace settings. diff --git a/docs/getting_started/installation/README.md b/docs/getting_started/installation/README.md index 448ca81f26e..9d2d11e94c5 100644 --- a/docs/getting_started/installation/README.md +++ b/docs/getting_started/installation/README.md @@ -1,7 +1,7 @@ # Installation !!! important - vLLM-Omni is released against the matching upstream vLLM major/minor version. vLLM-Omni 0.28.x requires vLLM 0.28.x; the stable v0.28.0 instructions pin vLLM 0.28.0. + vLLM-Omni is released against the matching upstream vLLM major/minor version. The 0.29 development line requires vLLM 0.29.x; the source-install instructions pin vLLM 0.29.0. Published vLLM-Omni 0.28.0 wheels and images still require vLLM 0.28.x. vLLM-Omni supports the following hardware platforms: diff --git a/docs/getting_started/installation/gpu.md b/docs/getting_started/installation/gpu.md index 7b87782967f..e714a52bfa7 100644 --- a/docs/getting_started/installation/gpu.md +++ b/docs/getting_started/installation/gpu.md @@ -34,6 +34,8 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr ### Pre-built wheels +The source-install instructions target vLLM 0.29.0. Until vLLM-Omni 0.29.0 artifacts are published, the CUDA/ROCm pre-built examples below remain on the matching 0.28.0 release pair. + Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, 0.26.0, and 0.28.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source). === "NVIDIA CUDA" diff --git a/docs/getting_started/installation/gpu/cuda.inc.md b/docs/getting_started/installation/gpu/cuda.inc.md index 5e6cd47c9ce..f6a1004a656 100644 --- a/docs/getting_started/installation/gpu/cuda.inc.md +++ b/docs/getting_started/installation/gpu/cuda.inc.md @@ -5,7 +5,7 @@ # --8<-- [end:requirements] # --8<-- [start:set-up-using-python] -vLLM-Omni depends on the matching major/minor release of vLLM. The vLLM-Omni 0.28.x release line uses vLLM 0.28.x. +vLLM-Omni depends on the matching major/minor release of vLLM. The 0.29 development line uses vLLM 0.29.x. Published 0.28.0 wheels and images use vLLM 0.28.x. !!! note PyTorch installed via `conda` will statically link `NCCL` library, which can cause issues when vLLM tries to use `NCCL`. See for more details. @@ -18,6 +18,8 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new** #### Installation of vLLM +These pre-built wheel instructions install the published vLLM-Omni 0.28.0 release. For the 0.29 development line, use the source-install instructions below. + vLLM-Omni is built based on vLLM. Please install it with command below. ```bash uv pip install vllm==0.28.0 --torch-backend=auto @@ -26,12 +28,12 @@ uv pip install vllm==0.28.0 --torch-backend=auto #### Installation of vLLM-Omni ```bash -uv pip install vllm-omni +uv pip install vllm-omni==0.28.0 ``` To run Gradio demos, also install the optional extras: ```bash -uv pip install 'vllm-omni[demo]' +uv pip install 'vllm-omni[demo]==0.28.0' ``` # --8<-- [end:pre-built-wheels] @@ -39,13 +41,13 @@ uv pip install 'vllm-omni[demo]' # --8<-- [start:build-wheel-from-source] #### Installation of vLLM -If you do not need to modify source code of vLLM, you can directly install the stable 0.28.0 release version of the library +If you do not need to modify source code of vLLM, you can directly install the stable 0.29.0 release version of the library ```bash -uv pip install vllm==0.28.0 --torch-backend=auto +uv pip install vllm==0.29.0 --torch-backend=auto ``` -The 0.28.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead. +The 0.29.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead. #### Installation of vLLM-Omni Since vllm-omni is rapidly evolving, it's recommended to install it from source @@ -66,12 +68,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -git checkout v0.28.0 +git checkout v0.29.0 ``` Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file. ```bash -#For CUDA 13.0 (the default for v0.28.0; the wheel filename has no `+cu130` suffix) -export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.28.0/vllm-0.28.0-cp38-abi3-manylinux_2_28_x86_64.whl +#For CUDA 13.0 (the default for v0.29.0; the wheel filename has no `+cu130` suffix) +export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.29.0/vllm-0.29.0-cp38-abi3-manylinux_2_28_x86_64.whl ``` Install vllm with command below (If you have no existing PyTorch). ```bash @@ -131,7 +133,7 @@ If you want to specify the base vLLM version: ```bash DOCKER_BUILDKIT=1 docker build \ -f docker/Dockerfile.cuda \ - --build-arg BASE_IMAGE=vllm/vllm-openai:v0.28.0 \ + --build-arg BASE_IMAGE=vllm/vllm-openai:v0.29.0 \ -t vllm-omni-cuda . ``` diff --git a/docs/getting_started/installation/gpu/rocm.inc.md b/docs/getting_started/installation/gpu/rocm.inc.md index e0f1614aecf..d15bba2eb44 100644 --- a/docs/getting_started/installation/gpu/rocm.inc.md +++ b/docs/getting_started/installation/gpu/rocm.inc.md @@ -11,6 +11,8 @@ vLLM-Omni current recommends the steps in under setup through Docker Images. #### Installation of vLLM +These pre-built wheel instructions install the published vLLM-Omni 0.28.0 release. For the 0.29 development line, use the source-install instructions below. + vLLM-Omni is built based on vLLM. Please install it with command below. ```bash uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723 @@ -22,7 +24,7 @@ uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/roc # we need to add --no-build-isolation as the torch # is not obtained from pypi, we have to install using the # torch installed in our environment -uv pip install vllm-omni +uv pip install vllm-omni==0.28.0 # Optional if want to run Qwen3 TTS uv pip uninstall onnxruntime # should be removed before we can install onnxruntime-rocm @@ -34,13 +36,13 @@ uv pip install onnxruntime-rocm # --8<-- [start:build-wheel-from-source] #### Installation of vLLM -If you do not need to modify source code of vLLM, you can directly install the stable 0.28.0 release version of the library +If you do not need to modify source code of vLLM, you can directly install the stable 0.29.0 release version of the library ```bash -uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723 +uv pip install vllm==0.29.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.29.0/rocm723 ``` -The pre-built 0.28.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead. +The pre-built 0.29.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead. #### Installation of vLLM-Omni Since vllm-omni is rapidly evolving, it's recommended to install it from source @@ -58,7 +60,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -git checkout v0.28.0 +git checkout v0.29.0 python3 -m pip install -r requirements/rocm.txt python3 setup.py develop ``` diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index 036c2da1766..8363e05ff54 100644 --- a/docs/getting_started/quickstart.md +++ b/docs/getting_started/quickstart.md @@ -19,10 +19,10 @@ uv venv --python 3.12 --seed source .venv/bin/activate # On CUDA -uv pip install vllm==0.28.0 --torch-backend=auto +uv pip install vllm==0.29.0 --torch-backend=auto # On ROCm -uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723 +uv pip install vllm==0.29.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.29.0/rocm723 git clone https://github.com/vllm-project/vllm-omni.git cd vllm-omni @@ -34,7 +34,7 @@ For additional installation methods — please see the [installation guide](inst !!! note It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni. - If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.28.0` and vLLM Omni `0.28.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue. + If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.29.0` and vLLM Omni `0.29.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue. ## Offline Inference diff --git a/tests/attention/test_fish_kvcache_attn.py b/tests/attention/test_fish_kvcache_attn.py index d38380e13e7..4a8ea17e825 100644 --- a/tests/attention/test_fish_kvcache_attn.py +++ b/tests/attention/test_fish_kvcache_attn.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 pytest import torch @@ -214,11 +214,20 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, sc assert fish_kvcache_backend.get_fish_kvcache_attn_stats()["small_hit_count"] == 1 -def _vllm_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size): +def _vllm_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size, dtype=torch.float16): # Source of truth for the paged KV cache layout the backend must unpack. - from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend - - return FlashAttentionBackend.get_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size) + # vLLM 0.29 removed AttentionBackend.get_kv_cache_shape; the per-layer shape + # now comes from the spec, with the content dim expressed in bytes. + from vllm.v1.kv_cache_interface import FullAttentionSpec, compute_layer_kv_cache_shape_bytes + + spec = FullAttentionSpec( + block_size=block_size, + num_kv_heads=num_kv_heads, + head_size=head_size, + dtype=dtype, + ) + num_pages, num_heads, num_states, content_bytes = compute_layer_kv_cache_shape_bytes(spec, num_blocks) + return (num_pages, num_heads, num_states, content_bytes // dtype.itemsize) def test_fish_kvcache_backend_unbinds_kv_on_vllm_cache_layout(monkeypatch): diff --git a/tests/buildkite/test_upload_pipeline.py b/tests/buildkite/test_upload_pipeline.py index 9011a1aefeb..7492a4b8281 100644 --- a/tests/buildkite/test_upload_pipeline.py +++ b/tests/buildkite/test_upload_pipeline.py @@ -12,6 +12,7 @@ from skip_ci import resolve_ci_decision # noqa: E402 from upload_pipeline import ( # noqa: E402 + CUDA_HF_TOKEN_EXPORT, _expand_mirror_hardwares, _get_mirror_hw_selector, _load_bootstrap_steps, @@ -143,6 +144,10 @@ def test_mirror_hardwares_l4_1_expands_to_agents_and_plugins(monkeypatch: pytest container = step["plugins"][0]["kubernetes"]["podSpec"]["containers"][0] assert container["image"].endswith("$BUILDKITE_COMMIT") assert container["resources"]["limits"]["nvidia.com/gpu"] == 1 + env_names = {item["name"] for item in container["env"]} + assert "VLLM_CI_HF_TOKEN" in env_names + assert "HF_TOKEN" not in env_names + assert step["commands"] == [CUDA_HF_TOKEN_EXPORT, "pytest -sv tests/example"] def test_mirror_hardwares_l4_preserves_explicit_retry() -> None: @@ -182,6 +187,33 @@ def test_mirror_hardwares_a2b3_npu_4_expands_agents_image_and_plugins() -> None: assert step["plugins"][0]["kubernetes"]["podSpecPatch"]["imagePullSecrets"] == [ {"name": "swr-secret"}, ] + assert step["commands"] == ["pytest -sv tests/example"] + + +def test_all_cuda_mirror_hardwares_restore_hf_token_at_runtime() -> None: + for name in ( + "l4_1", + "l4_2", + "l4_3", + "l4_4", + "h100_1", + "h100_2", + "h100_3", + "h100_4", + "b200_1", + "b200_2", + "b200_3", + "b200_4", + ): + step = _expand_mirror_hardwares( + {"label": name, "mirror_hardwares": name, "commands": ["pytest -sv tests/example"]}, + ) + assert step is not None + container = step["plugins"][0]["kubernetes"]["podSpec"]["containers"][0] + env_names = {item["name"] for item in container["env"]} + assert "VLLM_CI_HF_TOKEN" in env_names + assert "HF_TOKEN" not in env_names + assert step["commands"][0] == CUDA_HF_TOKEN_EXPORT def _gpu_limit(step: dict) -> int: diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 3e499dc3843..40229f52154 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -215,6 +215,16 @@ def test_from_pipeline_config_normalizes_stage_engine_extras_without_expanding_s assert stage.diffusion_config.model_config["default_robot_embodiment"] == "roboarena" +@pytest.mark.parametrize("disabled", [True, False]) +def test_frontend_log_stats_flag_is_not_an_unowned_stage_argument(disabled): + from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config + + config = _from_pipeline_key("dots_tts", cli_overrides={"disable_log_stats": disabled}) + assert config.stage_configs + engine_args = build_engine_args_dict_from_omni_stage_config(config.stage_by_id(0), model="test-model") + assert "disable_log_stats" not in engine_args + + def test_from_pipeline_config_applies_cli_overrides_without_stage_config_runtime_bridge(): omni_config = _from_pipeline_key( "qwen3_tts", diff --git a/tests/core/sched/test_omni_scheduler_forces_v1_model_runner.py b/tests/core/sched/test_omni_scheduler_forces_v1_model_runner.py new file mode 100644 index 00000000000..fb8546e2dc0 --- /dev/null +++ b/tests/core/sched/test_omni_scheduler_forces_v1_model_runner.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The omni scheduler must stay on the v1 model runner, like every omni worker. + +vLLM 0.28 resolved this config to v1; 0.29 ends its fallback chain with v2. When +only the workers forced v1, the scheduler took its v2 fast path -- which omits +resumed-request token ids -- while the v1 runner still read them, so resuming a +request raised ``KeyError: `` in ``_update_states`` and killed the engine +core (Multi-Replica Startup Test). + +The default is upstream's to change, so this pins Omni's side of the contract +rather than the default itself. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Scheduler(OmniSchedulerMixin): + """Minimal carrier for the shared init hook both omni schedulers call.""" + + def __init__(self, *, use_v2_model_runner: bool) -> None: + self.use_v2_model_runner = use_v2_model_runner + self.vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + async_chunk=False, + stage_id=0, + pooling_output_decoder=None, + ) + ) + self._init_omni_io_scheduling_state() + + +@pytest.mark.parametrize("initial", [True, False]) +def test_scheduler_ends_on_v1_model_runner(initial: bool) -> None: + """Whatever upstream resolves, the omni scheduler runs v1.""" + assert _Scheduler(use_v2_model_runner=initial).use_v2_model_runner is False + + +def test_v2_default_does_not_reach_the_resumed_request_fast_path() -> None: + """The v2 fast path is what dropped resumed-request token ids. + + ``OmniGenerationScheduler`` branches on ``use_v2_model_runner`` when it + assembles SchedulerOutput; forcing it false keeps the v1 assembly that + carries ``all_token_ids`` for resumed requests, which the v1 runner reads. + """ + scheduler = _Scheduler(use_v2_model_runner=True) + + assert not scheduler.use_v2_model_runner, ( + "scheduler would take the v2 fast path while omni workers run v1; " + "resumed requests then KeyError in _update_states" + ) diff --git a/tests/diffusion/diffusion_kv/test_block_tables.py b/tests/diffusion/diffusion_kv/test_block_tables.py index f2002f1c6a7..aa8e0730774 100644 --- a/tests/diffusion/diffusion_kv/test_block_tables.py +++ b/tests/diffusion/diffusion_kv/test_block_tables.py @@ -279,17 +279,17 @@ def init_backend(kv_cache_config, vllm_config, device): group = SimpleNamespace(get_metadata_builder=lambda _index: SimpleNamespace(reorder_batch_threshold=None)) return [[group], [group]], "cg-support", [4, 8] + # vLLM 0.29 dropped attn_groups and cache_dtype from init_kv_cache: the + # layout now rides on the resolved CacheConfig.kv_cache_layout instead. def init_cache( runner_kv_caches, forward_context, kv_cache_config, - attn_groups, device, - cache_dtype, kernel_block_sizes, vllm_config, ): - events.append(("cache", kv_cache_config, attn_groups, kernel_block_sizes)) + events.append(("cache", kv_cache_config, kernel_block_sizes)) runner_kv_caches.extend(["cache-0", "cache-1"]) return {"layer-0": "tensor-0", "layer-1": "tensor-1"} diff --git a/tests/diffusion/diffusion_kv/test_initialization.py b/tests/diffusion/diffusion_kv/test_initialization.py index 861f9564d9c..b5f914b11a3 100644 --- a/tests/diffusion/diffusion_kv/test_initialization.py +++ b/tests/diffusion/diffusion_kv/test_initialization.py @@ -299,8 +299,10 @@ def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None ) assert vllm_config.model_config.original_max_model_len == -1 - # Sixteen physical pages provide ``16 * block_size`` tokens. The native - # backend chooses the block geometry, so keep this assertion backend + # Sixteen physical pages provide ``(16 - 1) * block_size`` tokens: vLLM 0.29 + # subtracts the null block that BlockPool permanently holds back before + # sizing (check_enough_kv_cache_memory), so only 15 pages are usable. The + # native backend chooses the block geometry, so keep this assertion backend # independent instead of assuming the generic 16-token block size. - expected_max_model_len = min(model_limit, spec.block_size * 16) + expected_max_model_len = min(model_limit, spec.block_size * (16 - 1)) assert vllm_config.model_config.max_model_len == expected_max_model_len diff --git a/tests/diffusion/diffusion_kv/test_layout.py b/tests/diffusion/diffusion_kv/test_layout.py new file mode 100644 index 00000000000..713849e95bf --- /dev/null +++ b/tests/diffusion/diffusion_kv/test_layout.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Physical KV cache layout selection and propagation for the Diffusion cache.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from vllm.config import CacheConfig +from vllm.v1.kv_cache_interface import FullAttentionSpec +from vllm.v1.kv_cache_layout import KVCacheLayout + +from vllm_omni.diffusion.diffusion_kv.layout import ( + DEFAULT_DIFFUSION_KV_CACHE_LAYOUT, + adopt_kv_cache_layout, + assert_backend_layout_supported, + build_kv_cache_tensor, + resolve_diffusion_kv_cache_layout, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _config(): + return SimpleNamespace(cache_config=CacheConfig()) + + +class _DenseBackend: + @classmethod + def indexes_kv_by_block_stride(cls) -> bool: + return False + + +class _BlockStrideBackend: + @classmethod + def indexes_kv_by_block_stride(cls) -> bool: + return True + + +def test_default_layout_is_not_block_outermost() -> None: + """The default reproduces the pre-0.29 indexes_kv_by_block_stride=False.""" + vllm_config = _config() + layout = resolve_diffusion_kv_cache_layout(vllm_config) + + assert layout is DEFAULT_DIFFUSION_KV_CACHE_LAYOUT + assert layout.is_block_outermost is False + assert vllm_config.cache_config.get_resolved_kv_cache_layout() is layout + + +def test_block_stride_backend_selects_block_outermost_layout() -> None: + layout = resolve_diffusion_kv_cache_layout(_config(), indexes_kv_by_block_stride=True) + + assert layout.is_block_outermost is True + + +def test_resolution_is_idempotent() -> None: + vllm_config = _config() + + assert resolve_diffusion_kv_cache_layout(vllm_config) is resolve_diffusion_kv_cache_layout(vllm_config) + + +def test_contradicting_requirement_raises_instead_of_misaddressing() -> None: + vllm_config = _config() + resolve_diffusion_kv_cache_layout(vllm_config) + + with pytest.raises(ValueError, match="contradicts the attention"): + resolve_diffusion_kv_cache_layout(vllm_config, indexes_kv_by_block_stride=True) + + +def test_preset_layout_is_honoured() -> None: + vllm_config = _config() + vllm_config.cache_config.kv_cache_layout = KVCacheLayout.BLNHC.name + + layout = resolve_diffusion_kv_cache_layout(vllm_config, indexes_kv_by_block_stride=True) + + assert layout is KVCacheLayout.BLNHC + + +def test_backend_assertion_is_quiet_until_the_layout_is_resolved() -> None: + """Workers collect specs before the control plane hands them the layout.""" + vllm_config = _config() + + assert_backend_layout_supported(vllm_config, _BlockStrideBackend) # unresolved: nothing to check + + resolve_diffusion_kv_cache_layout(vllm_config) + assert_backend_layout_supported(vllm_config, _DenseBackend) + assert_backend_layout_supported(vllm_config, None) + with pytest.raises(ValueError, match="contradicts the attention"): + assert_backend_layout_supported(vllm_config, _BlockStrideBackend) + + +def test_worker_adopts_the_layout_carried_on_the_kv_cache_config() -> None: + """Diffusion workers receive only a KVCacheConfig, so the name rides on it.""" + worker = _config() + shipped = SimpleNamespace(kv_cache_layout=KVCacheLayout.BLNHC.name) + + assert adopt_kv_cache_layout(worker, shipped) is KVCacheLayout.BLNHC + assert worker.cache_config.get_resolved_kv_cache_layout() is KVCacheLayout.BLNHC + + +def test_worker_without_a_carried_layout_falls_back_to_the_same_default() -> None: + worker = _config() + + assert adopt_kv_cache_layout(worker, SimpleNamespace(kv_cache_layout=None)) is DEFAULT_DIFFUSION_KV_CACHE_LAYOUT + + +def test_worker_disagreeing_with_the_control_plane_raises() -> None: + worker = _config() + resolve_diffusion_kv_cache_layout(worker) + + with pytest.raises(ValueError, match="disagrees with the layout"): + adopt_kv_cache_layout(worker, SimpleNamespace(kv_cache_layout=KVCacheLayout.BLNHC.name)) + + +def test_block_stride_requirement_propagates_control_plane_to_worker() -> None: + """A block-stride backend must reach the worker as a block-outermost layout. + + Covers the whole seam: the control plane resolves with the requirement, the + name rides on each rank-local KVCacheConfig, the worker adopts it, and the + backend contract check then passes instead of raising. + """ + control_plane = _config() + layout = resolve_diffusion_kv_cache_layout(control_plane, indexes_kv_by_block_stride=True) + assert layout.is_block_outermost is True + + shipped = SimpleNamespace(kv_cache_layout=layout.name) + + worker = _config() + assert adopt_kv_cache_layout(worker, shipped) is layout + # The worker's own backend requirement now agrees with what it was handed. + assert_backend_layout_supported(worker, _BlockStrideBackend) + with pytest.raises(ValueError, match="contradicts the attention"): + assert_backend_layout_supported(worker, _DenseBackend) + + +def test_tensor_strides_come_from_the_upstream_derivation() -> None: + spec = FullAttentionSpec(block_size=16, num_kv_heads=2, head_size=8, dtype=torch.bfloat16) + + tensor = build_kv_cache_tensor(spec, 4, ["layer0", "layer1"]) + + assert tensor.layers == ["layer0", "layer1"] + assert tensor.size == spec.page_size_bytes * 2 * 4 + # Layer-outermost: each layer owns a contiguous run of its own blocks. + assert tensor.layer_stride == spec.page_size_bytes * 4 + assert tensor.block_stride == spec.page_size_bytes diff --git a/tests/diffusion/diffusion_kv/test_manager.py b/tests/diffusion/diffusion_kv/test_manager.py index 5152c2a6746..ed39f652329 100644 --- a/tests/diffusion/diffusion_kv/test_manager.py +++ b/tests/diffusion/diffusion_kv/test_manager.py @@ -1,11 +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 import pytest import torch import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec +from tests.helpers.kv_layout import build_kv_cache_tensor from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVAdmissionError, DiffusionKVCacheManager from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest @@ -24,7 +25,7 @@ def _config(num_blocks: int) -> KVCacheConfig: ) return KVCacheConfig( num_blocks=num_blocks, - kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])], + kv_cache_tensors=[build_kv_cache_tensor(spec, num_blocks, ["layer0"])], kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)], ) diff --git a/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py b/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py index 1f242420d57..535c1b41163 100644 --- a/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py +++ b/tests/diffusion/diffusion_kv/test_paged_attention_adapter.py @@ -15,6 +15,7 @@ from torch import nn from vllm.v1.kv_cache_interface import FullAttentionSpec +from tests.helpers.kv_layout import layout_for_backend from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.diffusion_kv import paged_attention_adapter as adapter_module @@ -1239,7 +1240,8 @@ def select_backend(**_kwargs): assert native_layer.num_heads == 4 assert native_layer.num_kv_heads == 2 assert native_layer.spec.num_kv_heads == 2 - assert native_layer.spec.indexes_kv_by_block_stride is True + # 0.29: the flag no longer rides on the spec, it selects the layout. + assert layout_for_backend(native_layer.attn_backend).is_block_outermost is True assert native_layer._q_scale_float == 1.0 assert native_layer._k_scale_float == 1.0 assert native_layer._v_scale_float == 1.0 diff --git a/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py b/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py index ef5d75f4d4d..d890dcc7b44 100644 --- a/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.py +++ b/tests/diffusion/diffusion_kv/test_paged_attention_adapter_gpu.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 __future__ import annotations @@ -10,12 +10,14 @@ from torch import nn from torch.nn import functional as F from vllm.config import AttentionConfig, CacheConfig, VllmConfig, set_current_vllm_config -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec from vllm.v1.worker.gpu.attn_utils import init_attn_backend, init_kv_cache from vllm.v1.worker.gpu.block_table import BlockTables +from tests.helpers.kv_layout import build_kv_cache_tensor from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl +from vllm_omni.diffusion.diffusion_kv.layout import resolve_diffusion_kv_cache_layout from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import ( DiffusionPagedAttentionAdapter, DiffusionPagedAttentionLayerAdapter, @@ -88,12 +90,7 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N canonical_spec = native_layer.spec kv_cache_config = KVCacheConfig( num_blocks=4, - kv_cache_tensors=[ - KVCacheTensor( - size=canonical_spec.page_size_bytes * 4, - shared_by=[_LAYER_NAME], - ) - ], + kv_cache_tensors=[build_kv_cache_tensor(canonical_spec, 4, [_LAYER_NAME])], kv_cache_groups=[ KVCacheGroupSpec( layer_names=[_LAYER_NAME], @@ -101,6 +98,9 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N ) ], ) + # 0.29 reads the resolved physical layout when allocating; the diffusion + # control plane pins it, so this standalone harness must too. + resolve_diffusion_kv_cache_layout(vllm_config) with set_current_vllm_config(vllm_config): attn_groups, _, kernel_block_sizes = init_attn_backend( kv_cache_config, @@ -108,13 +108,13 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N device, ) runner_kv_caches: list[torch.Tensor | list[torch.Tensor]] = [] + # vLLM 0.29 dropped attn_groups and cache_dtype here: the layout now + # comes from the resolved CacheConfig.kv_cache_layout inside the call. init_kv_cache( runner_kv_caches, vllm_config.compilation_config.static_forward_context, kv_cache_config, - attn_groups, device, - vllm_config.cache_config.cache_dtype, kernel_block_sizes, vllm_config, ) diff --git a/tests/diffusion/diffusion_kv/test_request.py b/tests/diffusion/diffusion_kv/test_request.py index ae5609b2a2c..ac9185e097a 100644 --- a/tests/diffusion/diffusion_kv/test_request.py +++ b/tests/diffusion/diffusion_kv/test_request.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 pytest import torch @@ -10,10 +10,10 @@ FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, - KVCacheTensor, ) from vllm.v1.request import RequestStatus +from tests.helpers.kv_layout import build_kv_cache_tensor from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -119,7 +119,7 @@ def _manager(*, enable_caching: bool) -> KVCacheManager: group = KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec) config = KVCacheConfig( num_blocks=8, - kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["layer0"])], + kv_cache_tensors=[build_kv_cache_tensor(spec, 8, ["layer0"])], kv_cache_groups=[group], ) return KVCacheManager( diff --git a/tests/diffusion/diffusion_kv/test_spec_discovery.py b/tests/diffusion/diffusion_kv/test_spec_discovery.py index 9485ed046e5..1ccef7a9f85 100644 --- a/tests/diffusion/diffusion_kv/test_spec_discovery.py +++ b/tests/diffusion/diffusion_kv/test_spec_discovery.py @@ -8,9 +8,10 @@ import pytest import torch from torch import nn -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module +from tests.helpers.kv_layout import build_kv_cache_tensor, layout_for_backend from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner @@ -79,7 +80,9 @@ def test_runner_discovers_native_spec_from_loaded_attention() -> None: assert spec.num_kv_heads == 2 assert spec.head_size == 8 assert spec.dtype is torch.bfloat16 - assert spec.indexes_kv_by_block_stride is True + # vLLM 0.29 moved the block-stride decision off the spec onto the resolved + # physical layout; the backend's declaration must still select one. + assert layout_for_backend(_Backend).is_block_outermost is True assert spec.non_causal is True @@ -112,7 +115,7 @@ def test_runner_retains_matching_rank_local_config() -> None: spec = runner.get_kv_cache_spec()["image_attention"] config = KVCacheConfig( num_blocks=8, - kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["image_attention"])], + kv_cache_tensors=[build_kv_cache_tensor(spec, 8, ["image_attention"])], kv_cache_groups=[KVCacheGroupSpec(layer_names=["image_attention"], kv_cache_spec=spec)], ) @@ -126,7 +129,7 @@ def test_runner_rejects_rank_local_config_for_different_layers() -> None: spec = runner.get_kv_cache_spec()["image_attention"] config = KVCacheConfig( num_blocks=8, - kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["other_attention"])], + kv_cache_tensors=[build_kv_cache_tensor(spec, 8, ["other_attention"])], kv_cache_groups=[KVCacheGroupSpec(layer_names=["other_attention"], kv_cache_spec=spec)], ) diff --git a/tests/diffusion/layers/test_adalayernorm.py b/tests/diffusion/layers/test_adalayernorm.py index 5e41b7a26d0..94cadf5acc1 100644 --- a/tests/diffusion/layers/test_adalayernorm.py +++ b/tests/diffusion/layers/test_adalayernorm.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 """Unit tests for shared AdaLayerNorm layers used by FLUX and other models.""" import os @@ -46,7 +46,7 @@ def _force_default_gemm(monkeypatch): monkeypatch.setattr( "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", - lambda: default_unquantized_gemm, + lambda *_args, **_kwargs: default_unquantized_gemm, ) diff --git a/tests/diffusion/layers/test_rope_broadcast.py b/tests/diffusion/layers/test_rope_broadcast.py index 284049347c1..d7d26350de8 100644 --- a/tests/diffusion/layers/test_rope_broadcast.py +++ b/tests/diffusion/layers/test_rope_broadcast.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 """ Shape broadcasting regression tests for RotaryEmbedding.forward_native. @@ -16,7 +16,11 @@ import pytest import torch -from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rotary_emb_mindiesd +from vllm_omni.diffusion.layers.rope import ( + RotaryEmbedding, + RotaryEmbeddingWan, + apply_rotary_emb_mindiesd, +) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -239,6 +243,25 @@ def test_wan_forward_native_unchanged(self) -> None: assert not torch.isnan(output).any() +def test_cuda_paths_fall_back_without_vendored_rotary_kernel() -> None: + """Source/CPU installs need not contain the wheel-vendored CUDA kernel.""" + x = torch.randn(1, 5, 2, 8) + + rope = RotaryEmbedding(is_neox_style=True) + rope.apply_rotary_emb_vllm_flash_attn = None + output = rope.forward_cuda(x, torch.randn(5, 4), torch.randn(5, 4)) + assert output.shape == x.shape + + wan_rope = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True) + wan_rope.apply_rotary_emb_vllm_flash_attn = None + output = wan_rope.forward_cuda( + x, + torch.randn(1, 5, 1, 4), + torch.randn(1, 5, 1, 4), + ) + assert output.shape == x.shape + + def test_mindie_full_dim_rope_keeps_h3_cos_sin_layout(monkeypatch: pytest.MonkeyPatch) -> None: """H3 passes complete rotary dims to MindIE instead of the CUDA half-dim layout.""" captured: dict[str, object] = {} diff --git a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py index 4296c1cde79..43e2f12f4cf 100644 --- a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py +++ b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py @@ -59,7 +59,7 @@ def _force_default_gemm(monkeypatch): monkeypatch.setattr( "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", - lambda: default_unquantized_gemm, + lambda *_args, **_kwargs: default_unquantized_gemm, ) diff --git a/tests/diffusion/models/dreamzero/test_qkv_fusion.py b/tests/diffusion/models/dreamzero/test_qkv_fusion.py index 16b78693f69..d67a01e485e 100644 --- a/tests/diffusion/models/dreamzero/test_qkv_fusion.py +++ b/tests/diffusion/models/dreamzero/test_qkv_fusion.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 """Equivalence test for DreamZero self-attn QKV fusion. @@ -52,7 +52,7 @@ def _force_default_gemm(monkeypatch): monkeypatch.setattr( "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", - lambda: default_unquantized_gemm, + lambda *_args, **_kwargs: default_unquantized_gemm, ) diff --git a/tests/diffusion/models/glm_image/test_glm_image_quantization.py b/tests/diffusion/models/glm_image/test_glm_image_quantization.py index 8ab046f0cef..1a6170ca20c 100644 --- a/tests/diffusion/models/glm_image/test_glm_image_quantization.py +++ b/tests/diffusion/models/glm_image/test_glm_image_quantization.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 """Unit tests for GLM-Image quantization support (W4A16/AutoRound). These tests verify that the GLM-Image DiT transformer correctly accepts and uses @@ -69,7 +69,7 @@ def _force_default_gemm(monkeypatch): monkeypatch.setattr( "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", - lambda: default_unquantized_gemm, + lambda *_args, **_kwargs: default_unquantized_gemm, ) diff --git a/tests/diffusion/models/mistral_encoder/test_mistral_encoder_tp.py b/tests/diffusion/models/mistral_encoder/test_mistral_encoder_tp.py index 20024abe799..62bcf9ce8bb 100644 --- a/tests/diffusion/models/mistral_encoder/test_mistral_encoder_tp.py +++ b/tests/diffusion/models/mistral_encoder/test_mistral_encoder_tp.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 import pytest @@ -64,6 +64,13 @@ def setup_tp_group(monkeypatch, mocker): lambda: 0, ) + # vLLM 0.29 added a CUDA-only fused embedding kernel that engages whenever + # tp_size > 1; these tests drive CPU tensors, so keep the portable masked path. + monkeypatch.setattr( + "vllm.model_executor.layers.vocab_parallel_embedding.current_platform.is_cuda", + lambda: False, + ) + mock_tp_group = mocker.MagicMock() mock_tp_group.world_size = 2 mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_tp_group) @@ -94,7 +101,7 @@ def setup_tp_group(monkeypatch, mocker): monkeypatch.setattr( "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", - lambda: default_unquantized_gemm, + lambda *_args, **_kwargs: default_unquantized_gemm, ) with set_current_vllm_config(VllmConfig(device_config=device_config)): diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py b/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py index 52f1b6aa3dc..4f3a81491cd 100644 --- a/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py +++ b/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py @@ -1,12 +1,39 @@ import pytest import torch -from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import QwenEmbedRope +from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( + QwenEmbedRope, + _apply_qwen_image_rotary_emb, +) from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("noncontiguous", [False, True]) +def test_qwen_rotary_preserves_frequency_precision(dtype, noncontiguous): + from diffusers.models.transformers.transformer_qwenimage import apply_rotary_emb_qwen + + generator = torch.Generator().manual_seed(42) + x = torch.randn(2, 17, 3, 128, generator=generator).to(dtype) + if noncontiguous: + x = x.transpose(1, 2).contiguous().transpose(1, 2) + angles = torch.randn(17, 64, generator=generator) + freqs = torch.polar(torch.ones_like(angles), angles) + expected = apply_rotary_emb_qwen(x, freqs, use_real=False) + actual = _apply_qwen_image_rotary_emb(x, freqs) + assert actual.dtype == dtype + assert actual.shape == x.shape + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + if dtype == torch.bfloat16: + # A regression to early BF16 frequency rounding must be observable. + rounded_freqs = torch.complex(freqs.real.to(dtype).float(), freqs.imag.to(dtype).float()) + rounded = _apply_qwen_image_rotary_emb(x, rounded_freqs) + assert not torch.equal(actual, rounded) + + def test_txt_seq_lens_from_embeds_uses_padded_width_not_valid_token_count(): prompt_embeds = torch.zeros(2, 64, 8) prompt_embeds_mask = torch.zeros(2, 64, dtype=torch.bool) diff --git a/tests/diffusion/models/sana_video/test_transformer_sana_video.py b/tests/diffusion/models/sana_video/test_transformer_sana_video.py index 0bc2ddb8e41..9e6404b21ad 100644 --- a/tests/diffusion/models/sana_video/test_transformer_sana_video.py +++ b/tests/diffusion/models/sana_video/test_transformer_sana_video.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 os @@ -22,7 +22,7 @@ def _init_distributed(monkeypatch): monkeypatch.setattr( "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", - lambda: default_unquantized_gemm, + lambda *_args, **_kwargs: default_unquantized_gemm, ) os.environ.setdefault("MASTER_ADDR", "localhost") os.environ.setdefault("MASTER_PORT", "29501") diff --git a/tests/diffusion/models/sana_video/test_transformer_sana_video_sp.py b/tests/diffusion/models/sana_video/test_transformer_sana_video_sp.py index 3f4c620754e..6a0f0548e23 100644 --- a/tests/diffusion/models/sana_video/test_transformer_sana_video_sp.py +++ b/tests/diffusion/models/sana_video/test_transformer_sana_video_sp.py @@ -41,8 +41,8 @@ def tp1_group(): os.environ.setdefault("MASTER_ADDR", "localhost") os.environ.setdefault("MASTER_PORT", "29501") - init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://") - initialize_model_parallel() + init_distributed_environment(world_size=1, rank=0, local_rank=0, distributed_init_method="env://", backend="gloo") + initialize_model_parallel(backend="gloo") yield cleanup_dist_env_and_memory() @@ -54,7 +54,7 @@ def force_default_gemm(monkeypatch): monkeypatch.setattr( "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", - lambda: default_unquantized_gemm, + lambda *_args, **_kwargs: default_unquantized_gemm, ) @@ -482,8 +482,12 @@ def batch_isend_irecv(self, ops): def test_tiny_transformer_sp2_lockstep_matches_dense(tp1_group, force_default_gemm, mocker, task) -> None: """Two lockstep rank threads over uneven frame shards must reproduce the dense tiny-transformer output, and must actually have sharded the work.""" + from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend from vllm_omni.diffusion.models.sana_video import SanaVideoTransformer3DModel + # These CPU tensor tests must not probe a CUDA device during construction. + mocker.patch("vllm_omni.diffusion.attention.selector._cached_get_backend_cls", return_value=SDPABackend) + torch.manual_seed(3) model = SanaVideoTransformer3DModel(**_TINY_SP_CONFIG).eval() for _, param in sorted(model.named_parameters()): diff --git a/tests/diffusion/models/sana_video/test_transformer_sana_video_tp.py b/tests/diffusion/models/sana_video/test_transformer_sana_video_tp.py index 41ca717d62e..a44c901469f 100644 --- a/tests/diffusion/models/sana_video/test_transformer_sana_video_tp.py +++ b/tests/diffusion/models/sana_video/test_transformer_sana_video_tp.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 """TP unit tests for SANA-Video: distributed RMSNorm, sharded attention and sharded GLUMB temporal conv. @@ -51,7 +51,7 @@ def force_default_gemm(monkeypatch): monkeypatch.setattr( "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", - lambda: default_unquantized_gemm, + lambda *_args, **_kwargs: default_unquantized_gemm, ) diff --git a/tests/diffusion/models/t5_encoder/test_t5_encoder_tp.py b/tests/diffusion/models/t5_encoder/test_t5_encoder_tp.py index b36acc68aa6..cca73fb0895 100644 --- a/tests/diffusion/models/t5_encoder/test_t5_encoder_tp.py +++ b/tests/diffusion/models/t5_encoder/test_t5_encoder_tp.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 pytest import torch from transformers import T5Config @@ -54,6 +54,13 @@ def setup_tp_group(monkeypatch, mocker): lambda: 0, ) + # vLLM 0.29 added a CUDA-only fused embedding kernel that engages whenever + # tp_size > 1; these tests drive CPU tensors, so keep the portable masked path. + monkeypatch.setattr( + "vllm.model_executor.layers.vocab_parallel_embedding.current_platform.is_cuda", + lambda: False, + ) + # TP group mock_tp_group = mocker.MagicMock() mock_tp_group.world_size = 2 diff --git a/tests/diffusion/quantization/test_component_routing.py b/tests/diffusion/quantization/test_component_routing.py index 94341177911..6aebd6b60d7 100644 --- a/tests/diffusion/quantization/test_component_routing.py +++ b/tests/diffusion/quantization/test_component_routing.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for component routing for quantization.""" from unittest.mock import MagicMock @@ -372,7 +372,8 @@ def test_maps_compressed_tensors_component_and_default(self): "thinker.model": "thinker.language_model.model", "talker.model": "talker.language_model.model", } - ).get_unstacked_mapper() + # vLLM 0.29 renamed get_unstacked_mapper to get_rename_mapper. + ).get_rename_mapper() config.apply_vllm_mapper(mapper) diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index d6d17c18f02..20a2f0384f0 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/tests/diffusion/test_diffusion_scheduler.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 asyncio import queue @@ -10,8 +10,9 @@ import torch import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers from pytest_mock import MockerFixture -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec +from tests.helpers.kv_layout import build_kv_cache_tensor from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode @@ -104,7 +105,7 @@ def _initialize_paged_scheduler( ) config = KVCacheConfig( num_blocks=num_blocks, - kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=["layer0"])], + kv_cache_tensors=[build_kv_cache_tensor(spec, num_blocks, ["layer0"])], kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)], ) scheduler.initialize( diff --git a/tests/e2e/accuracy/test_diffusers_backend_similarity.py b/tests/e2e/accuracy/test_diffusers_backend_similarity.py index be3d7f84cb6..b12e9c1e725 100644 --- a/tests/e2e/accuracy/test_diffusers_backend_similarity.py +++ b/tests/e2e/accuracy/test_diffusers_backend_similarity.py @@ -164,6 +164,7 @@ def _run_diffusers_wan22_i2v(*, model: str, output_path: Path, conditioning_imag ) pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=FLOW_SHIFT) pipe.to("cuda") + _set_matched_attention_backend(pipe) _diffusers_dummy_run(pipe) diff --git a/tests/e2e/online_serving/test_dots_tts_expansion.py b/tests/e2e/online_serving/test_dots_tts_expansion.py index 34bc54b8e4b..e5df55a33fb 100644 --- a/tests/e2e/online_serving/test_dots_tts_expansion.py +++ b/tests/e2e/online_serving/test_dots_tts_expansion.py @@ -85,5 +85,7 @@ def test_text_only_nonstreaming_pcm_001(omni_server, openai_client) -> None: "voice": "default", "min_audio_bytes": _MIN_AUDIO_BYTES, "min_hnr_db": -2.0, + # dots.tts-soar emits native 48 kHz PCM, which has no rate header. + "expected_sample_rate": 48_000, } openai_client.send_audio_speech_request(request_config) diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index e6f7f8a8ba9..27677bd546b 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -1593,73 +1593,58 @@ def test_model_path_resolver_is_generic_and_model_owned(tmp_path): assert "model_path_resolver" not in engine_args -def test_build_stage0_input_processor_uses_omni_input_preprocessor(monkeypatch): +def test_build_stage0_input_processor_uses_omni_renderer_subclass(monkeypatch): + from vllm.renderers import BaseRenderer + import vllm_omni.engine.stage_init_utils as init_mod + from vllm_omni.inputs.preprocess import OmniRenderer, omni_renderer_cls - class DummyInputProcessor: - def __init__(self, vllm_config, renderer=None): - self.vllm_config = vllm_config - self.renderer = renderer or object() - self.input_preprocessor = None + class _Base(BaseRenderer): + def __init__(self, config, tokenizer): + self.config, self.tokenizer = config, tokenizer + + def render_messages(self, messages, params): # pragma: no cover - abstract stub + raise NotImplementedError + + config = types.SimpleNamespace(model_config=types.SimpleNamespace(try_get_generation_config=lambda: {})) + built = omni_renderer_cls(_Base)(config, "tok") + seen = {} - class DummyOmniInputPreprocessor: + class DummyInputProcessor: def __init__(self, vllm_config, renderer=None): - self.vllm_config = vllm_config + seen["renderer"] = renderer self.renderer = renderer monkeypatch.setattr(init_mod, "InputProcessor", DummyInputProcessor) - monkeypatch.setattr(init_mod, "OmniInputPreprocessor", DummyOmniInputPreprocessor) - - input_processor = build_stage0_input_processor( - types.SimpleNamespace(model_config=types.SimpleNamespace(try_get_generation_config=lambda: {})) - ) - - assert isinstance(input_processor.input_preprocessor, DummyOmniInputPreprocessor) - assert input_processor.input_preprocessor.renderer is input_processor.renderer + monkeypatch.setattr(init_mod, "build_omni_renderer", lambda cfg: built if cfg is config else None) + processor = build_stage0_input_processor(config) + assert seen["renderer"] is built + assert isinstance(processor.renderer, OmniRenderer) + assert isinstance(processor.renderer, _Base) + assert not hasattr(processor, "input_preprocessor") -def test_build_stage0_input_processor_does_not_resolve_tokenizer_when_skipped( - monkeypatch, -): +def test_build_stage0_input_processor_does_not_resolve_tokenizer_when_skipped(monkeypatch): import vllm_omni.engine.stage_init_utils as init_mod - token_only_renderer = object() - seen = {} + original = object() class DummyInputProcessor: def __init__(self, vllm_config, renderer=None): - if renderer is None: - raise AssertionError("skip_tokenizer_init must supply a renderer") - seen["renderer"] = renderer + assert renderer is original self.renderer = renderer - self.input_preprocessor = None - class DummyOmniInputPreprocessor: - def __init__(self, vllm_config, renderer=None): - seen["preprocessor_renderer"] = renderer + def _must_not_resolve(_cfg): + raise AssertionError("tokenizer must not be resolved when skip_tokenizer_init=True") monkeypatch.setattr(init_mod, "InputProcessor", DummyInputProcessor) - monkeypatch.setattr( - init_mod, - "_build_token_only_renderer", - lambda _config: token_only_renderer, - ) - monkeypatch.setattr( - init_mod, - "OmniInputPreprocessor", - DummyOmniInputPreprocessor, - ) - + monkeypatch.setattr(init_mod, "_build_token_only_renderer", lambda _: original) + monkeypatch.setattr(init_mod, "build_omni_renderer", _must_not_resolve) config = types.SimpleNamespace( - model_config=types.SimpleNamespace( - skip_tokenizer_init=True, - try_get_generation_config=lambda: {}, - ) + model_config=types.SimpleNamespace(skip_tokenizer_init=True, try_get_generation_config=lambda: {}) ) - build_stage0_input_processor(config) - - assert seen["renderer"] is token_only_renderer - assert seen["preprocessor_renderer"] is token_only_renderer + processor = build_stage0_input_processor(config) + assert processor.renderer is original def test_inject_kv_stage_info_infers_sender_tp_topology(): diff --git a/tests/entrypoints/openai/test_batch_chat_completions.py b/tests/entrypoints/openai/test_batch_chat_completions.py index 9fc72b2c116..c3866dbc81e 100644 --- a/tests/entrypoints/openai/test_batch_chat_completions.py +++ b/tests/entrypoints/openai/test_batch_chat_completions.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import asyncio from unittest.mock import AsyncMock @@ -12,7 +15,7 @@ ChatCompletionResponseChoice, ChatMessage, ) -from vllm.entrypoints.openai.engine.protocol import ErrorResponse, UsageInfo +from vllm.entrypoints.serve.engine.protocol import ErrorResponse, UsageInfo from vllm_omni.entrypoints.openai.batch_serving import OmniOpenAIServingChatBatch diff --git a/tests/entrypoints/openai_api/test_audio_format.py b/tests/entrypoints/openai_api/test_audio_format.py index a9bf76319bd..9c7eb9d8975 100644 --- a/tests/entrypoints/openai_api/test_audio_format.py +++ b/tests/entrypoints/openai_api/test_audio_format.py @@ -248,7 +248,7 @@ def test_pcm16_mapped_to_pcm(self, serving_chat): assert result == "pcm" def test_invalid_format_returns_error(self, serving_chat): - from vllm.entrypoints.openai.engine.protocol import ErrorResponse + from vllm.entrypoints.serve.engine.protocol import ErrorResponse request = self._make_request({"format": "aac", "voice": "alloy"}) result = serving_chat._resolve_audio_format(request) @@ -256,7 +256,7 @@ def test_invalid_format_returns_error(self, serving_chat): assert "aac" in result.error.message def test_all_supported_formats_accepted(self, serving_chat): - from vllm.entrypoints.openai.engine.protocol import ErrorResponse + from vllm.entrypoints.serve.engine.protocol import ErrorResponse for fmt in SUPPORTED_CHAT_AUDIO_FORMATS: request = self._make_request({"format": fmt, "voice": "alloy"}) diff --git a/tests/entrypoints/openai_api/test_modality_validation.py b/tests/entrypoints/openai_api/test_modality_validation.py index fde0c0ca1cd..91198b2e2f3 100644 --- a/tests/entrypoints/openai_api/test_modality_validation.py +++ b/tests/entrypoints/openai_api/test_modality_validation.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Regression tests for output modality validation in chat completions. Covers: @@ -13,7 +14,7 @@ import pytest from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest -from vllm.entrypoints.openai.engine.protocol import ErrorResponse +from vllm.entrypoints.serve.engine.protocol import ErrorResponse pytestmark = [pytest.mark.core_model, pytest.mark.cpu] diff --git a/tests/entrypoints/openai_api/test_serving_chat_metrics.py b/tests/entrypoints/openai_api/test_serving_chat_metrics.py index 2fdf9f6469c..6e796f480d6 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_metrics.py +++ b/tests/entrypoints/openai_api/test_serving_chat_metrics.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for OmniChatCompletionResponse/StreamResponse metrics field.""" from types import SimpleNamespace @@ -19,7 +20,7 @@ def test_omni_chat_completion_response_metrics(): """Test OmniChatCompletionResponse metrics field works correctly.""" - from vllm.entrypoints.openai.engine.protocol import UsageInfo + from vllm.entrypoints.serve.engine.protocol import UsageInfo from vllm_omni.entrypoints.openai.protocol.chat_completion import ( OmniChatCompletionResponse, diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 6a5ebcf7d3b..3922ab4a765 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -25,7 +25,7 @@ from fastapi.testclient import TestClient from pydantic import ValidationError from pytest_mock import MockerFixture -from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse +from vllm.entrypoints.serve.engine.protocol import ErrorInfo, ErrorResponse from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler diff --git a/tests/entrypoints/openai_api/test_video_stream_handler.py b/tests/entrypoints/openai_api/test_video_stream_handler.py index e382b5e3279..31717d0b728 100644 --- a/tests/entrypoints/openai_api/test_video_stream_handler.py +++ b/tests/entrypoints/openai_api/test_video_stream_handler.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 """Runtime behavior tests for VideoStreamHandler. Every test creates a mock WebSocket, drives handle_session through a @@ -281,7 +281,7 @@ async def test_query_with_no_frames(self): @pytest.mark.asyncio async def test_query_returns_error_response(self): """create_chat_completion returns ErrorResponse → error forwarded.""" - from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse + from vllm.entrypoints.serve.engine.protocol import ErrorInfo, ErrorResponse err = ErrorResponse( error=ErrorInfo( diff --git a/tests/entrypoints/test_stream_finish_reason.py b/tests/entrypoints/test_stream_finish_reason.py index 97ed2370bd1..87aeb787ebb 100644 --- a/tests/entrypoints/test_stream_finish_reason.py +++ b/tests/entrypoints/test_stream_finish_reason.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for multi-modal streaming finish_reason behavior (commit 44c799bc). Verifies that the /v1/chat/completions streaming endpoint emits exactly one @@ -411,7 +412,7 @@ async def result_generator(): @pytest.mark.asyncio async def test_audio_choice_error_response_is_not_iterated_as_choices(): """Any ErrorResponse from the audio path is skipped, not iterated.""" - from vllm.entrypoints.openai.engine.protocol import ErrorResponse + from vllm.entrypoints.serve.engine.protocol import ErrorResponse serving_chat = build_serving_chat() request = make_request(modalities=["text", "audio"]) diff --git a/tests/helpers/kv_layout.py b/tests/helpers/kv_layout.py new file mode 100644 index 00000000000..8351de8da1e --- /dev/null +++ b/tests/helpers/kv_layout.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Shared access to the production ``KVCacheTensor`` layout derivation. + +Tests deliberately go through the same code production uses +(``vllm_omni.diffusion.diffusion_kv.layout``) rather than recomputing strides, +so a regression in that derivation fails the suite instead of being masked by a +parallel test-only copy. +""" + +from types import SimpleNamespace + +from vllm.config import CacheConfig +from vllm.v1.kv_cache_layout import KVCacheLayout + +from vllm_omni.diffusion.diffusion_kv.layout import ( + build_kv_cache_tensor, + resolve_diffusion_kv_cache_layout, +) + +__all__ = ["build_kv_cache_tensor", "layout_for_backend"] + + +def layout_for_backend(attn_backend) -> KVCacheLayout: + """The physical layout ``attn_backend``'s block-stride declaration selects. + + Before 0.29 a backend's ``indexes_kv_by_block_stride()`` was copied onto + ``AttentionSpec``; asserting on the resolved layout is the surviving way to + check that declaration still reaches the cache. + """ + + vllm_config = SimpleNamespace(cache_config=CacheConfig()) + return resolve_diffusion_kv_cache_layout( + vllm_config, + indexes_kv_by_block_stride=attn_backend.indexes_kv_by_block_stride(), + ) diff --git a/tests/helpers/serving_chat.py b/tests/helpers/serving_chat.py index f1f2710b569..45add4b63bb 100644 --- a/tests/helpers/serving_chat.py +++ b/tests/helpers/serving_chat.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Shared scaffolding for OmniOpenAIServingChat unit tests. Provides factories for a minimal serving-chat instance, chat requests, and @@ -8,11 +9,11 @@ import json from unittest.mock import MagicMock +from vllm.entrypoints.generate.base.protocol import DeltaMessage from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionResponseStreamChoice, ) -from vllm.entrypoints.openai.engine.protocol import DeltaMessage from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.outputs import CompletionOutput, RequestOutput diff --git a/tests/inputs/test_mm_processor.py b/tests/inputs/test_mm_processor.py new file mode 100644 index 00000000000..813d7f436e2 --- /dev/null +++ b/tests/inputs/test_mm_processor.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from transformers import BatchFeature + +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize("counts", [{}, {"audio": 1, "image": 0}]) +@pytest.mark.parametrize("original_text", [None, "", "voice cloning text"]) +def test_custom_processor_bridge_preserves_prompt_kwargs_and_passthrough(counts, original_text): + class Processor(OmniMultiModalProcessor): + def _get_mm_fields_config(self, *args): + return {} + + def _get_prompt_updates(self, *args): + return [] + + processor = object.__new__(Processor) + processor.dummy_inputs = SimpleNamespace(get_dummy_text=Mock(return_value="dummy")) + processor._get_hf_mm_data = Mock(return_value=({"audios": ["wave"]}, {"embedding": "preserve"})) + processor._call_hf_processor = Mock(return_value=BatchFeature({"features": "processed"})) + items = Mock() + items.get_all_counts.return_value = counts + kwargs = {"sampling_rate": 16000} + if original_text is not None: + kwargs[processor._OMNI_PROMPT_TEXT_KEY] = original_text + before = dict(kwargs) + result = processor._apply_hf_processor_main(items, kwargs) + items.select.assert_called_once_with({k for k, v in counts.items() if v > 0}) + processor._call_hf_processor.assert_called_once_with( + "dummy" if original_text is None else original_text, + {"audios": ["wave"]}, + {"sampling_rate": 16000}, + {}, + ) + assert dict(result) == {"features": "processed", "embedding": "preserve"} + assert kwargs == before diff --git a/tests/inputs/test_preprocess.py b/tests/inputs/test_preprocess.py index db8817cf2d1..3fb48a7cce4 100644 --- a/tests/inputs/test_preprocess.py +++ b/tests/inputs/test_preprocess.py @@ -1,35 +1,157 @@ # SPDX-License-Identifier: Apache-2.0 -"""Unit tests for OmniInputPreprocessor._process_text routing.""" +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Exercise Omni routing through the active upstream rendering pipeline.""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock import pytest -from pytest_mock import MockerFixture +import torch +from vllm.inputs import tokens_input +from vllm.pooling_params import PoolingParams +from vllm.renderers import BaseRenderer +from vllm.renderers.params import TokenizeParams +from vllm.v1.engine.input_processor import InputProcessor + +from vllm_omni.inputs import preprocess as preprocess_mod +from vllm_omni.inputs.preprocess import OmniRenderer, build_omni_renderer, omni_renderer_cls pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -class TestProcessTextMmProcessorKwargsRouting: - """Presence-based routing: an explicit empty ``mm_processor_kwargs`` dict - on the prompt still routes through ``_process_multimodal``. Required for - AR-based image generation (e.g. GLM-Image t2i) where the HF processor - supplies the image-generation scaffold from its own defaults and no extra - kwargs are needed from the caller. - """ - - @pytest.fixture - def preprocessor(self, mocker: MockerFixture): - from vllm_omni.inputs.preprocess import OmniInputPreprocessor - - instance = object.__new__(OmniInputPreprocessor) - instance._process_multimodal = mocker.MagicMock(return_value={}) - instance._tokenize_prompt = mocker.MagicMock(return_value=[1, 2, 3]) - return instance - - def test_empty_mm_processor_kwargs_routes_to_multimodal(self, preprocessor): - preprocessor._process_text({"prompt": "hello", "mm_processor_kwargs": {}}) - assert preprocessor._process_multimodal.called - assert not preprocessor._tokenize_prompt.called - - def test_missing_mm_processor_kwargs_routes_to_tokenize(self, preprocessor): - preprocessor._process_text({"prompt": "hello"}) - assert preprocessor._tokenize_prompt.called - assert not preprocessor._process_multimodal.called +class _Renderer(BaseRenderer): + def __init__(self): + self.model_config = SimpleNamespace(is_encoder_decoder=False, enable_prompt_embeds=True) + self.tokenizer = None + self.default_cmpl_tok_params = TokenizeParams(max_total_tokens=128) + self._tokenize_prompt = lambda prompt, params: {**prompt, "prompt_token_ids": [1, 2, 3]} + self._tokenize_prompt_async = AsyncMock(side_effect=self._tokenize_prompt) + self._process_multimodal = Mock(return_value=tokens_input([1, 2, 3, 99])) + self._process_multimodal_async = AsyncMock(side_effect=self._process_multimodal) + + def render_messages(self, messages, params): + return messages, {"prompt": "hello"} + + +@pytest.fixture +def renderer(): + return omni_renderer_cls(_Renderer)() + + +@pytest.mark.parametrize("kwargs", [{}, {"target_h": 512, "target_w": 768}]) +@pytest.mark.parametrize("tokenized", [False, True]) +def test_process_inputs_routes_no_media_processor_kwargs(renderer, kwargs, tokenized): + # Keep real process_inputs/render_cmpl; stub unrelated config validation. + processor = object.__new__(InputProcessor) + processor.renderer = renderer + processor.model_config = renderer.model_config + processor.vllm_config = SimpleNamespace( + parallel_config=SimpleNamespace( + data_parallel_size=1, + data_parallel_size_local=1, + local_engines_only=False, + ) + ) + processor._validate_params = Mock() + processor._validate_lora = Mock() + processor._validate_model_inputs = Mock() + prompt: dict[str, object] = {"prompt_token_ids": [1, 2, 3]} if tokenized else {"prompt": "hello"} + prompt.update(mm_processor_kwargs=kwargs, cache_salt="salt") + request = processor.process_inputs("image-request", prompt, PoolingParams(), ("embed",)) + assert request.prompt_token_ids == [1, 2, 3, 99] + assert request.cache_salt == "salt" + renderer._process_multimodal.assert_called_once_with( + [1, 2, 3], + {}, + mm_processor_kwargs=kwargs, + mm_uuids=None, + skip_mm_cache=False, + ) + + +@pytest.mark.parametrize("async_mode", [False, True]) +@pytest.mark.parametrize("kind", ["text", "tokens", "media", "kwargs", "embeds"]) +def test_renderer_preserves_routes_extras_and_cache_policy(renderer, async_mode, kind): + prompt: dict[str, object] = { + "prompt": "hello", + "additional_information": {"speaker": 1}, + "model_intermediate_buffer": {"x": 2}, + } + if kind == "tokens": + prompt["prompt_token_ids"] = [1, 2, 3] + elif kind == "media": + prompt["multi_modal_data"] = {"image": "image"} + elif kind == "kwargs": + prompt["mm_processor_kwargs"] = {} + elif kind == "embeds": + prompt["prompt_embeds"] = torch.ones(3, 4) + if async_mode: + (result,) = asyncio.run(renderer.render_cmpl_async([prompt], skip_mm_cache=True)) + else: + (result,) = renderer.render_cmpl([prompt], skip_mm_cache=True) + assert result["additional_information"] == {"speaker": 1} + assert result["model_intermediate_buffer"] == {"x": 2} + assert result["prompt"] == "hello" + if kind in ("media", "kwargs"): + assert result["prompt_token_ids"] == [1, 2, 3, 99] + assert renderer._process_multimodal.call_args.kwargs["skip_mm_cache"] is True + else: + renderer._process_multimodal.assert_not_called() + if kind == "embeds": + torch.testing.assert_close(result["prompt_embeds"], prompt["prompt_embeds"]) + else: + assert result["prompt_token_ids"] == [1, 2, 3] + + +def test_omni_renderer_is_a_real_subclass_created_once(): + cls = omni_renderer_cls(_Renderer) + assert issubclass(cls, OmniRenderer) and issubclass(cls, _Renderer) + assert omni_renderer_cls(_Renderer) is cls + assert omni_renderer_cls(cls) is cls + assert cls.__name__ == "Omni_Renderer" + instance = cls() + assert isinstance(instance, _Renderer) + assert not hasattr(instance, "_renderer") + + +def test_concrete_renderer_overrides_are_not_suppressed(): + """The old proxy bound BaseRenderer.render_cmpl itself and hid subclass overrides.""" + + class _Custom(_Renderer): + def render_cmpl(self, prompts, *args, **kwargs): + return ["custom"] + + assert omni_renderer_cls(_Custom)().render_cmpl([{"prompt": "x"}]) == ["custom"] + + +class _CtorRenderer(BaseRenderer): + def __init__(self, config, tokenizer): + self.config = config + self.tokenizer = tokenizer + + def render_messages(self, messages, params): + raise NotImplementedError + + +def test_build_omni_renderer_with_explicit_class_skips_registry(monkeypatch): + monkeypatch.setattr(preprocess_mod, "cached_tokenizer_from_config", Mock(side_effect=AssertionError("resolved"))) + config = SimpleNamespace(model_config=SimpleNamespace()) + built = build_omni_renderer(config, renderer_cls=_CtorRenderer, tokenizer=None) + assert isinstance(built, OmniRenderer) and isinstance(built, _CtorRenderer) + assert built.config is config and built.tokenizer is None + + +def test_build_omni_renderer_resolves_like_upstream(monkeypatch): + config = SimpleNamespace(model_config=SimpleNamespace()) + tokenizer = object() + monkeypatch.setattr( + preprocess_mod, "cached_tokenizer_from_config", lambda mc: tokenizer if mc is config.model_config else None + ) + monkeypatch.setattr(preprocess_mod, "tokenizer_args_from_config", lambda mc: ("fake-mode", None)) + load_cls = Mock(return_value=_CtorRenderer) + monkeypatch.setattr(preprocess_mod.RENDERER_REGISTRY, "load_renderer_cls", load_cls) + built = build_omni_renderer(config) + load_cls.assert_called_once_with("fake-mode") + assert isinstance(built, OmniRenderer) and isinstance(built, _CtorRenderer) + assert built.tokenizer is tokenizer diff --git a/tests/model_executor/models/glm_image/test_glm_image_ar.py b/tests/model_executor/models/glm_image/test_glm_image_ar.py index 5ccc9fc62b0..ae10608397e 100644 --- a/tests/model_executor/models/glm_image/test_glm_image_ar.py +++ b/tests/model_executor/models/glm_image/test_glm_image_ar.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 """Unit tests for GLM-Image AR model: DataParser, processor, and M-RoPE.""" import importlib.util @@ -91,6 +91,41 @@ def _load_target_classes(): pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@pytest.mark.parametrize("num_images", [0, 1, 2]) +@pytest.mark.parametrize("preformatted", [False, True]) +def test_generation_prompt_ids_preserve_hf_target_scaffold(num_images, preformatted): + from types import SimpleNamespace + + from vllm.multimodal.parse import MultiModalDataItems + from vllm.multimodal.processing import BaseMultiModalProcessor, ProcessorInputs + + processor = object.__new__(GlmImageMultiModalProcessor) + # Only the public token attributes of GlmImageProcessor are needed; the + # expected suffix below is what HF's own formatter produces for 512x768. + hf_processor = SimpleNamespace(grid_bos_token="", grid_eos_token="", bos_token="") + tokenizer = SimpleNamespace(encode=lambda text, **kwargs: list(text.encode())) + processor.info = SimpleNamespace( + get_tokenizer=lambda: tokenizer, + get_hf_processor=lambda: hf_processor, + get_hf_config=lambda: _make_hf_config(), + ) + items = MagicMock(spec=MultiModalDataItems) + items.get_all_counts.return_value = {"image": num_images} + suffix = "16 24" + ("13 19" if not num_images else "") + "" + text_ids = list(b"a red bird") + expected = [167855] * num_images + text_ids + list(suffix.encode()) + inputs = ProcessorInputs( + prompt=expected if preformatted else text_ids, + mm_data_items=items, + hf_processor_mm_kwargs={"target_h": 512, "target_w": 768}, + ) + with patch.object(BaseMultiModalProcessor, "apply", side_effect=lambda inputs, timing: inputs) as upstream: + result = processor.apply(inputs, MagicMock()) + assert result.prompt == expected + assert inputs.prompt == (expected if preformatted else text_ids) + upstream.assert_called_once() + + def test_glm_image_ar_omits_hidden_pooler_payload(): assert GlmImageForConditionalGeneration.omni_pooler_payload_include_hidden is False diff --git a/tests/model_executor/models/minicpmo_4_5/test_audio_embeds_input.py b/tests/model_executor/models/minicpmo_4_5/test_audio_embeds_input.py index d25b58e4858..48d3bd9a339 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_audio_embeds_input.py +++ b/tests/model_executor/models/minicpmo_4_5/test_audio_embeds_input.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Regression tests for the MiniCPM-o 4.5 ``audio_embeds`` input path. Two defects made the OpenAI ``audio_embeds`` content part unusable with this model: @@ -71,6 +71,12 @@ def encode(self, text: str, add_special_tokens: bool = True) -> str: def decode(self, tokens: str) -> str: return tokens + def convert_tokens_to_ids(self, token: str) -> int: + # vLLM 0.29 prompt updates select placeholder positions by token id, + # so the processor resolves ```` once; the value is irrelevant here. + assert token == "" + return 0 + def _processing_info() -> MiniCPMO45OmniLLMProcessingInfo: processor = _FakeProcessor() diff --git a/tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py b/tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py index 3e8ee2a634f..eaa2ff2b329 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py +++ b/tests/model_executor/models/minicpmo_4_5/test_long_audio_processing.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Regression tests for MiniCPM-o 4.5 >30s audio preprocessing. Covers: @@ -43,6 +43,34 @@ _CHUNK_FRAMES = 3000 +@pytest.mark.parametrize("version", [(2, 5), (4, 5)]) +def test_prompt_batching_delegates_upstream_without_tts_kwargs(version): + from unittest.mock import MagicMock + + info = SimpleNamespace( + get_model_version=lambda: version, + get_hf_processor=MagicMock(return_value=object()), + ctx=SimpleNamespace(call_hf_processor=MagicMock()), + ) + info.ctx.call_hf_processor.side_effect = ( + [{"features": [1]}, {"features": [2]}] if version == (2, 5) else [{"features": [1, 2]}] + ) + kwargs = {"use_tts": True, "sampling_rate": 16000} + result = MiniCPMO45OmniLLMMultiModalProcessor._call_hf_processor_on_prompts( + SimpleNamespace(info=info), ["one", "two"], {"audios": [10, 20]}, kwargs, out_keys={"features"} + ) + assert result == {"features": [1, 2]} + assert kwargs == {"use_tts": True, "sampling_rate": 16000} + for call in info.ctx.call_hf_processor.call_args_list: + assert call.args[2] == {"sampling_rate": 16000} + payloads = [call.args[1] for call in info.ctx.call_hf_processor.call_args_list] + assert payloads == ( + [{"text": "one", "audios": 10}, {"text": "two", "audios": 20}] + if version == (2, 5) + else [{"text": ["one", "two"], "audios": [10, 20]}] + ) + + class TestMiniCPMOFieldConfig: def test_short_audios_stay_batched(self) -> None: # two <=30s audios: one feature entry + one 1-element lens per audio @@ -99,10 +127,10 @@ def _make_processor( processor.info = SimpleNamespace(audio_pattern="()") processor.data_parser = MiniCPMOMultiModalDataParser(target_sr=_SAMPLE_RATE) - def fake_base_call_hf_processor(prompts, mm_data, mm_kwargs, tok_kwargs, *, out_keys): + def fake_call_hf_processor_on_prompts(prompts, mm_data, mm_kwargs, *, out_keys): return {key: fake_hf_outputs[key] for key in out_keys} - processor._base_call_hf_processor = fake_base_call_hf_processor + processor._call_hf_processor_on_prompts = fake_call_hf_processor_on_prompts return processor @@ -120,7 +148,7 @@ def test_long_audio_unpads_per_chunk(self) -> None: } processor = _make_processor(fake_outputs) - result = processor.process_audios({"audios": [_audio_seconds(45)]}, {}, {}) + result = processor.process_audios({"audios": [_audio_seconds(45)]}, {}) features = result["audio_features"] assert [feat.shape[-1] for feat in features] == chunk_lens @@ -138,7 +166,7 @@ def test_multiple_long_audios(self) -> None: } processor = _make_processor(fake_outputs) - result = processor.process_audios({"audios": [_audio_seconds(42), _audio_seconds(38)]}, {}, {}) + result = processor.process_audios({"audios": [_audio_seconds(42), _audio_seconds(38)]}, {}) widths = [feat.shape[-1] for feat in result["audio_features"]] assert widths == [_CHUNK_FRAMES, 1200, _CHUNK_FRAMES, 800] @@ -151,7 +179,7 @@ def test_short_audio_unchanged(self) -> None: } processor = _make_processor(fake_outputs) - result = processor.process_audios({"audios": [_audio_seconds(9.8)]}, {}, {}) + result = processor.process_audios({"audios": [_audio_seconds(9.8)]}, {}) features = result["audio_features"] assert len(features) == 1 diff --git a/tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py b/tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py new file mode 100644 index 00000000000..9c6533c7644 --- /dev/null +++ b/tests/model_executor/models/omnivoice/test_processor_prompt_tokenization.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""OmniVoice processor: the prompt is tokenized exactly once per request. + +Upstream ``_cached_apply_hf_processor`` calls ``_apply_hf_processor_main`` for +every request, including text-only requests and requests whose reference audio +is already cached (empty ``mm_items``). These tests pin that only ``apply`` +touches the text tokenizer and that the media path only encodes the audio. +""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch +from vllm.multimodal.parse import MultiModalDataItems +from vllm.multimodal.processing import ProcessorInputs + +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor +from vllm_omni.model_executor.models.omnivoice.omnivoice import OmniVoiceMultiModalProcessor + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_PROMPT_IDS = [5, 6, 7] + + +def _make_processor(monkeypatch: pytest.MonkeyPatch) -> OmniVoiceMultiModalProcessor: + processor = object.__new__(OmniVoiceMultiModalProcessor) + tokenizer = SimpleNamespace(decode=Mock(return_value="hello world")) + processor.info = SimpleNamespace( + ctx=SimpleNamespace( + get_hf_config=lambda *_args, **_kwargs: SimpleNamespace(sample_rate=24000), + model_config=SimpleNamespace(model="fake-model"), + ), + get_tokenizer=lambda: tokenizer, + ) + processor.text_tokenizer = Mock(return_value=SimpleNamespace(input_ids=torch.tensor([_PROMPT_IDS]))) + processor.audio_tokenizer = SimpleNamespace(encode=Mock(return_value=torch.zeros(8, 4, dtype=torch.long))) + processor.feature_extractor = None + monkeypatch.setattr(processor, "_ensure_cached_runtime_components", lambda *_args, **_kwargs: None) + return processor + + +def _items(counts: dict[str, int], mm_data: dict[str, object]) -> SimpleNamespace: + return SimpleNamespace( + get_all_counts=lambda: counts, + select=lambda _keys: mm_data, + ) + + +def test_apply_tokenizes_the_prompt_once(monkeypatch: pytest.MonkeyPatch) -> None: + processor = _make_processor(monkeypatch) + captured: dict[str, ProcessorInputs] = {} + + def fake_parent_apply(self, inputs, timing_ctx): + captured["inputs"] = inputs + return "sentinel" + + monkeypatch.setattr(OmniMultiModalProcessor, "apply", fake_parent_apply) + inputs = ProcessorInputs( + prompt=[1, 2], + mm_data_items=MultiModalDataItems({}), + hf_processor_mm_kwargs={"ref_text": "reference"}, + ) + + assert processor.apply(inputs, timing_ctx=None) == "sentinel" + assert captured["inputs"].prompt == _PROMPT_IDS + assert processor.text_tokenizer.call_count == 1 + (text_prompt,), _ = processor.text_tokenizer.call_args + assert text_prompt.endswith("<|text_start|>reference hello world<|text_end|>") + assert captured["inputs"].hf_processor_mm_kwargs["ref_text"] == "reference" + + +def test_media_path_encodes_audio_without_retokenizing(monkeypatch: pytest.MonkeyPatch) -> None: + processor = _make_processor(monkeypatch) + audio = np.zeros(2400, dtype=np.float32) + monkeypatch.setattr(processor, "_get_hf_mm_data", lambda _items: ({"audios": [audio]}, {})) + + out = processor._apply_hf_processor_main(_items({"audio": 1}, {}), {"ref_text": "reference"}) + + assert tuple(out["ref_audio_tokens"].shape) == (8, 4) + assert list(out["ref_audio_len"]) == [4] + assert "input_ids" not in out + assert processor.audio_tokenizer.encode.call_count == 1 + assert processor.text_tokenizer.call_count == 0 + + +def test_media_path_with_no_audio_touches_no_tokenizer(monkeypatch: pytest.MonkeyPatch) -> None: + """Text-only requests and fully cached audio both reach this path with empty items.""" + processor = _make_processor(monkeypatch) + monkeypatch.setattr(processor, "_get_hf_mm_data", lambda _items: ({}, {})) + + out = processor._apply_hf_processor_main(_items({}, {}), {"ref_text": "reference"}) + + assert len(out.keys()) == 0 + assert processor.text_tokenizer.call_count == 0 + assert processor.audio_tokenizer.encode.call_count == 0 + + +def test_legacy_call_hf_processor_is_gone() -> None: + assert not hasattr(OmniVoiceMultiModalProcessor, "_call_hf_processor") diff --git a/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py b/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py index d88d918bc98..54efc99e47c 100644 --- a/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py +++ b/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py @@ -266,8 +266,6 @@ def test_cached_apply_hf_processor_expands_pair_miss_for_audio_hit_video_miss(): inputs = MagicMock() inputs.mm_data_items = mm_data_items inputs.hf_processor_mm_kwargs = {"use_audio_in_video": True} - inputs.tokenization_kwargs = {} - inputs.prompt = [1, 2, 3] inputs.get_mm_hashes = MagicMock(return_value={"video": ["v0"], "audio": ["a0"]}) timing_ctx = MagicMock() @@ -307,7 +305,7 @@ def capture_merge(cache_arg, **kwargs): ) fake_self._paired_cache_keys = Qwen2_5OmniThinkerMultiModalProcessor._paired_cache_keys fake_self.info.parse_mm_data = MagicMock(return_value=SimpleNamespace(name="missing")) - fake_self._apply_hf_processor_main = MagicMock(return_value=([9], {"video_grid_thw": object()}, False)) + fake_self._apply_hf_processor_main = MagicMock(return_value={"video_grid_thw": object()}) fake_self._get_mm_fields_config = MagicMock(return_value={}) fake_self._get_mm_prompt_updates = MagicMock(return_value={"video": [[object()]], "audio": [[object()]]}) fake_self._merge_mm_kwargs = capture_merge @@ -316,13 +314,12 @@ def capture_merge(cache_arg, **kwargs): "vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.MultiModalKwargsItems.from_hf_inputs", return_value={"video": [object()], "audio": [object()]}, ): - prompt_ids, mm_info, _ = Qwen2_5OmniThinkerMultiModalProcessor._cached_apply_hf_processor( + mm_info = Qwen2_5OmniThinkerMultiModalProcessor._cached_apply_hf_processor( fake_self, inputs, timing_ctx, ) - assert prompt_ids == [9] # Both modalities were reprocessed as one unit after pair-miss expansion. assert fake_self.info.parse_mm_data.call_args.args[0] == { "video": ["video-0"], @@ -623,10 +620,17 @@ def test_derive_audio_from_video_placeholders_only_pairs_true_videos(): def test_qwen3_processor_inherits_vllm_omni_per_video_helpers(): + from vllm.model_executor.models.qwen3_omni_moe_thinker import ( + Qwen3OmniMoeThinkerMultiModalProcessor as UpstreamQwen3Processor, + ) + assert issubclass( Qwen3OmniMoeThinkerMultiModalProcessor, Qwen2_5OmniThinkerMultiModalProcessor, ) + assert Qwen3OmniMoeThinkerMultiModalProcessor._get_hf_mm_data is UpstreamQwen3Processor._get_hf_mm_data + mro = Qwen3OmniMoeThinkerMultiModalProcessor.__mro__ + assert mro.index(Qwen2_5OmniThinkerMultiModalProcessor) < mro.index(UpstreamQwen3Processor) def test_qwen3_prompt_updates_prefer_hf_second_per_grid_ts_from_outputs(): diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py index 75d81eed548..4ff6d5096a7 100644 --- a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py @@ -113,7 +113,9 @@ def _save_peft_adapter(tmp_path: Path) -> Path: return adapter_dir -def test_peft_expert_lora_loads_into_thinker(tmp_path: Path): +def test_peft_expert_lora_loads_into_thinker(tmp_path: Path, monkeypatch): + # This CPU harness tests weight mapping, not CUDA pinned-memory allocation. + monkeypatch.setattr("vllm.lora.lora_model.PIN_MEMORY", False) adapter_dir = _save_peft_adapter(tmp_path) peft_helper = PEFTHelper.from_local_dir(str(adapter_dir), max_position_embeddings=128) lora_model = LoRAModel.from_local_checkpoint( diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py index 45b40d02652..64f9ac8de33 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.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 from unittest.mock import patch @@ -157,7 +157,10 @@ class _FakeAutoWeightsLoader: def __init__(self, *_: object, **__: object): pass - def load_weights(self, _weights: object) -> set[str]: + def load_weights(self, _weights: object, *, mapper: object | None = None) -> set[str]: + # vLLM 0.29 dropped skip_prefixes from AutoWeightsLoader; the model + # now passes the stage exclusions as a WeightsMapper instead. + del mapper return {"decoder.fake_weight"} with ( diff --git a/tests/model_executor/models/test_weight_loader.py b/tests/model_executor/models/test_weight_loader.py new file mode 100644 index 00000000000..34d1cd28a3a --- /dev/null +++ b/tests/model_executor/models/test_weight_loader.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch +from torch import nn +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _TinyModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.kept = nn.Linear(1, 1, bias=False) + self.skipped = nn.Linear(1, 1, bias=False) + self.rotary_embed = nn.Module() + self.rotary_embed.register_parameter("inv_freq", nn.Parameter(torch.zeros(1))) + + +def test_upstream_mapper_preserves_stage_exclusions_after_renaming() -> None: + model = _TinyModel() + model.skipped.weight.data.zero_() + + loaded = AutoWeightsLoader(model).load_weights( + [ + ("checkpoint.kept.weight", torch.ones(1, 1)), + ("checkpoint.skipped.weight", torch.full((1, 1), 2.0)), + ("rotary_embed.inv_freq", torch.ones(1)), + ], + mapper=WeightsMapper(orig_to_new_prefix={"checkpoint.": ""}) + | WeightsMapper(orig_to_new_prefix={"skipped.": None}, orig_to_new_substr={"rotary_embed.inv_freq": None}), + ) + + assert loaded == {"kept.weight"} + torch.testing.assert_close(model.kept.weight, torch.ones(1, 1)) + torch.testing.assert_close(model.skipped.weight, torch.zeros(1, 1)) + torch.testing.assert_close(model.rotary_embed.inv_freq, torch.zeros(1)) diff --git a/tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py b/tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py index 0e9f23de60c..985c6411453 100644 --- a/tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py +++ b/tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for the NVFP4 W4A4 weight_scale NaN clamp installed by ``vllm_omni.patch``. @@ -198,7 +198,7 @@ def test_reload_does_not_nest_wrapper(): print("NESTED" if nested else "FLAT") """ ) - result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=120) + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=600) if result.stdout.strip().startswith("SKIP:"): pytest.skip(result.stdout.strip()) assert "FLAT" in result.stdout, ( @@ -265,9 +265,12 @@ def test_env_var_escape_hatch_disables_install(env_value, acceptable): ) env = os.environ.copy() env["VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP"] = env_value - # 120s rather than 60s: cold CI containers can take 30-50s just to - # import vllm_omni → vllm → torch before the subprocess does any work. - result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=120, env=env) + # Generous timeout: cold CI containers can take 30-50s just to import + # vllm_omni → vllm → torch, and heavily loaded shared machines have + # been observed taking >2min for the same import under disk I/O + # contention. The subprocess does a single import + flag print, so 600s + # only bounds pathological hangs, not normal runtime. + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=600, env=env) if result.stdout.strip().startswith("SKIP:"): pytest.skip("modelopt not available in subprocess") # Exact token match on the final stdout line — not a substring test, since diff --git a/tests/model_executor/test_upstream_prompt_update_api.py b/tests/model_executor/test_upstream_prompt_update_api.py new file mode 100644 index 00000000000..da31980c2bf --- /dev/null +++ b/tests/model_executor/test_upstream_prompt_update_api.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Guard that Omni only calls prompt-update helpers upstream still provides. + +vLLM 0.29 removed ``PromptUpdateDetails.select_text`` in favour of the +token-oriented ``select_token_id`` / ``select_token_ids``. Three MiniCPM-o call +sites survived the rebase and only failed at engine-core startup, where the +traceback surfaced as an opaque "Engine core initialization failed" across eight +CI jobs. + +These helpers are referenced from replacement closures that no unit test drives, +so a static check is what actually catches the next removal. It resolves the +attribute against the installed class rather than a hard-coded list, so it +follows upstream automatically. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest +from vllm.multimodal.processing.processor import PromptUpdateDetails + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_OMNI_ROOT = Path(__file__).resolve().parents[2] / "vllm_omni" +_GUARDED = "PromptUpdateDetails" + + +def _referenced_attributes() -> dict[str, set[str]]: + """Map ``PromptUpdateDetails.`` -> source files referencing it.""" + found: dict[str, set[str]] = {} + for path in _OMNI_ROOT.rglob("*.py"): + try: + tree = ast.parse(path.read_text(encoding="utf-8")) + except (SyntaxError, UnicodeDecodeError): # pragma: no cover - defensive + continue + for node in ast.walk(tree): + if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name) and node.value.id == _GUARDED: + found.setdefault(node.attr, set()).add(str(path.relative_to(_OMNI_ROOT.parent))) + return found + + +def test_omni_only_uses_prompt_update_helpers_that_exist_upstream() -> None: + referenced = _referenced_attributes() + assert referenced, "expected Omni to reference PromptUpdateDetails somewhere" + + missing = {attr: sorted(files) for attr, files in referenced.items() if not hasattr(PromptUpdateDetails, attr)} + + assert not missing, ( + "Omni calls PromptUpdateDetails helpers that the installed vLLM does not " + f"provide: {missing}. Upstream moved prompt updates to token ids; use " + "select_token_id/select_token_ids with the encoded text." + ) + + +def test_token_oriented_helpers_are_the_ones_upstream_exposes() -> None: + """Pin the replacements, so their removal fails here and not at startup.""" + assert hasattr(PromptUpdateDetails, "select_token_id") + assert hasattr(PromptUpdateDetails, "select_token_ids") + assert not hasattr(PromptUpdateDetails, "select_text"), ( + "select_text is back upstream; re-check the MiniCPM-o replacements." + ) diff --git a/tests/platforms/test_platform_detection.py b/tests/platforms/test_platform_detection.py new file mode 100644 index 00000000000..6673a2376bc --- /dev/null +++ b/tests/platforms/test_platform_detection.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import logging + +import pytest + +import vllm_omni.platforms as platforms + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_failing_oot_plugin_is_logged_and_falls_back(monkeypatch, caplog): + plugin_name = "raising_oot_plugin" + + def raising_plugin(): + raise RuntimeError("injected platform detection failure") + + monkeypatch.setattr(platforms, "builtin_omni_platform_plugins", {}) + monkeypatch.setattr( + platforms, + "load_omni_plugins_by_group", + lambda _group: {plugin_name: raising_plugin}, + ) + + target_logger = logging.getLogger(platforms.__name__) + target_logger.addHandler(caplog.handler) + previous_level = target_logger.level + target_logger.setLevel(logging.DEBUG) + try: + platform_cls = platforms.resolve_current_omni_platform_cls_qualname() + finally: + target_logger.removeHandler(caplog.handler) + target_logger.setLevel(previous_level) + + assert platform_cls == "vllm_omni.platforms.interface.UnspecifiedOmniPlatform" + assert plugin_name in caplog.text + assert "RuntimeError: injected platform detection failure" in caplog.text diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index 10f0853e088..3313e691f1f 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -2975,7 +2975,9 @@ def measured_ttft(output: RequestFuncOutput) -> float | None: "duration": benchmark_duration, "completed": metrics.completed, "total_input_tokens": metrics.total_input, + "total_input_sequences": metrics.total_input_sequences, "request_throughput": metrics.request_throughput, + "input_sequence_throughput": metrics.input_sequence_throughput, "total_token_throughput": metrics.total_token_throughput, "input_lens": [output.prompt_len for output in outputs], "errors": [output.error for output in outputs], diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index aa71245a057..9123b794d7c 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -107,6 +107,17 @@ class OmniSchedulerMixin: def _init_omni_io_scheduling_state(self) -> None: """Initialize scheduler state shared by AR and generation stages.""" + # Every omni worker forces the v1 model runner (the v2 runner carries no + # omni hooks), so the scheduler has to agree or the two disagree about + # what the SchedulerOutput carries. vLLM 0.29 defaults this to v2 where + # 0.28 did not, so the scheduler took its v2 fast path -- the one that + # deliberately omits resumed-request token ids -- while the v1 runner + # still read them, and resuming a request raised ``KeyError: `` + # in ``_update_states``. + if getattr(self, "use_v2_model_runner", False): + logger.warning("OMNI scheduler forces v1 model runner for omni hooks.") + self.use_v2_model_runner = False + model_config = self.vllm_config.model_config self.chunk_transfer_adapter = ( OmniChunkTransferAdapter(self.vllm_config) if getattr(model_config, "async_chunk", False) else None diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index 5b5b8a32802..e903960d6d7 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -121,6 +121,13 @@ def indexes_kv_by_block_stride(cls) -> bool: when it unifies cache layouts across layers. Dense diffusion backends conservatively keep the default ``False``; a paged backend should override this only when its kernel actually follows that layout. + + Since vLLM 0.29 this no longer rides on ``AttentionSpec``. It selects + the block-outermost physical layout pinned on + ``CacheConfig.kv_cache_layout`` -- see + ``vllm_omni.diffusion.diffusion_kv.layout`` -- whose + ``is_block_outermost`` is what upstream reads when it sizes + ``KVCacheTensor`` regions. """ return False diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index ac541808b3e..4b8ca8f4980 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -24,6 +24,7 @@ from vllm_omni.diffusion.attention.selector import get_attn_backend_for_role from vllm_omni.diffusion.config import get_current_diffusion_config_or_none from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.diffusion_kv.layout import assert_backend_layout_supported from vllm_omni.diffusion.distributed.parallel_state import get_sp_group from vllm_omni.diffusion.forward_context import ( get_forward_context, @@ -261,15 +262,19 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: return None 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. + # upstream vLLM's attention-spec collector. vLLM 0.29 moved the + # block-stride decision off the spec and onto the single physical + # ``CacheConfig.kv_cache_layout`` resolved before the cache is built, + # so the backend's preference is enforced there (see + # ``vllm_omni.diffusion.diffusion_kv.initialization``) rather than + # carried per layer. with set_current_vllm_config(vllm_config): - indexes_kv_by_block_stride = self.attn_backend.indexes_kv_by_block_stride() + assert_backend_layout_supported(vllm_config, self.attn_backend) return FullAttentionSpec( block_size=vllm_config.cache_config.block_size, num_kv_heads=self.num_kv_heads, head_size=self.head_size, dtype=dtype, - indexes_kv_by_block_stride=indexes_kv_by_block_stride, non_causal=not self.causal, ) diff --git a/vllm_omni/diffusion/diffusion_kv/initialization.py b/vllm_omni/diffusion/diffusion_kv/initialization.py index 31aaaa46e37..d1a79d93467 100644 --- a/vllm_omni/diffusion/diffusion_kv/initialization.py +++ b/vllm_omni/diffusion/diffusion_kv/initialization.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 __future__ import annotations @@ -16,6 +16,7 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.diffusion_kv.layout import resolve_diffusion_kv_cache_layout from vllm_omni.diffusion.vllm_config import create_diffusion_vllm_config from vllm_omni.platforms import current_omni_platform @@ -28,14 +29,35 @@ def build_native_kv_cache_configs( vllm_config: VllmConfig, worker_specs: list[dict[str, KVCacheSpec]], available_memory: list[int], + *, + indexes_kv_by_block_stride: bool = False, ) -> tuple[list[KVCacheConfig], KVCacheConfig]: - """Build rank-local and Scheduler-native configs using vLLM utilities.""" + """Build rank-local and Scheduler-native configs using vLLM utilities. + + ``indexes_kv_by_block_stride`` carries the native backend requirement that + 0.29 removed from ``AttentionSpec``. Every diffusion backend in tree + declares ``False``, so the default reproduces current behaviour; a caller + that learns otherwise from its workers passes ``True`` here and the + block-outermost layout propagates to every rank-local config. + """ from vllm.v1.core import single_type_kv_cache_manager + # 0.29 sizes every KVCacheTensor region from the resolved physical layout, + # so it must be pinned before get_kv_cache_configs. Idempotent, so callers + # that already resolved it (the control plane) are unaffected. + layout = resolve_diffusion_kv_cache_layout( + vllm_config, + indexes_kv_by_block_stride=indexes_kv_by_block_stride, + ) single_type_kv_cache_manager.register_all_kvcache_specs(vllm_config) worker_configs = get_kv_cache_configs(vllm_config, worker_specs, available_memory) - return worker_configs, generate_scheduler_kv_cache_config(worker_configs) + scheduler_config = generate_scheduler_kv_cache_config(worker_configs) + # Workers only receive a KVCacheConfig, so the resolved layout travels on it + # and is adopted into each worker's own CacheConfig before allocation. + for config in (*worker_configs, scheduler_config): + config.kv_cache_layout = layout.name + return worker_configs, scheduler_config def initialize_diffusion_kv_control_plane( diff --git a/vllm_omni/diffusion/diffusion_kv/layout.py b/vllm_omni/diffusion/diffusion_kv/layout.py new file mode 100644 index 00000000000..960df395ded --- /dev/null +++ b/vllm_omni/diffusion/diffusion_kv/layout.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Physical KV cache layout selection for the native Diffusion paged cache. + +vLLM 0.29 removed ``AttentionSpec.indexes_kv_by_block_stride`` and replaced it +with a single physical layout per model, resolved once and recorded on +``CacheConfig.kv_cache_layout``. ``KVCacheLayout.is_block_outermost`` is the +direct successor of the old per-spec flag: upstream derives +``interleaved_block_stride`` from it when it lays out ``KVCacheTensor`` regions +(``vllm/v1/core/kv_cache_utils.py``). + +Diffusion never reaches vLLM's engine-core resolution. Its backends subclass +Omni's own ``AttentionBackend`` ABC, which upstream's +``get_supported_kv_cache_layouts`` never sees, so the layout is pinned here +before ``get_kv_cache_configs`` runs. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING + +from vllm.v1.kv_cache_interface import KVCacheSpec, KVCacheTensor, compute_layout_strides +from vllm.v1.kv_cache_layout import KVCacheLayout + +if TYPE_CHECKING: + from vllm.config import VllmConfig + +# Layer-outermost, matching the head of upstream's ``_DEFAULT_LAYOUT_PREFERENCE`` +# and reproducing the ``indexes_kv_by_block_stride=False`` that every dense +# diffusion backend declared before 0.29. +DEFAULT_DIFFUSION_KV_CACHE_LAYOUT = KVCacheLayout.LBNHC +# Block-outermost counterpart, for a paged backend whose kernel really does read +# K/V pages by the runtime block stride. +BLOCK_STRIDE_DIFFUSION_KV_CACHE_LAYOUT = KVCacheLayout.BLNHC + + +def resolve_diffusion_kv_cache_layout( + vllm_config: VllmConfig, + *, + indexes_kv_by_block_stride: bool = False, +) -> KVCacheLayout: + """Pin the physical KV layout for the Diffusion cache and return it. + + A layout already present on the config wins, mirroring upstream's + ``resolve_kv_cache_layout``; it is validated against the requirement rather + than silently overwritten, because an inconsistent layout reads K/V + from the wrong offsets instead of raising. + """ + + required_block_outermost = bool(indexes_kv_by_block_stride) + cache_config = getattr(vllm_config, "cache_config", None) + if cache_config is None: + # A config double with nowhere to record the choice; report what the + # requirement implies without pretending it was pinned. + return BLOCK_STRIDE_DIFFUSION_KV_CACHE_LAYOUT if required_block_outermost else DEFAULT_DIFFUSION_KV_CACHE_LAYOUT + + current = getattr(cache_config, "kv_cache_layout", None) + if current is not None: + layout = KVCacheLayout[current] + if layout.is_block_outermost != required_block_outermost: + raise ValueError( + "Diffusion KV cache layout " + f"{layout.name} (is_block_outermost=" + f"{layout.is_block_outermost}) contradicts the attention " + "backend, which requires is_block_outermost=" + f"{required_block_outermost}." + ) + return layout + + layout = BLOCK_STRIDE_DIFFUSION_KV_CACHE_LAYOUT if required_block_outermost else DEFAULT_DIFFUSION_KV_CACHE_LAYOUT + cache_config.kv_cache_layout = layout.name + return layout + + +def adopt_kv_cache_layout(vllm_config: VllmConfig, kv_cache_config) -> KVCacheLayout: + """Adopt the control plane's resolved layout inside a worker. + + The engine core resolves the layout once and every worker adopts it before + allocating (``CacheConfig.kv_cache_layout``). Diffusion workers receive only + a ``KVCacheConfig``, so the name rides along on ``kv_cache_config`` and is + copied onto the worker's own config here. Without this, anything calling + ``get_resolved_kv_cache_layout`` in a worker -- ``init_kv_cache`` does -- + raises. + """ + + name = getattr(kv_cache_config, "kv_cache_layout", None) + if name is None: + # Older configs, or a standalone harness: fall back to the same default + # the control plane pins so both sides still agree. + return resolve_diffusion_kv_cache_layout(vllm_config) + + cache_config = getattr(vllm_config, "cache_config", None) + if cache_config is None: + return KVCacheLayout[name] + + current = getattr(cache_config, "kv_cache_layout", None) + if current is None: + cache_config.kv_cache_layout = name + elif current != name: + raise ValueError( + f"Worker KV cache layout {current!r} disagrees with the layout resolved by the control plane ({name!r})." + ) + return KVCacheLayout[name] + + +def build_kv_cache_tensor( + spec: KVCacheSpec, + num_blocks: int, + layers: Sequence[str], + *, + layout: KVCacheLayout | None = None, + offset: int = 0, +) -> KVCacheTensor: + """A ``KVCacheTensor`` covering ``layers`` for ``num_blocks`` blocks. + + Before 0.29 a tensor named the layers sharing it + (``KVCacheTensor(size=..., shared_by=[...])``); 0.29 replaced that with an + explicit layout, so layer ``l``'s block ``b`` starts at + ``offset + l * layer_stride + b * block_stride``. Note this lays the layers + out at distinct offsets rather than aliasing them, which is what upstream + now does for a single cache group. + + The strides come from upstream's own ``compute_layout_strides``, called the + way ``vllm.v1.core.kv_cache_utils`` calls it, because hand-computed strides + read K/V from the wrong offsets silently instead of raising. + """ + + layout = layout or DEFAULT_DIFFUSION_KV_CACHE_LAYOUT + layer_names = list(layers) + # Mirrors kv_cache_utils: the block stride is pinned up front only for + # block-outermost layouts, where every block repeats the same packing. + bytes_per_block = spec.page_size_bytes * len(layer_names) + interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None + layer_stride, block_stride, _, _, _ = compute_layout_strides( + spec, + num_blocks, + len(layer_names), + layout, + fixed_strides=(None, interleaved_block_stride, None, None, None), + ) + return KVCacheTensor( + size=bytes_per_block * num_blocks, + layers=layer_names, + layer_stride=layer_stride, + block_stride=block_stride, + offset=offset, + ) + + +def assert_backend_layout_supported(vllm_config: VllmConfig, attn_backend: type | None) -> None: + """Fail loudly when a backend's block-stride need contradicts the layout. + + Skipped while the layout is still unresolved -- worker processes collect + specs before the control plane hands them the resolved name -- so this + enforces the contract wherever it is knowable and never crashes a rank that + has simply not been told yet. + """ + + if attn_backend is None: + return + # Spec discovery runs against lightweight config doubles in tests and against + # worker configs that have not been handed the resolved name yet; in both + # cases there is simply nothing to check. + cache_config = getattr(vllm_config, "cache_config", None) + if cache_config is None or getattr(cache_config, "kv_cache_layout", None) is None: + return + indexes_by_block_stride = getattr(attn_backend, "indexes_kv_by_block_stride", None) + if indexes_by_block_stride is None: + return + resolve_diffusion_kv_cache_layout( + vllm_config, + indexes_kv_by_block_stride=indexes_by_block_stride(), + ) diff --git a/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py b/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py index c5ab3fb1610..20f4db1653a 100644 --- a/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py +++ b/vllm_omni/diffusion/diffusion_kv/model_runner_backend.py @@ -18,6 +18,10 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode, is_scheduler_paged_kv_mode +from vllm_omni.diffusion.diffusion_kv.layout import ( + adopt_kv_cache_layout, + assert_backend_layout_supported, +) from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import ( DiffusionPagedAttentionAdapter, @@ -285,13 +289,16 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: layers=self._kv_cache_layer_adapters, resolve_row=self._resolve_paged_attention_row, ) + # vLLM 0.29 reads the resolved physical layout when allocating, + # and dropped attn_groups/cache_dtype from the signature. + adopt_kv_cache_layout(self.vllm_config, kv_cache_config) + for layer_adapter in self._kv_cache_layer_adapters.values(): + assert_backend_layout_supported(self.vllm_config, getattr(layer_adapter, "attn_backend", None)) init_kv_cache( kv_caches, self.vllm_config.compilation_config.static_forward_context, kv_cache_config, - attn_groups, self.device, - self.vllm_config.cache_config.cache_dtype, kernel_block_sizes, self.vllm_config, ) diff --git a/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py b/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py index 6a786266957..a76a5434c50 100644 --- a/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py +++ b/vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py @@ -105,10 +105,12 @@ def __init__( attn_backend, ulysses_degree=ulysses_degree, ) + # vLLM 0.29 carries the block-stride decision on the resolved physical + # ``CacheConfig.kv_cache_layout`` instead of the spec; the backend's + # requirement is enforced against that layout rather than copied here. canonical_spec = replace( spec, num_kv_heads=num_kv_heads, - indexes_kv_by_block_stride=attn_backend.indexes_kv_by_block_stride(), ) self.layer_name = layer_name self.spec = canonical_spec diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py index 59d2c9a621d..37b2ed83582 100644 --- a/vllm_omni/diffusion/layers/rope.py +++ b/vllm_omni/diffusion/layers/rope.py @@ -1,3 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from importlib import import_module from importlib.util import find_spec import torch @@ -98,6 +102,7 @@ def __init__(self, is_neox_style: bool = False, half_head_dim: bool = True) -> N self.interleaved = not is_neox_style self.half_head_dim = half_head_dim self.apply_rotary_emb_flash_attn = None + self.apply_rotary_emb_vllm_flash_attn = None self.has_mindie = False # ``find_spec("flash_attn")`` is True as long as *any* package publishes # the ``flash_attn`` namespace — including ``flash-attn-4``, which ships @@ -111,6 +116,17 @@ def __init__(self, is_neox_style: bool = False, half_head_dim: bool = True) -> N self.apply_rotary_emb_flash_attn = apply_rotary except ImportError: pass + # CUDA wheels vendor this kernel under vllm_flash_attn, while source + # and CPU-only installations may not contain the generated package. + # Resolve it lazily and retain the native implementation as a portable + # fallback for those environments. + if current_omni_platform.is_cuda(): + try: + self.apply_rotary_emb_vllm_flash_attn = import_module( + "vllm.vllm_flash_attn.layers.rotary" + ).apply_rotary_emb + except ImportError: + pass if find_spec("mindiesd") is not None: self.has_mindie = True @@ -135,12 +151,18 @@ def forward_cuda( cos: torch.Tensor, sin: torch.Tensor, ) -> torch.Tensor: - from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb - cos, sin = self._prepare_half_head_dim_cos_sin(cos, sin) + if self.apply_rotary_emb_vllm_flash_attn is None: + return apply_rotary_emb_torch( + x, + cos, + sin, + interleaved=self.interleaved, + ) + x, squeezed = _ensure_batch_dim(x) - output = apply_rotary_emb( + output = self.apply_rotary_emb_vllm_flash_attn( x, cos, sin, @@ -231,13 +253,14 @@ def forward_cuda( cos: torch.Tensor, sin: torch.Tensor, ) -> torch.Tensor: - from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb + if self.apply_rotary_emb_vllm_flash_attn is None: + return self.forward_native(x, cos, sin) if cos.dim() > 2: cos = cos.reshape(-1, cos.shape[-1]) sin = sin.reshape(-1, sin.shape[-1]) - return apply_rotary_emb( + return self.apply_rotary_emb_vllm_flash_attn( x, cos, sin, diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py index 6598525da38..fcfd48b4f35 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py @@ -454,11 +454,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Note: guidance_emb and timestep_r_emb are no longer skipped # to support HunyuanImage-3.0-Distil and MeanFlow distilled models - loader = AutoWeightsLoader( - self, - skip_prefixes=skip_prefixes, + loader = AutoWeightsLoader(self) + return loader.load_weights( + weights, mapper=WeightsMapper(orig_to_new_prefix={name: None for name in (skip_prefixes or ())}) ) - return loader.load_weights(weights) def prepare_seed(self, seed=None, batch_size=1): # random seed diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py b/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py index 6b0433bbee3..7228fa6552e 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py @@ -24,7 +24,7 @@ tensor_model_parallel_all_reduce, ) from vllm.logger import init_logger -from vllm.model_executor.layers.batch_invariant import linear_batch_invariant +from vllm.model_executor.determinism.batch_invariant import linear_batch_invariant from vllm.model_executor.layers.linear import ( ColumnParallelLinear, QKVParallelLinear, diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index c82ace9c4a7..696728b9043 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -50,6 +50,17 @@ logger = init_logger(__name__) +def _apply_qwen_image_rotary_emb(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """Rotate interleaved pairs before rounding back to the activation dtype. + + Qwen-Image's reference uses complex FP32 multiplication. Rounding the + frequencies to BF16 before rotation loses positional precision; those + errors accumulate across the denoising steps. + """ + paired = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2)) + return torch.view_as_real(paired * freqs.unsqueeze(1)).flatten(3).to(x.dtype) + + def _normalize_qwen_image_weight_name(name: str) -> str: name = name.removeprefix("transformer.") if ".to_out.0." in name: @@ -642,15 +653,23 @@ def forward( txt_query = self.norm_added_q(txt_query) txt_key = self.norm_added_k(txt_key) - img_cos = torch.real(vid_freqs).to(img_query.dtype) - img_sin = torch.imag(vid_freqs).to(img_query.dtype) - txt_cos = torch.real(txt_freqs).to(txt_query.dtype) - txt_sin = torch.imag(txt_freqs).to(txt_query.dtype) - - img_query = self.rope(img_query, img_cos, img_sin) - img_key = self.rope(img_key, img_cos, img_sin) - txt_query = self.rope(txt_query, txt_cos, txt_sin) - txt_key = self.rope(txt_key, txt_cos, txt_sin) + if img_query.device.type == "cuda": + img_query = _apply_qwen_image_rotary_emb(img_query, vid_freqs) + img_key = _apply_qwen_image_rotary_emb(img_key, vid_freqs) + txt_query = _apply_qwen_image_rotary_emb(txt_query, txt_freqs) + txt_key = _apply_qwen_image_rotary_emb(txt_key, txt_freqs) + else: + # Retain the platform-specific kernels on other accelerators, + # which may not support complex tensors. + img_cos = torch.real(vid_freqs).to(img_query.dtype) + img_sin = torch.imag(vid_freqs).to(img_query.dtype) + txt_cos = torch.real(txt_freqs).to(txt_query.dtype) + txt_sin = torch.imag(txt_freqs).to(txt_query.dtype) + + img_query = self.rope(img_query, img_cos, img_sin) + img_key = self.rope(img_key, img_cos, img_sin) + txt_query = self.rope(txt_query, txt_cos, txt_sin) + txt_key = self.rope(txt_key, txt_cos, txt_sin) seq_len_txt = encoder_hidden_states.shape[1] joint_query = torch.cat([txt_query, img_query], dim=1) diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index bd6a7156215..44bd9058c7d 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -443,8 +443,8 @@ def create_model_config(self) -> OmniModelConfig: @dataclass class OmniAsyncEngineArgs(AsyncEngineArgs, OmniEngineArgs): @classmethod - def add_cli_args(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser: - parser = AsyncEngineArgs.add_cli_args(parser) + def add_cli_args(cls, parser: argparse.ArgumentParser, async_args_only: bool = False) -> argparse.ArgumentParser: + parser = AsyncEngineArgs.add_cli_args(parser, async_args_only=async_args_only) parser = OmniEngineArgs._add_omni_specific_args(parser) return parser diff --git a/vllm_omni/engine/stage_engine_core_proc_manager.py b/vllm_omni/engine/stage_engine_core_proc_manager.py index b14c62d8f60..94a3e58b183 100644 --- a/vllm_omni/engine/stage_engine_core_proc_manager.py +++ b/vllm_omni/engine/stage_engine_core_proc_manager.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 """Process manager for omni stage engine subprocesses. @@ -83,6 +83,12 @@ def __init__( if local_engine_count <= 0: raise ValueError(f"local_engine_count must be > 0, got {local_engine_count}") + # Mirrors the vLLM 0.29 parent __init__: the inherited shutdown() reads + # this to bound how long in-flight requests may drain. Omitting it makes + # shutdown raise AttributeError, which leaves the engine core + # subprocesses alive and hangs interpreter exit until the job timeout. + self._request_shutdown_timeout = vllm_config.shutdown_timeout + context = get_mp_context() common_kwargs: dict[str, object] = { "vllm_config": vllm_config, diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 765549bcb37..de97f7aa38d 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -54,7 +54,7 @@ from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices from vllm_omni.entrypoints.utils import filter_dataclass_kwargs from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams -from vllm_omni.inputs.preprocess import OmniInputPreprocessor +from vllm_omni.inputs.preprocess import build_omni_renderer, omni_renderer_cls from vllm_omni.outputs.output_processor import MultimodalOutputProcessor from vllm_omni.platforms import current_omni_platform from vllm_omni.quantization.inc_config import OmniINCConfig @@ -1549,25 +1549,24 @@ def render_messages(self, messages, params): def _build_token_only_renderer(stage_vllm_config: Any) -> BaseRenderer: - return _TokenOnlyRenderer(stage_vllm_config, tokenizer=None) + return omni_renderer_cls(_TokenOnlyRenderer)(stage_vllm_config, tokenizer=None) def build_stage0_input_processor(stage_vllm_config: Any) -> InputProcessor: - """Build the shared stage-0 input processor.""" + """Build the shared stage-0 input processor. + + The renderer is the Omni subclass of the upstream renderer class that + ``renderer_from_config`` would have picked (or of the token-only renderer + when the stage skips tokenizer initialization), so upstream's + ``InputProcessor`` runs unmodified on top of it. + """ patch_generation_config_if_needed(stage_vllm_config.model_config) if bool(getattr(stage_vllm_config.model_config, "skip_tokenizer_init", False)): - input_processor = InputProcessor( - vllm_config=stage_vllm_config, - renderer=_build_token_only_renderer(stage_vllm_config), - ) + renderer = _build_token_only_renderer(stage_vllm_config) else: - input_processor = InputProcessor(vllm_config=stage_vllm_config) - input_processor.input_preprocessor = OmniInputPreprocessor( - vllm_config=stage_vllm_config, - renderer=input_processor.renderer, - ) - return input_processor + renderer = build_omni_renderer(stage_vllm_config) + return InputProcessor(vllm_config=stage_vllm_config, renderer=renderer) def device_init_lock_path(device_id: int, lock_dir: str = "/tmp") -> str: diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 948f7e53a82..0012039e0fb 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -51,9 +51,9 @@ from vllm_omni.outputs import OmniRequestOutput if TYPE_CHECKING: - from vllm.inputs.preprocess import InputPreprocessor from vllm.tokenizers import TokenizerLike from vllm.v1.engine import PauseMode + from vllm.v1.engine.input_processor import InputProcessor from vllm_omni.inputs.data import OmniInteractionPrompt, OmniPromptType @@ -1714,7 +1714,7 @@ def dead_error(self) -> BaseException: # ==================== EngineClient Interface ==================== - async def get_input_preprocessor(self) -> InputPreprocessor: + async def get_input_preprocessor(self) -> InputProcessor: """Get input preprocessor.""" return self.input_processor @@ -1755,10 +1755,12 @@ async def start_weight_update(self, is_checkpoint_format: bool = True) -> None: """ logger.debug("Weight update start requested (no-op in omni)") - async def finish_weight_update(self) -> None: + async def finish_weight_update(self, weight_version: str | None = None) -> None: """Finish the current weight update. Omni does not currently support weight transfer, so this is a no-op. + ``weight_version`` is accepted for upstream ``EngineClient`` protocol + compatibility (RLHF weight-transfer routers pass it positionally). """ logger.debug("Weight update finish requested (no-op in omni)") diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 78b212f98f6..5c0204a3345 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -18,7 +18,7 @@ import uvloop from vllm.entrypoints.cli.types import CLISubcommand -from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args +from vllm.entrypoints.launchers.cli_args import make_arg_parser, validate_parsed_serve_args from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG from vllm.logger import init_logger diff --git a/vllm_omni/entrypoints/duplex/chat_fallback.py b/vllm_omni/entrypoints/duplex/chat_fallback.py index 6e4d0a0e0e4..f41d7205bb6 100644 --- a/vllm_omni/entrypoints/duplex/chat_fallback.py +++ b/vllm_omni/entrypoints/duplex/chat_fallback.py @@ -9,7 +9,7 @@ from typing import Any from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest -from vllm.entrypoints.openai.engine.protocol import ErrorResponse +from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.logger import init_logger from vllm_omni.entrypoints.duplex.protocol import DuplexSession diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index d976eb18170..e020be7bf87 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -34,7 +34,10 @@ from vllm.engine.protocol import EngineClient from vllm.entrypoints.anthropic.serving import AnthropicServingMessages from vllm.entrypoints.chat_utils import ChatTemplateConfig, load_chat_template -from vllm.entrypoints.launcher import serve_http, terminate_if_errored +from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata +from vllm.entrypoints.launchers.cli_args import make_arg_parser +from vllm.entrypoints.launchers.launcher import serve_http, terminate_if_errored +from vllm.entrypoints.launchers.utils.server_utils import get_uvicorn_log_config from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer from vllm.entrypoints.openai.api_server import build_app as build_openai_app from vllm.entrypoints.openai.api_server import setup_server as setup_openai_server @@ -43,19 +46,11 @@ ChatCompletionRequest, ChatCompletionResponse, ) -from vllm.entrypoints.openai.cli_args import make_arg_parser # yapf conflicts with isort for this block # yapf: disable # yapf: enable from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion -from vllm.entrypoints.openai.engine.protocol import ( - ErrorResponse, - ModelCard, - ModelList, - ModelPermission, - RequestResponseMetadata, -) from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses @@ -64,14 +59,13 @@ from vllm.entrypoints.pooling.pooling.serving import ServingPooling from vllm.entrypoints.pooling.scoring.serving import ServingScores from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens - -# vLLM < 0.28 keeps create_error_response under serve.utils (it is not -# re-exported from vllm.entrypoints.serve); 0.28+ moved it under -# serve.exception_handling and re-exports it from the package root. -try: - from vllm.entrypoints.serve import create_error_response -except ImportError: - from vllm.entrypoints.serve.utils.error_response import create_error_response +from vllm.entrypoints.serve import create_error_response +from vllm.entrypoints.serve.engine.protocol import ( + ErrorResponse, + ModelCard, + ModelList, + ModelPermission, +) # vLLM moved `base` from openai.basic.api_router to serve.instrumentator.basic. # Keep a fallback for older/newer upstream layouts during rebase windows. @@ -85,7 +79,6 @@ ) from vllm.entrypoints.serve.utils.orca_metrics import metrics_header from vllm.entrypoints.serve.utils.request_logger import RequestLogger -from vllm.entrypoints.serve.utils.server_utils import get_uvicorn_log_config from vllm.entrypoints.speech_to_text.realtime.serving import OpenAIServingRealtime from vllm.entrypoints.speech_to_text.transcription.serving import ( OpenAIServingTranscription, diff --git a/vllm_omni/entrypoints/openai/batch_serving.py b/vllm_omni/entrypoints/openai/batch_serving.py index bc5ef36cab2..74d78dff71d 100644 --- a/vllm_omni/entrypoints/openai/batch_serving.py +++ b/vllm_omni/entrypoints/openai/batch_serving.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 asyncio import time @@ -11,7 +11,7 @@ ChatCompletionResponse, ChatCompletionResponseChoice, ) -from vllm.entrypoints.openai.engine.protocol import ( +from vllm.entrypoints.serve.engine.protocol import ( ErrorResponse, UsageInfo, ) diff --git a/vllm_omni/entrypoints/openai/realtime_connection.py b/vllm_omni/entrypoints/openai/realtime_connection.py index aad7ea658fb..86a8ba0319d 100644 --- a/vllm_omni/entrypoints/openai/realtime_connection.py +++ b/vllm_omni/entrypoints/openai/realtime_connection.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 asyncio @@ -8,7 +11,7 @@ from uuid import uuid4 import numpy as np -from vllm.entrypoints.openai.engine.protocol import UsageInfo +from vllm.entrypoints.serve.engine.protocol import UsageInfo from vllm.entrypoints.speech_to_text.realtime.connection import RealtimeConnection as VllmRealtimeConnection from vllm.entrypoints.speech_to_text.realtime.protocol import TranscriptionDelta, TranscriptionDone from vllm.logger import init_logger diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index bfa241a40ea..f5cb24efd29 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -54,8 +54,17 @@ soundfile = None +from vllm.entrypoints.generate.base.protocol import ( + DeltaFunctionCall, + DeltaMessage, + DeltaToolCall, + FunctionCall, + FunctionDefinition, + RequestResponseMetadata, + ToolCall, +) from vllm.entrypoints.generate.base.serving import clamp_prompt_logprobs -from vllm.entrypoints.launcher import terminate_if_errored +from vllm.entrypoints.launchers.launcher import terminate_if_errored from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, @@ -69,22 +78,15 @@ _get_mm_token_counts, _make_prompt_tokens_details, ) -from vllm.entrypoints.openai.engine.protocol import ( - DeltaFunctionCall, - DeltaMessage, - DeltaToolCall, - ErrorInfo, - ErrorResponse, - FunctionCall, - FunctionDefinition, - RequestResponseMetadata, - ToolCall, - UsageInfo, -) from vllm.entrypoints.openai.parser.harmony_utils import ( get_streamable_parser_for_assistant, ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.entrypoints.serve.engine.protocol import ( + ErrorInfo, + ErrorResponse, + UsageInfo, +) from vllm.entrypoints.serve.engine.typing import ChatLikeRequest from vllm.entrypoints.serve.utils.api_utils import should_include_usage from vllm.entrypoints.serve.utils.tool_calls_utils import maybe_filter_parallel_tool_calls @@ -700,7 +702,7 @@ async def _create_chat_completion( if negative_prompt is not None: tprompt["negative_prompt"] = negative_prompt # Always attach mm_processor_kwargs (possibly empty) so - # OmniInputPreprocessor._process_text routes through the + # OmniRenderer routes through the # multimodal processor path. Without it, the preprocessor # falls back to plain _tokenize_prompt and AR-based image-gen # models like GLM-Image never see their image-generation diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 88d9f090d39..a9256d95332 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -24,12 +24,10 @@ import torch from fastapi import HTTPException, Request, UploadFile from fastapi.responses import Response, StreamingResponse +from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata from vllm.entrypoints.generate.base.serving import GenerateBaseServing as OpenAIServing -from vllm.entrypoints.launcher import terminate_if_errored -from vllm.entrypoints.openai.engine.protocol import ( - ErrorResponse, - RequestResponseMetadata, -) +from vllm.entrypoints.launchers.launcher import terminate_if_errored +from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.logger import init_logger from vllm.multimodal.media import MediaConnector from vllm.utils import random_uuid diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index 35a39b2eacc..0f01fb28f91 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -156,7 +156,9 @@ def _decode_video_bytes( frames_array, metadata = loader.load_bytes( video_bytes, num_frames=num_frames, - backend="pyav", + # vLLM 0.29 removed the "pyav" decoder backend; "opencv" is + # upstream's default and the remaining CPU-only option. + backend="opencv", keep=keep, ) except Exception as exc: diff --git a/vllm_omni/entrypoints/openai/video_stream_session.py b/vllm_omni/entrypoints/openai/video_stream_session.py index ead3b1ec50c..f77fcf74e76 100644 --- a/vllm_omni/entrypoints/openai/video_stream_session.py +++ b/vllm_omni/entrypoints/openai/video_stream_session.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 """Video streaming session manager and WebSocket handler. Provides ``VideoStreamConfig``, ``VideoStreamSession`` (frame/audio buffer @@ -20,7 +20,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ) -from vllm.entrypoints.openai.engine.protocol import ErrorResponse +from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.logger import init_logger from vllm_omni.entrypoints.openai.video_frame_filter import FrameSimilarityFilter diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py index f18a00ada1c..f4cb416f64e 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Engine-level KV cache orchestration for AR-Diffusion models. This is the *body* of AR-Diffusion's KV management: it owns a vLLM ``KVCacheManager`` (a @@ -22,10 +23,10 @@ KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, - KVCacheTensor, ) from vllm.v1.request import RequestStatus +from vllm_omni.diffusion.diffusion_kv.layout import build_kv_cache_tensor from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig from vllm_omni.experimental.ar_diffusion.kv_cache.paged import ( @@ -134,7 +135,7 @@ def build_kv_manager( """ layer_names = list(layer_names) group = KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec) - tensors = [KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=layer_names)] + tensors = [build_kv_cache_tensor(spec, num_blocks, layer_names)] config = KVCacheConfig( num_blocks=num_blocks, kv_cache_tensors=tensors, diff --git a/vllm_omni/inputs/mm_processor.py b/vllm_omni/inputs/mm_processor.py new file mode 100644 index 00000000000..bb02fe9612e --- /dev/null +++ b/vllm_omni/inputs/mm_processor.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Multimodal data plumbing for Omni's custom, non-HF processors.""" + +from collections.abc import Mapping +from typing import TypeVar + +from transformers import BatchFeature +from vllm.multimodal.parse import MultiModalDataItems +from vllm.multimodal.processing import BaseMultiModalProcessor, BaseProcessingInfo + +_I = TypeVar("_I", bound=BaseProcessingInfo) + + +class OmniMultiModalProcessor(BaseMultiModalProcessor[_I]): + """Share data/passthrough handling, leaving model processing in its owner. + + These processors implement ``_call_hf_processor`` locally, sometimes + without any HF processor object. Prompt tokens are prepared separately + in ``apply``; the upstream base still owns caching and prompt updates. + """ + + _OMNI_PROMPT_TEXT_KEY = "_vllm_omni_original_prompt_text" + + def _apply_hf_processor_main( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> BatchFeature: + kwargs = dict(hf_processor_mm_kwargs) + prompt = kwargs.pop(self._OMNI_PROMPT_TEXT_KEY, None) + if prompt is None: + prompt = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) + valid_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0}) + mm_data, passthrough = self._get_hf_mm_data(valid_items) + result = self._call_hf_processor(str(prompt), dict(mm_data), kwargs, {}) + result.update(passthrough) + return result diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py index 3da1ef89715..dc4105b4db3 100644 --- a/vllm_omni/inputs/preprocess.py +++ b/vllm_omni/inputs/preprocess.py @@ -1,200 +1,118 @@ -from typing import Any +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Omni renderer: the upstream renderer class with Omni engine-input routing. -from typing_extensions import assert_never -from vllm.inputs import EmbedsInput, MultiModalInput, SingletonInput -from vllm.inputs.preprocess import InputPreprocessor -from vllm.logger import init_logger -from vllm.renderers.inputs import SingletonDictPrompt +``OmniRenderer`` is a mixin. ``omni_renderer_cls`` puts it in front of a +concrete ``BaseRenderer`` subclass, so the Omni instance *is* an upstream +renderer (tokenizer, executors, multimodal cache and every public entry point +come from upstream unchanged) and only ``_process_singleton`` / +``_process_singleton_async`` are extended. ``build_omni_renderer`` resolves +the concrete class the same way upstream ``renderer_from_config`` does. +""" -from vllm_omni.inputs.data import ( - OmniEmbedsPrompt, - OmniTextPrompt, - OmniTokenInputs, - OmniTokensPrompt, - token_inputs_omni, -) +from __future__ import annotations -logger = init_logger(__name__) +from typing import TYPE_CHECKING, cast +from vllm.renderers import BaseRenderer +from vllm.renderers.registry import RENDERER_REGISTRY +from vllm.tokenizers.registry import cached_tokenizer_from_config, tokenizer_args_from_config -class OmniInputPreprocessor(InputPreprocessor): - """Input preprocessor for omni models. +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.tokenizers import TokenizerLike - Extends the base InputPreprocessor to handle omni-specific input - types including prompt embeddings and additional information payloads. - Supports processing tokens, embeddings, text, and multimodal inputs. - """ - def _process_text( - self, - parsed_content: OmniTextPrompt, - tokenization_kwargs: dict[str, Any] | None = None, - *, - mm_uuids: Any | None = None, - ) -> OmniTokenInputs | MultiModalInput: - """Process text prompts with support for mm_processor_kwargs. - - Extends base class to support mm_processor_kwargs without multi_modal_data. - This is needed for models like GLM-Image where text-to-image generation - requires processor kwargs (target_h, target_w) to format the prompt. - """ - prompt_text = parsed_content["prompt"] - mm_processor_kwargs = parsed_content.get("mm_processor_kwargs") or {} - # When the deprecated raw-prompt path is used, process_inputs does - # not pass mm_uuids to preprocess(). Fall back to reading it from - # the prompt dict so the Renderer's _validate_mm_uuids can see it. - effective_mm_uuids = mm_uuids or parsed_content.get("multi_modal_uuids") - - inputs: OmniTokenInputs | MultiModalInput - if multi_modal_data := parsed_content.get("multi_modal_data"): - inputs = self._process_multimodal( - prompt_text, - multi_modal_data, - mm_processor_kwargs, - tokenization_kwargs=tokenization_kwargs, - mm_uuids=effective_mm_uuids, - ) - prompt_embeds = parsed_content.get("prompt_embeds") - if prompt_embeds is not None: - inputs["prompt_embeds"] = prompt_embeds - additional_information = parsed_content.get("additional_information") - if additional_information is not None: - inputs["additional_information"] = additional_information - model_intermediate_buffer = parsed_content.get("model_intermediate_buffer") - if model_intermediate_buffer is not None: - inputs["model_intermediate_buffer"] = model_intermediate_buffer - elif "mm_processor_kwargs" in parsed_content: - # Presence — not truthiness. An explicitly-set empty dict still - # signals "route through the multimodal processor" (needed for - # AR-based image-gen where the HF processor supplies its own - # defaults and scaffold). - inputs = self._process_multimodal( - prompt_text, - {}, - mm_processor_kwargs, - tokenization_kwargs=tokenization_kwargs, - mm_uuids=effective_mm_uuids, - ) - else: - prompt_token_ids = self._tokenize_prompt( - prompt_text, - tokenization_kwargs=tokenization_kwargs, - ) - inputs = token_inputs_omni( - prompt_token_ids, - prompt_embeds=parsed_content.get("prompt_embeds"), - additional_information=parsed_content.get("additional_information"), - model_intermediate_buffer=parsed_content.get("model_intermediate_buffer"), - ) - prompt_embeds = parsed_content.get("prompt_embeds") - if prompt_embeds is not None: - inputs["prompt_embeds"] = prompt_embeds - additional_information = parsed_content.get("additional_information") - if additional_information is not None: - inputs["additional_information"] = additional_information - model_intermediate_buffer = parsed_content.get("model_intermediate_buffer") - if model_intermediate_buffer is not None: - inputs["model_intermediate_buffer"] = model_intermediate_buffer - if cache_salt := parsed_content.get("cache_salt"): - inputs["cache_salt"] = cache_salt +class OmniRenderer: + """Mixin that routes Omni prompts through the upstream renderer. - return inputs + It has no base class of its own; ``omni_renderer_cls`` prepends it to a + concrete ``BaseRenderer`` subclass, so ``self`` and ``super()`` are that + renderer at runtime (hence the casts below). - def _process_tokens( - self, - parsed_content: OmniTokensPrompt, - tokenization_kwargs: dict[str, Any] | None = None, - ) -> OmniTokenInputs | MultiModalInput: - prompt_token_ids = self._truncate_inputs(parsed_content["prompt_token_ids"], tokenization_kwargs) - prompt_embeds = parsed_content.get("prompt_embeds") - additional_information = parsed_content.get("additional_information") - model_intermediate_buffer = parsed_content.get("model_intermediate_buffer") - - multi_modal_data = parsed_content.get("multi_modal_data") - - inputs: OmniTokenInputs | MultiModalInput - if multi_modal_data: - inputs = self._process_multimodal( - prompt_token_ids, - multi_modal_data, - parsed_content.get("mm_processor_kwargs"), - tokenization_kwargs=tokenization_kwargs, - mm_uuids=parsed_content.get("multi_modal_uuids"), - ) + Two extensions over upstream ``_process_singleton``: - else: - inputs = token_inputs_omni( - prompt_token_ids=prompt_token_ids, - prompt_embeds=prompt_embeds, - additional_information=additional_information, - model_intermediate_buffer=model_intermediate_buffer, - ) - if prompt_embeds is not None: - inputs["prompt_embeds"] = prompt_embeds - if additional_information is not None: - inputs["additional_information"] = additional_information - if model_intermediate_buffer is not None: - inputs["model_intermediate_buffer"] = model_intermediate_buffer - if prompt_text := parsed_content.get("prompt"): - inputs["prompt"] = prompt_text - if cache_salt := parsed_content.get("cache_salt"): - inputs["cache_salt"] = cache_salt + * a token prompt that carries ``mm_processor_kwargs`` but no + ``multi_modal_data`` still goes through the multimodal processor (AR + image-generation models such as GLM-Image read their target size from + those kwargs); + * Omni-only prompt keys (``prompt``, ``cache_salt``, + ``additional_information``, ``model_intermediate_buffer``) are copied + onto the engine input. + """ + @staticmethod + def _with_omni_extras(inputs, prompt): + for key in ("prompt", "cache_salt", "additional_information", "model_intermediate_buffer"): + if key in prompt: + inputs[key] = prompt[key] return inputs - def _process_embeds( - self, - parsed_content: OmniEmbedsPrompt, - ) -> EmbedsInput: - """Process embeddings prompt with omni-specific extensions. - - Extends base _process_embeds to handle additional_information payload - for direct transfer between pipeline stages. - """ - # Call parent implementation for base embeds processing - inputs = super()._process_embeds(parsed_content) - - # Add omni-specific additional_information if present - additional_information = parsed_content.get("additional_information") - if additional_information is not None: - inputs["additional_information"] = additional_information # type: ignore[typeddict-unknown-key] - model_intermediate_buffer = parsed_content.get("model_intermediate_buffer") - if model_intermediate_buffer is not None: - inputs["model_intermediate_buffer"] = model_intermediate_buffer # type: ignore[typeddict-unknown-key] + @staticmethod + def _routes_no_media_kwargs(prompt) -> bool: + return "prompt_embeds" not in prompt and "mm_processor_kwargs" in prompt and not prompt.get("multi_modal_data") - return inputs - - def _prompt_to_llm_inputs( - self, - prompt: SingletonDictPrompt, - tokenization_kwargs: dict[str, Any] | None = None, - *, - mm_uuids: Any | None = None, - ) -> SingletonInput: - """ - Extract the singleton inputs from a prompt. + def _process_singleton(self, prompt, *, skip_mm_cache: bool = False): + renderer = cast(BaseRenderer, self) + if self._routes_no_media_kwargs(prompt): + inputs = renderer._process_multimodal( + prompt["prompt_token_ids"], + {}, + mm_processor_kwargs=prompt["mm_processor_kwargs"], + mm_uuids=prompt.get("multi_modal_uuids"), + skip_mm_cache=skip_mm_cache, + ) + else: + inputs = cast(BaseRenderer, super())._process_singleton(prompt, skip_mm_cache=skip_mm_cache) + return self._with_omni_extras(inputs, prompt) + + async def _process_singleton_async(self, prompt, *, skip_mm_cache: bool = False): + renderer = cast(BaseRenderer, self) + if self._routes_no_media_kwargs(prompt): + inputs = await renderer._process_multimodal_async( + prompt["prompt_token_ids"], + {}, + mm_processor_kwargs=prompt["mm_processor_kwargs"], + mm_uuids=prompt.get("multi_modal_uuids"), + skip_mm_cache=skip_mm_cache, + ) + else: + inputs = await cast(BaseRenderer, super())._process_singleton_async(prompt, skip_mm_cache=skip_mm_cache) + return self._with_omni_extras(inputs, prompt) - Arguments: - * prompt: single encoder or decoder input prompt +_OMNI_RENDERER_CLASSES: dict[type[BaseRenderer], type[BaseRenderer]] = {} - Returns: - * [`SingletonInput`][vllm.inputs.engine.SingletonInput] instance - """ - if "prompt_embeds" in prompt: - return self._process_embeds(prompt) # type: ignore[arg-type] +def omni_renderer_cls(base_cls: type[BaseRenderer]) -> type[BaseRenderer]: + """Return the Omni subclass of ``base_cls`` (created once per base class).""" + if issubclass(base_cls, OmniRenderer): + return base_cls + cls = _OMNI_RENDERER_CLASSES.get(base_cls) + if cls is None: + name = f"Omni{base_cls.__name__}" + cls = type(name, (OmniRenderer, base_cls), {"__module__": __name__, "__qualname__": name}) + _OMNI_RENDERER_CLASSES[base_cls] = cls + return cls - if "prompt_token_ids" in prompt: - return self._process_tokens( - prompt, # type: ignore[arg-type] - ) - if "prompt" in prompt: - return self._process_text( - prompt, # type: ignore[arg-type] - tokenization_kwargs=tokenization_kwargs, - mm_uuids=mm_uuids, - ) +def build_omni_renderer( + vllm_config: VllmConfig, + *, + renderer_cls: type[BaseRenderer] | None = None, + tokenizer: TokenizerLike | None = None, +) -> BaseRenderer: + """Build the Omni renderer for ``vllm_config``. - assert_never(prompt) # type: ignore[arg-type] + With ``renderer_cls`` unset this mirrors upstream ``renderer_from_config``: + the tokenizer and renderer mode come from the model config and the class + from ``RENDERER_REGISTRY``. Pass ``renderer_cls`` (and optionally + ``tokenizer``) to skip that resolution, e.g. for tokenizer-less stages. + """ + model_config = vllm_config.model_config + if renderer_cls is None: + if tokenizer is None: + tokenizer = cached_tokenizer_from_config(model_config) + renderer_mode, *_ = tokenizer_args_from_config(model_config) + renderer_cls = RENDERER_REGISTRY.load_renderer_cls(renderer_mode) + return omni_renderer_cls(renderer_cls)(vllm_config, tokenizer) diff --git a/vllm_omni/model_executor/models/audex/audex_omni_processing.py b/vllm_omni/model_executor/models/audex/audex_omni_processing.py index 5e4c5390835..7b80f53f857 100644 --- a/vllm_omni/model_executor/models/audex/audex_omni_processing.py +++ b/vllm_omni/model_executor/models/audex/audex_omni_processing.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 # # Adapted from nvidia/Nemotron-Labs-Audex-2B (Apache-2.0): # inference_scripts_vllm/audioqa_scripts/audex_2b_vllm/audio_features.py @@ -33,13 +33,14 @@ from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, PromptReplacement, PromptUpdate, PromptUpdateDetails, ) +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor + SOUND_TOKEN = "" SOUND_START_TOKEN = "" SOUND_END_TOKEN = "" @@ -176,7 +177,7 @@ def get_dummy_mm_data( return {"audio": self._get_dummy_audios(length=audio_len, num_audios=num_audios, overrides=overrides)} -class AudexMultiModalProcessor(BaseMultiModalProcessor[AudexProcessingInfo]): +class AudexMultiModalProcessor(OmniMultiModalProcessor[AudexProcessingInfo]): def _call_hf_processor( self, prompt: str, diff --git a/vllm_omni/model_executor/models/bagel/bagel.py b/vllm_omni/model_executor/models/bagel/bagel.py index 495e0037f29..30dced9d053 100644 --- a/vllm_omni/model_executor/models/bagel/bagel.py +++ b/vllm_omni/model_executor/models/bagel/bagel.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from collections.abc import Iterable, Mapping, Sequence +from dataclasses import replace from math import isqrt from typing import Any @@ -20,7 +21,7 @@ from vllm.model_executor.models.bagel import BagelForConditionalGeneration from vllm.model_executor.models.interfaces import MultiModalEmbeddings from vllm.model_executor.models.qwen2 import Qwen2DecoderLayer, Qwen2MLP -from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, @@ -221,6 +222,20 @@ def _get_subparsers(self): class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingInfo]): IMG2IMG_PLACEHOLDER = "<|fim_middle|>" + def apply(self, inputs, timing_ctx): + num_img2img = inputs.mm_data_items.get_all_counts().get("img2img", 0) + if num_img2img: + token_id = self.info.get_tokenizer().get_vocab().get(self.IMG2IMG_PLACEHOLDER) + if token_id is not None: + prompt_ids = list(inputs.prompt) + missing = num_img2img - prompt_ids.count(token_id) + if missing > 0: + inputs = replace( + inputs, + prompt=[token_id] * missing + prompt_ids, + ) + return super().apply(inputs, timing_ctx) + @staticmethod def _mm_kwargs_for_bagel_img2img_hf(mm_kwargs: Mapping[str, object]) -> dict[str, object]: # OpenAI / GLM-style serving may pass target_h/target_w for output grid sizing. @@ -241,69 +256,51 @@ def _get_mm_fields_config(self, hf_inputs, hf_processor_mm_kwargs): "pixel_values_img2img": MultiModalFieldConfig.batched("img2img"), } - def _call_hf_processor( + def _apply_hf_processor_main( self, - prompt: str, - mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], - ) -> "BatchFeature": - has_image = "images" in mm_data - has_img2img = "pixel_values_img2img" in mm_data - - if has_img2img and self.IMG2IMG_PLACEHOLDER not in prompt: - prompt = f"{self.IMG2IMG_PLACEHOLDER}{prompt}" - - if has_image and has_img2img: - outputs = BatchFeature() - - img_data = dict(mm_data) - if "pixel_values_img2img" in img_data: - del img_data["pixel_values_img2img"] - kwargs_img = dict(mm_kwargs) - kwargs_img["is_img2img"] = False - out_img = super()._call_hf_processor(prompt, img_data, kwargs_img, tok_kwargs) - if "pixel_values" in out_img: - outputs["pixel_values"] = out_img["pixel_values"] - for k, v in out_img.items(): - if k != "pixel_values": - outputs[k] = v - - img2img_data = dict(mm_data) - if "images" in img2img_data: - del img2img_data["images"] - img2img_data["images"] = img2img_data.pop("pixel_values_img2img") - kwargs_img2img = self._mm_kwargs_for_bagel_img2img_hf(mm_kwargs) - kwargs_img2img["is_img2img"] = True - out_img2img = super()._call_hf_processor(prompt, img2img_data, kwargs_img2img, tok_kwargs) - if "pixel_values" in out_img2img: - outputs["pixel_values_img2img"] = out_img2img["pixel_values"] - for k, v in out_img2img.items(): - if k not in outputs: - outputs[k] = v - - return outputs - - elif has_img2img: - mm_data = dict(mm_data) - mm_data["images"] = mm_data.pop("pixel_values_img2img") - mm_kwargs = self._mm_kwargs_for_bagel_img2img_hf(mm_kwargs) - mm_kwargs["is_img2img"] = True - outputs = super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) - if "pixel_values" in outputs: - outputs["pixel_values_img2img"] = outputs.pop("pixel_values") - return outputs - - return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) - - def _hf_processor_applies_updates( - self, - prompt_text: str, mm_items: MultiModalDataItems, hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False + ) -> BatchFeature: + valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0}) + mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) + prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + + has_image = "images" in mm_data + has_img2img = "pixel_values_img2img" in mm_data + processed_data = BatchFeature() + + if has_image: + image_kwargs = {**hf_processor_mm_kwargs, "is_img2img": False} + image_outputs = self.info.ctx.call_hf_processor( + processor, + {"text": prompt_text, "images": mm_data["images"]}, + image_kwargs, + ) + processed_data.update(image_outputs) + + if has_img2img: + img2img_kwargs = self._mm_kwargs_for_bagel_img2img_hf(hf_processor_mm_kwargs) + img2img_kwargs["is_img2img"] = True + img2img_outputs = self.info.ctx.call_hf_processor( + processor, + { + "text": prompt_text, + "images": mm_data["pixel_values_img2img"], + }, + img2img_kwargs, + ) + pixel_values = img2img_outputs.pop("pixel_values", None) + if pixel_values is not None: + processed_data["pixel_values_img2img"] = pixel_values + for key, value in img2img_outputs.items(): + processed_data.setdefault(key, value) + + if not has_image and not has_img2img: + processed_data = BatchFeature(dict(passthrough_data)) + else: + processed_data.update(passthrough_data) + return processed_data def _get_prompt_updates( self, @@ -372,7 +369,8 @@ def get_img2img_replacement(item_idx: int): embed_mask = [True] * num_vae_total + [False] + [True] * num_vit_total return PromptUpdateDetails( full=tokens, - is_embed=lambda _tok, _seq, _m=embed_mask: torch.tensor(_m, dtype=torch.bool), + # vLLM 0.29 calls is_embed with just the full token list. + is_embed=lambda _full, _m=embed_mask: torch.tensor(_m, dtype=torch.bool), ) replacements.append( @@ -1102,11 +1100,12 @@ def _map_vae_weight_name(name: str) -> str: filtered_weights.append((mapped_name, tensor)) loader = AutoWeightsLoader( - self, - skip_prefixes=["vit_pos_embed.pos_embed"], - ignore_unexpected_prefixes=["vae.", "latent_pos_embed.", "time_embedder.", "vae2llm."], + self, ignore_unexpected_prefixes=["vae.", "latent_pos_embed.", "time_embedder.", "vae2llm."] + ) + loaded = loader.load_weights( + filtered_weights, + mapper=(self.hf_to_vllm_mapper) | WeightsMapper(orig_to_new_prefix={"vit_pos_embed.pos_embed": None}), ) - loaded = loader.load_weights(filtered_weights, mapper=self.hf_to_vllm_mapper) loaded |= self._load_moe_gen_weights(moe_gen_weights) diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index 84652cadabf..558e96b65c2 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import os from collections.abc import Iterable, Mapping, Sequence +from dataclasses import replace from functools import partial from math import gcd from threading import Lock @@ -24,7 +25,6 @@ from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, ProcessorInputs, PromptIndexTargets, @@ -38,6 +38,7 @@ from vllm.v1.sample.sampler import Sampler from vllm_omni.data_entry_keys import EmbeddingsStruct, OmniPayloadStruct, to_dict, to_struct +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer from vllm_omni.model_executor.models.cosyvoice3.utils import ( concat_text_with_prompt_ids, @@ -108,7 +109,45 @@ def get_data_parser(self): ) -class CosyVoice3MultiModalProcessor(BaseMultiModalProcessor[CosyVoice3MultiModalProcessingInfo]): +class CosyVoice3MultiModalProcessor(OmniMultiModalProcessor[CosyVoice3MultiModalProcessingInfo]): + def apply(self, inputs: ProcessorInputs, timing_ctx): + tokenizer = self.info.get_tokenizer() + prompt_text = tokenizer.decode(inputs.prompt, skip_special_tokens=False) + config = self.info.ctx.get_hf_config() + model_dir = self.info.ctx.model_config.model + self._ensure_cached_runtime_components(model_dir, config) + + text_token, text_token_len = extract_text_token( + prompt_text, + self.tokenizer, + config.allowed_special, + ) + if inputs.mm_data_items.get_all_counts().get("audio", 0): + reference_text = inputs.hf_processor_mm_kwargs.get("prompt_text") + if not isinstance(reference_text, str): + raise ValueError(f"prompt text is None : {reference_text}") + prompt_text_token, prompt_text_token_len = extract_text_token( + reference_text, + self.tokenizer, + config.allowed_special, + ) + text_token, _ = concat_text_with_prompt_ids( + text_token, + text_token_len, + prompt_text_token, + prompt_text_token_len, + ) + + inputs = replace( + inputs, + prompt=text_token.reshape(-1).tolist(), + hf_processor_mm_kwargs={ + **inputs.hf_processor_mm_kwargs, + self._OMNI_PROMPT_TEXT_KEY: prompt_text, + }, + ) + return super().apply(inputs, timing_ctx) + def _ensure_cached_runtime_components(self, model_dir: str, config: CosyVoice3Config) -> None: cached_model_dir = getattr(self, "_cached_model_dir", None) if cached_model_dir == model_dir: @@ -376,15 +415,6 @@ def _get_mm_fields_config( "embedding": MultiModalFieldConfig.batched("audio"), } - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False - def _get_prompt_updates( self, mm_items: MultiModalDataItems, diff --git a/vllm_omni/model_executor/models/covo_audio/covo_audio.py b/vllm_omni/model_executor/models/covo_audio/covo_audio.py index efbe7b0d667..108bd36c53d 100644 --- a/vllm_omni/model_executor/models/covo_audio/covo_audio.py +++ b/vllm_omni/model_executor/models/covo_audio/covo_audio.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # Copyright 2026 Tencent. from collections.abc import Iterable, Mapping, Sequence from typing import Any @@ -17,7 +20,6 @@ from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, ProcessorInputs, PromptReplacement, @@ -26,6 +28,7 @@ from vllm.sequence import IntermediateTensors from vllm.v1.sample.metadata import SamplingMetadata +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor 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.utils import add_prefix_to_loaded_weights @@ -106,10 +109,7 @@ def get_dummy_processor_inputs( ) -class CovoAudioMultiModalProcessor(BaseMultiModalProcessor[CovoAudioProcessingInfo]): - def _hf_processor_applies_updates(self, prompt_text, mm_items, hf_processor_mm_kwargs, tokenization_kwargs) -> bool: - return False - +class CovoAudioMultiModalProcessor(OmniMultiModalProcessor[CovoAudioProcessingInfo]): def _call_hf_processor( self, prompt: str, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, Any], tok_kwargs: Mapping[str, object] ) -> BatchFeature: diff --git a/vllm_omni/model_executor/models/gepard/gepard_talker.py b/vllm_omni/model_executor/models/gepard/gepard_talker.py index 02f5da533c1..37c1e66a98a 100644 --- a/vllm_omni/model_executor/models/gepard/gepard_talker.py +++ b/vllm_omni/model_executor/models/gepard/gepard_talker.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 """Gepard-1.0 native-AR talker. Single-stage AR TTS on a vLLM-native Qwen3.5 backbone. Each step samples one @@ -22,7 +22,7 @@ from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.models.qwen3_5 import Qwen3_5ForCausalLM -from vllm.model_executor.models.utils import AutoWeightsLoader, maybe_prefix +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper, maybe_prefix from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig @@ -199,8 +199,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # NanoCodec loads lazily on first decode, so load_format=dummy and # NeMo-less environments never pay for it. - loader = AutoWeightsLoader(self, skip_prefixes=["mtp.", "ref_compressor."]) - return loader.load_weights(iter(rest)) + loader = AutoWeightsLoader(self) + return loader.load_weights( + iter(rest), mapper=WeightsMapper(orig_to_new_prefix={"mtp.": None, "ref_compressor.": None}) + ) def _get_or_create_state(self, request_id: str) -> _GepardState: st = self._active_states.get(request_id) diff --git a/vllm_omni/model_executor/models/glm_image/glm_image_ar.py b/vllm_omni/model_executor/models/glm_image/glm_image_ar.py index 77aef398b85..e97a78c2535 100644 --- a/vllm_omni/model_executor/models/glm_image/glm_image_ar.py +++ b/vllm_omni/model_executor/models/glm_image/glm_image_ar.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 # Adapted from # https://github.com/huggingface/transformers/blob/main/src/transformers/models/glm_image/modeling_glm_image.py @@ -24,6 +24,7 @@ import math import os from collections.abc import Iterable, Mapping, Sequence +from dataclasses import replace from functools import lru_cache from typing import Annotated, Literal @@ -81,7 +82,6 @@ MultiModalKwargsItems, ) from vllm.multimodal.parse import ( - ImageProcessorItems, MultiModalDataItems, MultiModalDataParser, ) @@ -341,6 +341,30 @@ def get_dummy_mm_data( } +_GLM_IMAGE_GRID_FACTOR = 32 + + +def _build_target_shape_scaffold(processor, *, height: int, width: int, is_text_to_image: bool) -> str: + """Return the grid scaffold HF appends to a GLM-Image generation prompt. + + Mirrors ``GlmImageProcessor._build_prompt_with_target_shape`` (identical in + transformers 5.13 and 5.14) using only the processor's public token + attributes, so this does not depend on a private HF method that a stub or + a future release may not provide. + """ + factor = _GLM_IMAGE_GRID_FACTOR + token_h = (height // factor * factor) // factor + token_w = (width // factor * factor) // factor + grid_bos, grid_eos, bos = processor.grid_bos_token, processor.grid_eos_token, processor.bos_token + scaffold = f"{grid_bos}{token_h} {token_w}{grid_eos}" + if is_text_to_image: + ratio = token_h / token_w + prev_token_h = int(math.sqrt(ratio) * (factor // 2)) + prev_token_w = int(math.sqrt(1 / ratio) * (factor // 2)) + scaffold += f"{grid_bos}{prev_token_h} {prev_token_w}{grid_eos}" + return scaffold + bos + + class GlmImageMultiModalProcessor(BaseMultiModalProcessor[GlmImageProcessingInfo]): """ Multimodal processor for GLM-Image. @@ -358,265 +382,38 @@ def _cached_apply_hf_processor(self, inputs, timing_ctx): return self._apply_hf_processor(inputs, timing_ctx) return super()._cached_apply_hf_processor(inputs, timing_ctx) - def _call_hf_processor( - self, - prompt: str, - mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], - ) -> BatchFeature: - """ - Call the HuggingFace processor. - - For text-to-image mode (no images), we need to: - 1. Build the prompt with target grid dimensions - 2. Build the image_grid_thw tensor for M-RoPE position encoding - - For image-to-image mode: - 1. Process source images through the image processor - 2. Build prompt with image placeholders expanded - 3. Build image_grid_thw including source and target grids - """ - processor = self.info.get_hf_processor() - - # Get target dimensions from mm_kwargs or use defaults - target_h = mm_kwargs.get("target_h", 1024) if mm_kwargs else 1024 - target_w = mm_kwargs.get("target_w", 1024) if mm_kwargs else 1024 - - logger.debug( - f"_call_hf_processor: target dimensions for generation: {target_h}x{target_w}, mm_kwargs={mm_kwargs}" - ) - - if not mm_data or not mm_data.get("images"): - # Text-to-image mode - if processor is not None: - # Build messages format expected by processor - messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}] - - # Use apply_chat_template which handles target dimensions - hf_inputs = processor.apply_chat_template( - messages, - tokenize=True, - target_h=target_h, - target_w=target_w, - return_dict=True, - return_tensors="pt", - ) - - return hf_inputs - else: - # Fallback: just tokenize (this won't work properly for generation) - tokenizer = self.info.get_tokenizer() - prompt_ids = tokenizer.encode(prompt) - return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt") - - # Image-to-image mode - # NOTE: Use "images" (plural) - this is what vLLM's ImageProcessorItems.get_processor_data() returns - images = mm_data.get("images") - if not isinstance(images, list): - images = [images] - - logger.debug( - f"_call_hf_processor i2i: num_images={len(images)}, image_types={[type(img).__name__ for img in images]}" - ) - - if processor is not None: - # Build messages with image objects directly in content - # This is how GlmImageProcessor expects images - embedded in the content dict - # NOT as a separate images= parameter - # - # IMPORTANT: Remove <|image|> placeholders from prompt since apply_chat_template - # will automatically insert them for each image in content. Having both leads to - # index out of bounds when processing image_grid_thw. - clean_prompt = prompt.replace("<|image|>", "") - content = [] - for img in images: - content.append({"type": "image", "image": img}) - content.append({"type": "text", "text": clean_prompt}) - messages = [{"role": "user", "content": content}] - - logger.debug(f"_call_hf_processor: calling apply_chat_template with {len(images)} images in content") - - # Use apply_chat_template - processor will process images when they're in content - hf_inputs = processor.apply_chat_template( - messages, - tokenize=True, - target_h=target_h, - target_w=target_w, - return_dict=True, - return_tensors="pt", - ) - - logger.debug(f"_call_hf_processor: apply_chat_template returned keys: {list(hf_inputs.keys())}") - - # IMPORTANT (i2i): vLLM multimodal encoder must see source-only grids - # (matching pixel_values and number of images), but M-RoPE needs full - # grids (source + target) to compute correct decode positions. - image_grid_thw = hf_inputs.get("image_grid_thw") - if image_grid_thw is not None: - # Preserve full grids for M-RoPE. - hf_inputs["mrope_image_grid_thw"] = image_grid_thw - - # Expose source-only grids for MM. - # In most i2i requests, we process one prompt at a time here, - # so `len(images)` is the number of source images. - num_source_images = len(images) - if image_grid_thw.shape[0] != num_source_images: - source_grids = image_grid_thw[:num_source_images] - hf_inputs["image_grid_thw"] = source_grids - logger.debug( - "_call_hf_processor: adjusted image_grid_thw for MM from %s to %s (num_source_images=%d)", - tuple(image_grid_thw.shape), - tuple(source_grids.shape), - num_source_images, - ) - - # Debug: Analyze input_ids for image tokens - input_ids = hf_inputs.get("input_ids") - if input_ids is not None: - if hasattr(input_ids, "tolist"): - ids_list = input_ids.tolist() - if isinstance(ids_list[0], list): - ids_list = ids_list[0] # Unbatch - else: - ids_list = list(input_ids) - - # Get image token ID from config - hf_config = self.info.get_hf_config() - image_token_id = getattr(hf_config, "image_token_id", 167855) - - # Count image tokens - image_token_count = ids_list.count(image_token_id) - logger.debug( - f"_call_hf_processor: input_ids length={len(ids_list)}, " - f"image_token_id={image_token_id}, " - f"image_token_count={image_token_count}" - ) - - # Log first/last few tokens to understand structure - logger.debug(f"_call_hf_processor: first 20 tokens: {ids_list[:20]}") - logger.debug(f"_call_hf_processor: last 20 tokens: {ids_list[-20:]}") - - # Find positions of image tokens - image_positions = [i for i, t in enumerate(ids_list) if t == image_token_id] - if image_positions: - logger.debug(f"_call_hf_processor: image token positions (first 10): {image_positions[:10]}") - - return hf_inputs - else: - # Fallback without processor - this is not ideal but prevents crashes - logger.warning("GlmImageProcessor not available, using fallback for i2i") - tokenizer = self.info.get_tokenizer() - hf_config = self.info.get_hf_config() - - # Get image token - image_token_id = getattr(hf_config, "image_token_id", 167855) - try: - image_token = tokenizer.convert_ids_to_tokens(image_token_id) - except Exception: - image_token = "<|image|>" - - # Build prompt with image placeholders - image_placeholders = image_token * len(images) - full_prompt = f"{image_placeholders}{prompt}" - prompt_ids = tokenizer.encode(full_prompt) - return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt") - - def _apply_hf_processor_mm_only( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> BatchFeature: - """ - Apply the HF processor on the multi-modal data only. - - GLM-Image requires special handling because apply_chat_template always - adds a target <|image|> placeholder in addition to source image placeholders. - This causes an IndexError when the HF processor tries to find grid info - for the target placeholder (which doesn't exist for source-only processing). - - Solution: Call the image processor directly to get pixel_values and - image_grid_thw, bypassing apply_chat_template's target handling. - """ - mm_counts = mm_items.get_all_counts() - num_images = mm_counts.get("image", 0) - - if num_images == 0: - # No images - call parent implementation - return super()._apply_hf_processor_mm_only( - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, + def apply(self, inputs, timing_ctx): + """Add the generation scaffold before upstream expands source images.""" + num_images = inputs.mm_data_items.get_all_counts().get("image", 0) + prompt_ids = list(inputs.prompt) + if num_images: + image_token_id = getattr( + self.info.get_hf_config(), + "image_token_id", + 167855, ) + missing = num_images - prompt_ids.count(image_token_id) + if missing > 0: + prompt_ids = [image_token_id] * missing + prompt_ids - # For i2i mode, we need to process images directly with the image processor - # to avoid the apply_chat_template target placeholder issue processor = self.info.get_hf_processor() - image_processor = processor.image_processor - - # Get images from mm_items - images = mm_items.get_items("image", ImageProcessorItems) - image_list = [images.get(i) for i in range(images.get_count())] - - logger.debug(f"_apply_hf_processor_mm_only: processing {len(image_list)} images directly") - - # Process images directly with image processor - image_inputs = image_processor( - images=image_list, - return_tensors="pt", - ) - - # Get grid info for source images only (no target) - pixel_values = image_inputs.get("pixel_values") - image_grid_thw = image_inputs.get("image_grid_thw") - if image_grid_thw is not None and image_grid_thw.shape[0] != num_images: - # Be defensive: some processors may include extra target grids. - image_grid_thw = image_grid_thw[:num_images] - image_inputs["image_grid_thw"] = image_grid_thw - - logger.debug( - f"_apply_hf_processor_mm_only: pixel_values shape=\ - {pixel_values.shape if pixel_values is not None else None}, " - f"image_grid_thw shape={image_grid_thw.shape if image_grid_thw is not None else None}" - ) - - # Build input_ids with image token placeholders - # The _get_prompt_updates returns PromptReplacement(target=[image_token_id], ...) - # which needs to find image tokens in input_ids to replace them. - # We need to include one image_token_id per image so the replacement can work. tokenizer = self.info.get_tokenizer() - image_token_id = tokenizer.convert_tokens_to_ids("<|image|>") - - # Build input_ids: [image_token] * num_images + tokenized text - # This way _apply_prompt_updates can find the image tokens and replace them - dummy_text = self.dummy_inputs.get_dummy_text(mm_counts) - text_ids = tokenizer.encode(dummy_text, add_special_tokens=False) - input_ids = [image_token_id] * num_images + text_ids - - logger.debug( - f"_apply_hf_processor_mm_only: built input_ids with {num_images} image tokens + {len(text_ids)} text tokens" + # Append HF's target-shape scaffold without decoding/re-tokenizing the + # user's tokens. Text-to-image needs both target and preview grids; i2i + # needs only the target. Grid metadata alone does not supply these AR + # tokens. + target_grid = self._build_generation_grids(inputs.hf_processor_mm_kwargs)[0] + suffix = _build_target_shape_scaffold( + processor, + height=int(target_grid[1]) * 32, + width=int(target_grid[2]) * 32, + is_text_to_image=num_images == 0, ) - - return BatchFeature( - dict( - input_ids=[input_ids], - pixel_values=pixel_values, - image_grid_thw=image_grid_thw, - ), - tensor_type="pt", - ) - - def _apply_hf_processor_text_only( - self, prompt_text: str, hf_processor_mm_kwargs: Mapping[str, object], tokenization_kwargs: Mapping[str, object] - ) -> list[int]: - prompt_ids, _, _ = super()._apply_hf_processor_text_mm( - prompt_text=prompt_text, - mm_items=MultiModalDataItems({}), - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - return prompt_ids + suffix_ids = tokenizer.encode(suffix, add_special_tokens=False) + if prompt_ids[-len(suffix_ids) :] != suffix_ids: + prompt_ids.extend(suffix_ids) + inputs = replace(inputs, prompt=prompt_ids) + return super().apply(inputs, timing_ctx) def _build_generation_grids(self, hf_processor_mm_kwargs: Mapping[str, object]) -> torch.Tensor: """Build generation grids for M-RoPE decode positions. @@ -661,171 +458,39 @@ def _build_generation_grids(self, hf_processor_mm_kwargs: Mapping[str, object]) def _apply_hf_processor_main( self, - prompt: str | list[int], mm_items: MultiModalDataItems, hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - *, - enable_hf_prompt_update: bool, - ) -> tuple[list[int], BatchFeature, bool]: - """ - Override to handle GLM-Image i2i mode correctly. - - Problem: When vLLM processes cached mm items (enable_hf_prompt_update=False), - the base implementation: - 1. Gets prompt_ids from _apply_hf_processor_text_only (no image tokens) - 2. Gets mm_data from _apply_hf_processor_mm_only - 3. Returns is_update_applied=False - - This causes _apply_prompt_updates to fail because prompt_ids has no image tokens. - - Solution: For i2i mode, we build prompt_ids that include image placeholders, - and return is_update_applied=False so _apply_prompt_updates can expand them. - """ + ) -> BatchFeature: + """Process source images under vLLM's token-only prompt contract.""" mm_counts = mm_items.get_all_counts() num_images = mm_counts.get("image", 0) + if num_images == 0: + return BatchFeature({"mrope_image_grid_thw": self._build_generation_grids(hf_processor_mm_kwargs)}) - logger.debug(f"_apply_hf_processor_main: mm_counts={mm_counts}, num_images={num_images}") - - if num_images == 0 and isinstance(prompt, str): - # t2i mode or normal flow - use parent implementation - prompt_ids = self._apply_hf_processor_text_only( - prompt_text=prompt, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - mm_processed_data = self._apply_hf_processor_mm_only( - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - - # t2i has no source images, so mm features cannot provide image_grid_thw. - # Provide explicit generation grids for M-RoPE to avoid fallback token parsing - # (which can degrade high-resolution spatial positions, e.g. 1920x1920). - try: - mrope_grid_thw = self._build_generation_grids(hf_processor_mm_kwargs) - mm_processed_data["mrope_image_grid_thw"] = mrope_grid_thw - logger.debug( - "_apply_hf_processor_main t2i: mrope_image_grid_thw=%s", - mrope_grid_thw.tolist(), - ) - except Exception as e: - logger.warning("_apply_hf_processor_main t2i: failed to set mrope_image_grid_thw: %s", e) - - return prompt_ids, mm_processed_data, False - - # i2i mode: use unified HF processor path only. - # This avoids drift between duplicated manual/HF i2i implementations. - logger.debug( - "_apply_hf_processor_main: i2i mode (enable_hf_prompt_update=%s), num_images=%s", - enable_hf_prompt_update, - num_images, - ) - - if not isinstance(prompt, str): - # Online OpenAI chat preprocessing can arrive here with tokenized - # prompts (list[int]) before serving_chat replaces engine prompt - # with the clean text prompt. Do not fail the whole request. - logger.warning( - "_apply_hf_processor_main i2i: got tokenized prompt type=%s; " - "using compatibility path for preprocessing", - type(prompt).__name__, - ) - - prompt_ids = list(prompt) - mm_processed_data = self._apply_hf_processor_mm_only( - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) + valid_mm_items = mm_items.select({key for key, count in mm_counts.items() if count > 0}) + processor_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) + images = processor_data.get("images") + if not images: + return BatchFeature(dict(passthrough_data)) - # Preserve full grids for M-RoPE decode (source + target), while - # keeping image_grid_thw source-only for MM batching. - try: - source_grid_thw = mm_processed_data.get("image_grid_thw") - if source_grid_thw is not None and isinstance(source_grid_thw, torch.Tensor): - target_h = ( - hf_processor_mm_kwargs.get("target_h") - if isinstance(hf_processor_mm_kwargs.get("target_h"), int) - else None - ) - target_w = ( - hf_processor_mm_kwargs.get("target_w") - if isinstance(hf_processor_mm_kwargs.get("target_w"), int) - else None - ) - if target_h is None or target_w is None: - target_h = ( - hf_processor_mm_kwargs.get("height") - if isinstance(hf_processor_mm_kwargs.get("height"), int) - else 1024 - ) - target_w = ( - hf_processor_mm_kwargs.get("width") - if isinstance(hf_processor_mm_kwargs.get("width"), int) - else 1024 - ) - - factor = 32 - token_h = max(1, target_h // factor) - token_w = max(1, target_w // factor) - target_grid = torch.tensor([[1, token_h, token_w]], dtype=source_grid_thw.dtype) - mm_processed_data["mrope_image_grid_thw"] = torch.cat([source_grid_thw, target_grid], dim=0) - except Exception: - pass - - # Prompt updates will expand image placeholders in this compatibility path. - return prompt_ids, mm_processed_data, False - - images = mm_items.get_items("image", ImageProcessorItems) - image_list = [images.get(i) for i in range(images.get_count())] - if not image_list: - raise ValueError("GLM-Image i2i requires at least one source image in mm_items") - - hf_inputs = self._call_hf_processor( - prompt=prompt, - mm_data={"images": image_list}, - mm_kwargs=hf_processor_mm_kwargs, - tok_kwargs=tokenization_kwargs, + processor = self.info.get_hf_processor() + image_inputs = processor.image_processor( + images=list(images), + return_tensors="pt", ) - - input_ids = hf_inputs.get("input_ids") - if input_ids is None: - raise ValueError("HF i2i processor returned no input_ids") - - if isinstance(input_ids, torch.Tensor): - prompt_ids = input_ids[0].tolist() if input_ids.dim() > 1 else input_ids.tolist() - else: - prompt_ids = ( - input_ids[0] - if isinstance(input_ids, list) and input_ids and isinstance(input_ids[0], list) - else list(input_ids) + image_grid_thw = image_inputs.get("image_grid_thw") + if image_grid_thw is not None: + source_grid_thw = image_grid_thw[:num_images] + image_inputs["image_grid_thw"] = source_grid_thw + + target_grid = self._build_generation_grids(hf_processor_mm_kwargs)[:1].to(dtype=source_grid_thw.dtype) + image_inputs["mrope_image_grid_thw"] = torch.cat( + [source_grid_thw, target_grid], + dim=0, ) - mm_processed_data = BatchFeature(dict(), tensor_type="pt") - for key in ("pixel_values", "image_grid_thw", "mrope_image_grid_thw"): - value = hf_inputs.get(key) - if value is not None: - mm_processed_data[key] = value - - image_grid_thw = mm_processed_data.get("image_grid_thw") - mrope_grid_thw = mm_processed_data.get("mrope_image_grid_thw") - hf_config = self.info.get_hf_config() - image_token_id = getattr(hf_config, "image_token_id", 167855) - image_token_count = prompt_ids.count(image_token_id) - logger.debug( - "_apply_hf_processor_main i2i(HF): num_images=%s, prompt_len=%s, image_token_count=%s, " - "source_grid_shape=%s, mrope_grid_shape=%s", - num_images, - len(prompt_ids), - image_token_count, - tuple(image_grid_thw.shape) if image_grid_thw is not None else None, - tuple(mrope_grid_thw.shape) if mrope_grid_thw is not None else None, - ) - - # HF processor already expanded image placeholders in input_ids. - return prompt_ids, mm_processed_data, True + image_inputs.update(passthrough_data) + return image_inputs def _get_mm_fields_config( self, @@ -836,7 +501,7 @@ def _get_mm_fields_config( Get the multimodal field configuration. For GLM-Image i2i mode: - - image_grid_thw has been sliced in _call_hf_processor to only include source images + - image_grid_thw contains only source grids from _apply_hf_processor_main - pixel_values has shape [total_patches, C, H, W] - only for source images For t2i mode: @@ -844,7 +509,7 @@ def _get_mm_fields_config( """ result = {} - # Get image_grid_thw if present (already sliced in _call_hf_processor) + # Get the source-only image_grid_thw if present. image_grid_thw = hf_inputs.get("image_grid_thw") if "pixel_values" in hf_inputs and image_grid_thw is not None: @@ -861,45 +526,21 @@ def _get_mm_fields_config( result["pixel_values"] = MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes) - # Register image_grid_thw - it's been sliced in _call_hf_processor + # Register image_grid_thw - it's been sliced in _apply_hf_processor_main # to only include source image grids, so batching will work correctly result["image_grid_thw"] = MultiModalFieldConfig.batched("image") + if "mrope_image_grid_thw" in hf_inputs: + result["mrope_image_grid_thw"] = MultiModalFieldConfig.shared( + "image", + num_source_images, + keep_on_cpu=True, + ) + logger.debug(f"_get_mm_fields_config: result keys: {list(result.keys())}") return result - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - """ - Return whether the HF processor applies prompt updates. - - For GLM-Image i2i mode, the HF processor's apply_chat_template already - expands <|image|> to N tokens (e.g., 4096 for 64x64 grid). - - By returning True, we tell vLLM that HF processor DID apply prompt updates, - so vLLM will use _find_mm_placeholders to locate the expanded tokens - instead of trying to apply replacements. - - For t2i mode (no images), there are no image placeholders to expand. - """ - # Check if we have images (i2i mode) - num_images = mm_items.get_all_counts().get("image", 0) - if num_images > 0: - logger.debug( - f"_hf_processor_applies_updates: returning True for i2i mode " - f"(num_images={num_images}) - HF processor already expanded tokens" - ) - return True - - # For t2i mode (no images), use default behavior - return True - def _get_prompt_updates( self, mm_items: MultiModalDataItems, @@ -909,17 +550,8 @@ def _get_prompt_updates( """ Get prompt updates for image tokens. - For GLM-Image image-to-image mode, the HF processor's apply_chat_template - already expands each <|image|> placeholder to the correct number of - image tokens (grid_h * grid_w tokens per source image). - - The HF processor does: - 1. Replace each <|image|> with num_image_tokens copies of <|placeholder|> - 2. Replace all <|placeholder|> back to <|image|> - - So the tokenized input already has the expanded tokens. We use - target=[image_token_id] to match each occurrence of the image token, - similar to how Qwen2VL handles this pattern. + Expand each source-image placeholder into grid_h * grid_w image tokens + through vLLM's token-based prompt updates. We use image_grid_thw from out_mm_kwargs to get the actual processed grid size, following the Qwen2VL pattern. This is critical because the HF processor @@ -930,8 +562,7 @@ def _get_prompt_updates( """ hf_config = self.info.get_hf_config() - # Get image token ID - this is the token that appears multiple times - # in the tokenized input after HF processor expansion + # Match the source-image placeholder in the tokenized input. image_token_id = getattr(hf_config, "image_token_id", 167855) # Debug: log mm_items info diff --git a/vllm_omni/model_executor/models/glm_tts/glm_tts.py b/vllm_omni/model_executor/models/glm_tts/glm_tts.py index fb3517036d0..e98c227a4fe 100644 --- a/vllm_omni/model_executor/models/glm_tts/glm_tts.py +++ b/vllm_omni/model_executor/models/glm_tts/glm_tts.py @@ -39,7 +39,6 @@ from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, ProcessorInputs, PromptIndexTargets, @@ -51,6 +50,7 @@ from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import Sampler +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor from vllm_omni.model_executor.models.common.nucleus_ras_sampling import ras_sample_one as _ras_sample_one from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform @@ -364,7 +364,7 @@ def get_data_parser(self): ) -class GLMTTSMultiModalProcessor(BaseMultiModalProcessor[GLMTTSMultiModalProcessingInfo]): +class GLMTTSMultiModalProcessor(OmniMultiModalProcessor[GLMTTSMultiModalProcessingInfo]): """GLM-TTS voice-clone processor. Unlike CosyVoice3, GLM-TTS prompt speech tokens are normal Llama vocab IDs @@ -373,6 +373,53 @@ class GLMTTSMultiModalProcessor(BaseMultiModalProcessor[GLMTTSMultiModalProcessi handoff. """ + def apply(self, inputs: ProcessorInputs, timing_ctx): + source_tokenizer = self.info.get_tokenizer() + prompt_text = source_tokenizer.decode( + inputs.prompt, + skip_special_tokens=False, + ) + config = self.info.ctx.get_hf_config() + model_dir = self.info.ctx.model_config.model + self._ensure_cached_runtime_components( + model_dir, + config, + load_voice_clone=False, + ) + + normalized_text = _normalize_glm_tts_processor_text( + self.text_frontend, + prompt_text, + ) + text_ids = self._encode_text(normalized_text) + prompt_parts = [] + if inputs.mm_data_items.get_all_counts().get("audio", 0): + reference_text = inputs.hf_processor_mm_kwargs.get("prompt_text") + if not isinstance(reference_text, str) or not reference_text.strip(): + raise ValueError("GLM-TTS voice cloning requires mm_processor_kwargs['prompt_text'].") + normalized_reference = _normalize_glm_tts_processor_text( + self.text_frontend, + reference_text, + add_trailing_space=True, + ) + prompt_parts.append(self._encode_text(normalized_reference)) + prompt_parts.extend( + [ + text_ids, + torch.tensor([[int(self.special_ids["boa"])]], dtype=torch.long), + ] + ) + prompt_ids = torch.cat(prompt_parts, dim=1).reshape(-1).tolist() + inputs = replace( + inputs, + prompt=prompt_ids, + hf_processor_mm_kwargs={ + **inputs.hf_processor_mm_kwargs, + self._OMNI_PROMPT_TEXT_KEY: prompt_text, + }, + ) + return super().apply(inputs, timing_ctx) + def _ensure_cached_runtime_components( self, model_dir: str, @@ -599,15 +646,6 @@ def _get_mm_fields_config( if key in hf_inputs } - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False - def _cached_apply_hf_processor(self, inputs: ProcessorInputs, timing_ctx: Any): # GLM-TTS builds the actual AR prompt from both the request text and # mm_processor_kwargs["prompt_text"]. The base cache path separates text diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py index 8fbac95f7cf..738e469f134 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py +++ b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py @@ -1669,8 +1669,10 @@ def __init__(self, model, lm_head, config): self.config = config def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - from vllm.model_executor.models.utils import AutoWeightsLoader + from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper skip = ["lm_head."] if getattr(self.config, "tie_word_embeddings", False) else None - loader = AutoWeightsLoader(self, skip_prefixes=skip) - return loader.load_weights(weights) + loader = AutoWeightsLoader(self) + return loader.load_weights( + weights, mapper=WeightsMapper(orig_to_new_prefix={name: None for name in (skip or ())}) + ) diff --git a/vllm_omni/model_executor/models/hunyuan_image3/_hunyuan_v1_vendored.py b/vllm_omni/model_executor/models/hunyuan_image3/_hunyuan_v1_vendored.py new file mode 100644 index 00000000000..cd154650503 --- /dev/null +++ b/vllm_omni/model_executor/models/hunyuan_image3/_hunyuan_v1_vendored.py @@ -0,0 +1,625 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# +# Vendored from vllm/model_executor/models/hunyuan_v1.py at vllm@2cf0a6915c +# (the upstream base vllm-omni origin/main was rebased against). Upstream +# removed the native HunYuan V1 implementation in #53272 / #53615 (migrated +# to the Transformers modeling backend). vllm-omni's HunyuanImage3 model +# subclasses the removed HunYuanModel / HunYuanSparseMoeBlock / HunYuanMLP, +# so those internals are preserved here to keep the omni diffusion model +# working against current upstream. +# Weight loading, expert mapping, and QKV splitting are owned by Omni's +# HunyuanModel subclass; their unused upstream implementations are omitted. +"""Inference-only HunYuan V1 internals vendored for vllm-omni HunyuanImage3.""" + +from itertools import islice + +import torch +from torch import nn +from transformers import PretrainedConfig +from vllm.compilation.decorators import support_torch_compile +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config +from vllm.distributed import ( + get_ep_group, + get_pp_group, + get_tensor_model_parallel_world_size, +) +from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.fused_moe import ( + FusedMoEFactory, +) +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, +) +from vllm.model_executor.models.interfaces import EagleModelMixin +from vllm.model_executor.models.utils import ( + PPMissingLayer, + make_layers, +) +from vllm.sequence import IntermediateTensors +from vllm.v1.attention.backend import AttentionType + + +def _is_moe(config: PretrainedConfig) -> bool: + num_experts = getattr(config, "num_experts", None) + if isinstance(num_experts, int): + return num_experts > 1 + if isinstance(num_experts, list) and num_experts: + # Ensure all elements are integers before calling max. + if all(isinstance(e, int) for e in num_experts): + return max(num_experts) > 1 + else: + return False + return False + + +def _get_cla_factor(config: PretrainedConfig) -> int: + if not getattr(config, "use_cla", False): + return 1 + return getattr(config, "cla_share_factor", 1) + + +class HunYuanMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + bias: bool = False, + prefix: str = "", + reduce_results: bool = True, + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + input_size=hidden_size, + output_sizes=[intermediate_size] * 2, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + input_size=intermediate_size, + output_size=hidden_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + reduce_results=reduce_results, + ) + if hidden_act != "silu": + raise ValueError(f"Unsupported activation: {hidden_act}. Only silu is supported for now.") + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class HunYuanAttention(nn.Module): + def __init__( + self, + config: PretrainedConfig, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + max_position_embeddings: int = 8192, + quant_config: QuantizationConfig | None = None, + bias: bool = False, + cache_config: CacheConfig | None = None, + prefix: str = "", + layer_id: int = -1, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + + if hasattr(config, "head_dim") and config.head_dim: + self.head_dim = config.head_dim + elif hasattr(config, "attention_head_dim"): + self.head_dim = config.attention_head_dim + else: + self.head_dim = self.hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.max_position_embeddings = max_position_embeddings + self.use_qk_norm = getattr(config, "use_qk_norm", False) + self.layer_id = layer_id + + self.qkv_proj = QKVParallelLinear( + hidden_size=hidden_size, + head_size=self.head_dim, + total_num_heads=self.total_num_heads, + total_num_kv_heads=self.total_num_kv_heads, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.o_proj = RowParallelLinear( + input_size=self.total_num_heads * self.head_dim, + output_size=hidden_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.rotary_emb = get_rope( + self.head_dim, + max_position=max_position_embeddings, + rope_parameters=config.rope_parameters, + is_neox_style=True, + ) + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.attn", + ) + + if self.use_qk_norm: + self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + kv_states: tuple[torch.Tensor] | None = None, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + ori_k = k + if self.use_qk_norm: + q = self.query_layernorm( + q.view(-1, self.num_heads, self.head_dim), + ) + k = self.key_layernorm( + k.view(-1, self.num_kv_heads, self.head_dim), + ) + + attn_output = self.attn(q, k, v) + # For o_proj + attn_output = attn_output.view(q.shape[0], -1) + output, _ = self.o_proj(attn_output) + return output, (ori_k, v) + + +class HunYuanCrossAttention(nn.Module): + def __init__( + self, + config: PretrainedConfig, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + max_position_embeddings: int = 8192, + quant_config: QuantizationConfig | None = None, + bias: bool = False, + cache_config: CacheConfig | None = None, + prefix: str = "", + layer_id: int = -1, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + # MistralConfig has an optional head_dim introduced by Mistral-Nemo + if hasattr(config, "head_dim"): + self.head_dim = config.head_dim + elif hasattr(config, "attention_head_dim"): + self.head_dim = config.attention_head_dim + else: + self.head_dim = self.hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.max_position_embeddings = max_position_embeddings + self.use_qk_norm = getattr(config, "use_qk_norm", False) + self.layer_id = layer_id + + self.q_proj = ColumnParallelLinear( + hidden_size, + hidden_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.q_proj", + ) + + self.o_proj = RowParallelLinear( + input_size=self.total_num_heads * self.head_dim, + output_size=hidden_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.rotary_emb = get_rope( + self.head_dim, + max_position=max_position_embeddings, + rope_parameters=config.rope_parameters, + is_neox_style=True, + ) + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.attn", + attn_type=AttentionType.ENCODER_DECODER, + ) + + if self.use_qk_norm: + self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + kv_states: tuple[torch.Tensor] | None = None, + ) -> torch.Tensor: + assert kv_states is not None + ori_k, v = kv_states # use last layer kv, + k = ori_k + q, _ = self.q_proj(hidden_states) + k_tmp = torch.empty_like(k) # Todo: reduant rotary embedding + q, _ = self.rotary_emb(positions, q, k_tmp) + if self.use_qk_norm: + q = self.query_layernorm( + q.view(-1, self.num_heads, self.head_dim), + ) + k = self.key_layernorm( + k.view(-1, self.num_kv_heads, self.head_dim), + ) + + attn_output = self.attn(q, k, v) + # For o_proj + attn_output = attn_output.view(q.shape[0], -1) + output, _ = self.o_proj(attn_output) + return output, (ori_k, v) + + +class HunYuanSparseMoeBlock(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: QuantizationConfig | None = None, + layer_id: int = -1, + prefix: str = "", + enable_eplb: bool = False, + ): + super().__init__() + self.tp_size = get_tensor_model_parallel_world_size() + + self.ep_group = get_ep_group().device_group + self.ep_size = self.ep_group.size() + self.n_routed_experts = config.num_experts + + if self.tp_size > config.num_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than the number of experts {config.num_experts}." + ) + + # Get layer_id topk if config.moe_topk is a list + if isinstance(config.moe_topk, list): + assert layer_id >= 0 + assert len(config.moe_topk) > layer_id + top_k = config.moe_topk[layer_id] + else: + top_k = config.moe_topk + + # If it is moe, moe_intermediate_size is preferred + intermediate_size = config.intermediate_size + if config.moe_intermediate_size is not None: + intermediate_size = ( + config.moe_intermediate_size + if isinstance(config.moe_intermediate_size, int) + else config.moe_intermediate_size[layer_id] + ) + + # Load balancing settings. + vllm_config = get_current_vllm_config() + eplb_config = vllm_config.parallel_config.eplb_config + self.enable_eplb = enable_eplb + + self.n_logical_experts = self.n_routed_experts + self.n_redundant_experts = eplb_config.num_redundant_experts + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts + self.n_local_physical_experts = self.n_physical_experts // self.ep_size + self.gate = ReplicatedLinear( + config.hidden_size, + config.num_experts, + bias=False, + quant_config=None, + prefix=f"{prefix}.gate", + ) + if config.use_mixed_mlp_moe > 0: + # Get layer_id num_shared_expert if config.num_shared_expert is + # a list. + if isinstance(config.num_shared_expert, list): + assert layer_id >= 0 + assert len(config.num_shared_expert) > layer_id + num_shared_expert = config.num_shared_expert[layer_id] + else: + num_shared_expert = config.num_shared_expert + + self.shared_mlp = HunYuanMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size * num_shared_expert, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.shared_mlp", + ) + else: + self.shared_mlp = None + + self.experts = FusedMoEFactory( + shared_experts=self.shared_mlp, + num_experts=self.n_routed_experts, + top_k=top_k, + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + renormalize=top_k > 1, + quant_config=quant_config, + prefix=f"{prefix}.experts", + enable_eplb=self.enable_eplb, + num_redundant_experts=self.n_redundant_experts, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # NOTE: hidden_states can have either 1D or 2D shape. + orig_shape = hidden_states.shape + hidden_dim = hidden_states.shape[-1] + hidden_states = hidden_states.view(-1, hidden_dim) + + # router_logits: (num_tokens, n_experts) + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts(hidden_states=hidden_states, router_logits=router_logits) + + return final_hidden_states.view(orig_shape) + + +class HunYuanDecoderLayer(nn.Module): + def __init__( + self, + config: PretrainedConfig, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + layer_id: int = -1, + enable_eplb: bool = False, + ) -> None: + super().__init__() + assert layer_id >= 0 + self.layer_id = layer_id + self.hidden_size = config.hidden_size + self.intermediate_size = ( + config.intermediate_size + if isinstance(config.intermediate_size, int) + else config.intermediate_size[layer_id] + ) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + attention_bias = getattr(config, "attention_bias", False) or getattr(config, "bias", False) + cla_factor = _get_cla_factor(config) + attention_type = ( + AttentionType.ENCODER_DECODER if layer_id >= 0 and layer_id % cla_factor != 0 else AttentionType.DECODER + ) + if attention_type == AttentionType.DECODER: + self.self_attn = HunYuanAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=getattr(config, "num_key_value_heads", config.num_attention_heads), + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + bias=attention_bias, + cache_config=cache_config, + prefix=f"{prefix}.self_attn", + layer_id=layer_id, + ) + elif attention_type == AttentionType.ENCODER_DECODER: + self.self_attn = HunYuanCrossAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=getattr(config, "num_key_value_heads", config.num_attention_heads), + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + bias=attention_bias, + cache_config=cache_config, + prefix=f"{prefix}.self_attn", + layer_id=layer_id, + ) + else: + raise RuntimeError(f"Unsupported attention type: {attention_type}") + + if _is_moe(config): + self.mlp = HunYuanSparseMoeBlock( + config=config, + quant_config=quant_config, + layer_id=layer_id, + prefix=f"{prefix}.mlp", + enable_eplb=enable_eplb, + ) + else: + self.mlp = HunYuanMLP( + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + bias=getattr(config, "mlp_bias", False), + prefix=f"{prefix}.mlp", + ) + + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + kv_states: tuple[torch.Tensor] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states, ori_kv_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + kv_states=kv_states, + ) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual, ori_kv_states + + +@support_torch_compile( + dynamic_arg_dims={ + "input_ids": 0, + # positions is of shape (xd, seq_len) if xdrope is enabled for hunyuan-vl, + # otherwise (seq_len, ). + "positions": -1, + "intermediate_tensors": 0, + "inputs_embeds": 0, + } +) +class HunYuanModel(nn.Module, EagleModelMixin): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + + eplb_config = vllm_config.parallel_config.eplb_config + enable_eplb = vllm_config.parallel_config.enable_eplb + self.num_redundant_experts = eplb_config.num_redundant_experts + + self.config = config + self.quant_config = quant_config + + self.vocab_size = config.vocab_size + + if get_pp_group().is_first_rank or (config.tie_word_embeddings and get_pp_group().is_last_rank): + self.embed_tokens = VocabParallelEmbedding( + self.vocab_size, + config.hidden_size, + quant_config=quant_config, + ) + else: + self.embed_tokens = PPMissingLayer() + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: HunYuanDecoderLayer( + config=config, + layer_id=int(prefix.split(".")[-1]), + cache_config=cache_config, + quant_config=quant_config, + prefix=prefix, + enable_eplb=enable_eplb, + ), + prefix=f"{prefix}.layers", + ) + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer() + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + + cla_factor = _get_cla_factor(self.config) + prev_kv_states = None + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual) + for i, layer in enumerate(islice(self.layers, self.start_layer, self.end_layer)): + hidden_states, residual, kv_states = layer( + positions, + hidden_states, + residual, + prev_kv_states, + ) + + if getattr(self.config, "use_cla", False) and i % cla_factor == 0: + prev_kv_states = kv_states + else: + prev_kv_states = None + + self._maybe_add_hidden_state(aux_hidden_states, i + 1, hidden_states, residual) + + if not get_pp_group().is_last_rank: + return IntermediateTensors({"hidden_states": hidden_states, "residual": residual}) + + hidden_states, _ = self.norm(hidden_states, residual) + + if len(aux_hidden_states) > 0: + return hidden_states, aux_hidden_states + return hidden_states diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py index 58b67d8aaa4..b15a19982ae 100755 --- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py +++ b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.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 gc import math import typing @@ -45,13 +45,6 @@ default_weight_loader, maybe_remap_kv_scale_name, ) -from vllm.model_executor.models.hunyuan_v1 import ( - HunYuanMLP, - HunYuanModel, - HunYuanSparseMoeBlock, - _get_cla_factor, - _is_moe, -) from vllm.model_executor.models.interfaces import ( MultiModalEmbeddings, SupportsLoRA, @@ -63,6 +56,7 @@ from vllm.model_executor.models.utils import ( AutoWeightsLoader, PPMissingLayer, + WeightsMapper, _merge_multimodal_embeddings, is_pp_missing_parameter, maybe_prefix, @@ -78,7 +72,6 @@ ) from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, PromptReplacement, PromptUpdate, @@ -91,6 +84,14 @@ from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import Sampler +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor +from vllm_omni.model_executor.models.hunyuan_image3._hunyuan_v1_vendored import ( + HunYuanMLP, + HunYuanModel, + HunYuanSparseMoeBlock, + _get_cla_factor, + _is_moe, +) from vllm_omni.model_executor.models.hunyuan_image3.autoencoder_kl_3d import AutoencoderKLConv3D from vllm_omni.model_executor.models.hunyuan_image3.siglip2 import LightProjector, Siglip2VisionTransformer @@ -1027,7 +1028,7 @@ def get_dummy_mm_data( } -class HunyuanImage3MultiModalProcessor(BaseMultiModalProcessor[HunyuanImage3ProcessingInfo]): +class HunyuanImage3MultiModalProcessor(OmniMultiModalProcessor[HunyuanImage3ProcessingInfo]): """Multimodal processor for HunyuanImage3 model.""" def _call_hf_processor( @@ -1046,15 +1047,6 @@ def _call_hf_processor( batch_feature["vae_generator_seed"] = torch.full((len(images),), int(vae_generator_seed), dtype=torch.long) return batch_feature - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False - def _get_mm_fields_config( self, hf_inputs: BatchFeature, @@ -2180,12 +2172,11 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: "timestep_r_emb", ] skip_prefixes.extend(unexpected_keywords) - loader = AutoWeightsLoader( - self, - skip_prefixes=skip_prefixes, - ) + loader = AutoWeightsLoader(self) - loaded_params = loader.load_weights(weights) + loaded_params = loader.load_weights( + weights, mapper=WeightsMapper(orig_to_new_prefix={name: None for name in (skip_prefixes or ())}) + ) return loaded_params def get_language_model(self) -> torch.nn.Module: diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py index 4309d617796..4dbc9a4556d 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + # Copyright 2025 Xiaomi Corporation. import os from collections.abc import Iterable, Mapping, Sequence @@ -27,14 +30,13 @@ from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, ProcessorInputs, PromptReplacement, PromptUpdate, PromptUpdateDetails, ) -from vllm.multimodal.utils import group_mm_kwargs_by_modality +from vllm.multimodal.utils import group_and_batch_mm_kwargs from vllm.sequence import IntermediateTensors from vllm.utils.cache import LRUCache from vllm.utils.collection_utils import is_list_of @@ -42,6 +44,7 @@ from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import Sampler +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import ( NO_INTERLEAVE_NEXT_TOKEN_ID, @@ -348,7 +351,7 @@ def _parse_audio_data( return AudioProcessorItems(new_audios) -class MiMoAudioLLMMultiModalProcessor(BaseMultiModalProcessor[MiMoAudioLLMProcessingInfo]): +class MiMoAudioLLMMultiModalProcessor(OmniMultiModalProcessor[MiMoAudioLLMProcessingInfo]): def _call_hf_processor( self, prompt: str, @@ -484,7 +487,7 @@ def get_replacement_mimo_audio(item_idx: int): return [ PromptReplacement( modality="audio", - target=audio_token, + target=[audio_token_id], replacement=get_replacement_mimo_audio, ) ] @@ -626,7 +629,7 @@ def fused_thinker_talker_prefill(self, input_ids: torch.Tensor, input_embeds: to ) mm_kwargs.append((mm_feature.modality, mm_item)) - for modality, num_items, mm_kwargs_group in group_mm_kwargs_by_modality( + for modality, num_items, mm_kwargs_group in group_and_batch_mm_kwargs( mm_kwargs, device=self.device, pin_memory=self.pin_memory, diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py index c90adb4676f..bb01e522c6c 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py @@ -390,35 +390,25 @@ def get_dummy_mm_data( class MimoAudioMultiModalProcessor(BaseMultiModalProcessor[MimoAudioProcessingInfo]): - def _call_hf_processor( + def _preprocess_hf_mm_data( self, - prompt: str, mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, Any], - tok_kwargs: Mapping[str, object], - ) -> BatchFeature: + hf_processor_mm_kwargs: Mapping[str, object], + ) -> tuple[Mapping[str, object], Mapping[str, object]]: + mm_data, hf_processor_mm_kwargs = super()._preprocess_hf_mm_data( + mm_data, + hf_processor_mm_kwargs, + ) + mm_data = dict(mm_data) audios = mm_data.pop("audios", []) if audios: mm_data["audio"] = audios - - # Text-only input not supported in composite processor - if not mm_data.get("audio", []): - prompt_ids = self.info.get_tokenizer().encode(prompt) - prompt_ids = self._apply_hf_processor_tokens_only(prompt_ids) - return BatchFeature(dict(input_ids=[prompt_ids]), tensor_type="pt") - - feature_extractor = self.info.get_feature_extractor(**mm_kwargs) - mm_kwargs = dict( - **mm_kwargs, - sampling_rate=feature_extractor.sampling_rate, - ) - - return super()._call_hf_processor( - prompt=prompt, - mm_data=mm_data, - mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, - ) + feature_extractor = self.info.get_feature_extractor(**hf_processor_mm_kwargs) + hf_processor_mm_kwargs = { + **hf_processor_mm_kwargs, + "sampling_rate": feature_extractor.sampling_rate, + } + return mm_data, hf_processor_mm_kwargs def _get_mm_fields_config( self, @@ -480,7 +470,7 @@ def get_replacement_mimo_audio(item_idx: int): return [ PromptReplacement( modality="audio", - target=audio_token, + target=[audio_token_id], replacement=get_replacement_mimo_audio, ) ] diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py index d7e31beb300..d34cce46542 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright 2025 The vLLM-Omni team. # Copyright (c) Ant Group. All rights reserved. # Adapted from: @@ -21,7 +22,7 @@ from transformers.utils.hub import cached_file from vllm.config import VllmConfig from vllm.logger import init_logger -from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin @@ -509,12 +510,13 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if self._standalone: weights = self._iter_talker_safetensors() - loader = AutoWeightsLoader( - self, - skip_prefixes=["audio_vae."], # loaded separately - skip_substrs=["rotary_embed.inv_freq"], # non-persistent buffer + loader = AutoWeightsLoader(self) + loaded = loader.load_weights( + weights, + mapper=WeightsMapper( + orig_to_new_prefix={"audio_vae.": None}, orig_to_new_substr={"rotary_embed.inv_freq": None} + ), ) - loaded = loader.load_weights(weights) logger.info("Loaded %d talker weights from checkpoint", len(loaded)) if self.audio_vae is not None and self._vae_weight_source is not None: diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py index 03cf87abb4d..1338f26060f 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Copyright 2025 The vLLM-Omni team. # Copyright 2024 ANT Group and the HuggingFace Inc. team. # Adapted from Ming repository modeling_bailingmm2.py and processing_bailingmm2.py @@ -8,6 +9,7 @@ import os from collections.abc import Iterable, Iterator, Mapping, Sequence +from dataclasses import replace from typing import Annotated, Any import numpy as np @@ -388,7 +390,10 @@ def get_replacement_audio(item_idx: int) -> PromptUpdateDetails: updates.append( PromptReplacement( modality="image", - target=PLACEHOLDER_IMAGE_TOKEN_IN_TEXT, + target=tokenizer.encode( + PLACEHOLDER_IMAGE_TOKEN_IN_TEXT, + add_special_tokens=False, + ), replacement=get_replacement_image, ) ) @@ -396,7 +401,10 @@ def get_replacement_audio(item_idx: int) -> PromptUpdateDetails: updates.append( PromptReplacement( modality="video", - target=PLACEHOLDER_VIDEO_TOKEN_IN_TEXT, + target=tokenizer.encode( + PLACEHOLDER_VIDEO_TOKEN_IN_TEXT, + add_special_tokens=False, + ), replacement=get_replacement_video, ) ) @@ -404,7 +412,10 @@ def get_replacement_audio(item_idx: int) -> PromptUpdateDetails: updates.append( PromptReplacement( modality="audio", - target=PLACEHOLDER_AUDIO_TOKEN_IN_TEXT, + target=tokenizer.encode( + PLACEHOLDER_AUDIO_TOKEN_IN_TEXT, + add_special_tokens=False, + ), replacement=get_replacement_audio, ) ) @@ -451,90 +462,71 @@ def _get_mm_fields_config( return config - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False - - def _apply_hf_processor_main( - self, - prompt: str | list[int], - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - *, - enable_hf_prompt_update: bool, - ): - """Apply Ming-specific prompt rewrites before delegating to vllm's - default text / multimodal processing pipeline. - - This is the single point in the base MM-processor call graph where we - still have BOTH the full prompt and the request-level - ``hf_processor_mm_kwargs`` together. Downstream, vllm's MM caching - design splits the work into - :meth:`_apply_hf_processor_text_only` (prompt, empty kwargs) and - :meth:`_apply_hf_processor_mm_only` (dummy text, real kwargs), so a - per-request flag like ``modalities`` cannot reach the real - tokenization from inside :meth:`_call_hf_processor`. - - For Ming-flash-omni-2.0, image-generation requests carry - ``mm_processor_kwargs["modalities"]`` set to ``["image"]`` (t2i) or - ``["img2img"]`` (image edit). We expand the prompt here with the - ``*N`` query-token block so the expanded - form flows through both the tokenization and (eventual) MM paths. - For img2img, we also prepend ```` so the thinker-side prompt - replacement can locate the reference-image placeholder. - """ - modalities = hf_processor_mm_kwargs.get("modalities") or [] - is_image_gen = "image" in modalities or "img2img" in modalities - if isinstance(prompt, str) and is_image_gen: + def apply(self, inputs, timing_ctx): + """Add Ming image-generation query tokens before prompt updates.""" + modalities = inputs.hf_processor_mm_kwargs.get("modalities") or [] + if "image" in modalities or "img2img" in modalities: from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import ( DEFAULT_NUM_QUERY_TOKENS, - maybe_expand_image_gen_prompt, + IMAGE_PATCH_TOKEN, ) - # img2img: prepend ```` so the ref-image placeholder survives - # MM caching (when the ref image is cache-warm, the thinker sees 0 - # missing items and cannot inject the placeholder itself). + tokenizer = self.info.get_tokenizer() + prompt_ids = list(inputs.prompt) if "img2img" in modalities: - prompt = "" + prompt - - ig = getattr(self.info.ctx.model_config.hf_config, "image_gen_config", None) - num_query_tokens = getattr(ig, "num_query_tokens", DEFAULT_NUM_QUERY_TOKENS) - prompt = maybe_expand_image_gen_prompt(prompt, num_query_tokens=int(num_query_tokens)) - - return super()._apply_hf_processor_main( - prompt=prompt, - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - enable_hf_prompt_update=enable_hf_prompt_update, - ) + prompt_ids = ( + tokenizer.encode( + PLACEHOLDER_IMAGE_TOKEN_IN_TEXT, + add_special_tokens=False, + ) + + prompt_ids + ) - def _call_hf_processor( + image_patch_ids = tokenizer.encode( + IMAGE_PATCH_TOKEN, + add_special_tokens=False, + ) + if not image_patch_ids or not any( + prompt_ids[index : index + len(image_patch_ids)] == image_patch_ids + for index in range(len(prompt_ids) - len(image_patch_ids) + 1) + ): + image_gen_config = getattr( + self.info.ctx.model_config.hf_config, + "image_gen_config", + None, + ) + num_query_tokens = int( + getattr( + image_gen_config, + "num_query_tokens", + DEFAULT_NUM_QUERY_TOKENS, + ) + ) + suffix = ( + tokenizer.encode("", add_special_tokens=False) + + image_patch_ids * num_query_tokens + + tokenizer.encode("", add_special_tokens=False) + ) + prompt_ids.extend(suffix) + + inputs = replace(inputs, prompt=prompt_ids) + + return super().apply(inputs, timing_ctx) + + def _apply_hf_processor_main( self, - prompt: str, - mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], ) -> BatchFeature: - """Call sub-processors for multimodal inputs and tokenize. + """Call Ming's modality processors without expanding prompt text. We call the image/audio sub-processors directly (instead of going through `MingFlashOmniProcessor.__call__`) so that the high-level - placeholder tokens remain **unexpanded** in the tokenized output. - - NOTE: Image-gen prompt expansion is handled in - :meth:`_apply_hf_processor_main` — by the time we get here, vllm's - MM caching has already split text/mm processing and ``mm_kwargs`` - is empty on the text-only branch. + placeholder tokens remain under vLLM's prompt-update control. """ + valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0}) + mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) hf_processor = self.info.get_hf_processor() - tokenizer = self.info.get_tokenizer() data: dict[str, object] = {} @@ -577,11 +569,9 @@ def _call_hf_processor( ) data.update(audio_outputs) - # Tokenize text with placeholders still intact - text_outputs = tokenizer(prompt, return_tensors="pt", **tok_kwargs) - data.update(text_outputs) - - return BatchFeature(data=data) + processed_data = BatchFeature(data=data) + processed_data.update(passthrough_data) + return processed_data @MULTIMODAL_REGISTRY.register_processor( diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index 9dfff5589c7..315ae159d3c 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -21,9 +21,8 @@ import math import os import warnings -from collections import defaultdict from collections.abc import Callable, Iterable, Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace from functools import partial from typing import Annotated, Any, Literal, TypeAlias @@ -3403,84 +3402,46 @@ def _parse_audio_data( class MiniCPMO45OmniLLMMultiModalProcessor(BaseMultiModalProcessor[MiniCPMO45OmniLLMProcessingInfo]): """Multimodal processor for MiniCPM-o thinker stage.""" + _TTS_SUFFIX = "<|spk_bos|><|spk|><|spk_eos|><|tts_bos|>" + + def apply(self, inputs, timing_ctx): + hf_processor_mm_kwargs = inputs.hf_processor_mm_kwargs + if hf_processor_mm_kwargs.get("use_tts", False): + tokenizer = self.info.get_tokenizer() + tts_ids = tokenizer.convert_tokens_to_ids(["<|spk_bos|>", "<|spk|>", "<|spk_eos|>", "<|tts_bos|>"]) + inputs = replace( + inputs, + prompt=[*inputs.prompt, *tts_ids], + hf_processor_mm_kwargs={ + key: value for key, value in hf_processor_mm_kwargs.items() if key != "use_tts" + }, + ) + return super().apply(inputs, timing_ctx) + def _apply_hf_processor_main( self, - prompt: str | list[int], mm_items: MultiModalDataItems, hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - *, - enable_hf_prompt_update: bool, - ) -> tuple[list[int], BatchFeature, bool]: + ) -> BatchFeature: """ - MiniCPM-O reimplements this to avoid calling HF processor with text-only - when enable_hf_prompt_update=False. The MiniCPM processor asserts - len(image_tags) == len(image_sizes) and fails if given placeholder text - without corresponding image data (e.g. during profiling/cache-miss path). + Process each modality independently because the MiniCPM processor + asserts that image tags and image sizes have matching lengths. """ - use_tts = hf_processor_mm_kwargs.get("use_tts", False) - hf_processor_mm_kwargs = {k: v for k, v in hf_processor_mm_kwargs.items() if k != "use_tts"} - if isinstance(prompt, str): - if use_tts: - prompt = prompt + self._TTS_SUFFIX - if enable_hf_prompt_update: - return self._apply_hf_processor_text_mm( - prompt_text=prompt, - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - tokenizer = self.info.get_tokenizer() - prompt_ids = _encode_tokens(tokenizer, prompt) - else: - prompt_ids = self._apply_hf_processor_tokens_only(prompt) - if use_tts: - tokenizer = self.info.get_tokenizer() - tts_ids = tokenizer.convert_tokens_to_ids(["<|spk_bos|>", "<|spk|>", "<|spk_eos|>", "<|tts_bos|>"]) - prompt_ids = list(prompt_ids) + tts_ids - - mm_processed_data = self._apply_hf_processor_mm_only( - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - - return prompt_ids, mm_processed_data, False - - _TTS_SUFFIX = "<|spk_bos|><|spk|><|spk_eos|><|tts_bos|>" + valid_mm_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0}) + mm_data, passthrough_data = self._get_hf_mm_data(valid_mm_items) - def _call_hf_processor( - self, - prompt: str, - mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], - ) -> BatchFeature: tokenizer = self.info.get_tokenizer() - - use_tts = mm_kwargs.get("use_tts", False) - mm_kwargs = {k: v for k, v in mm_kwargs.items() if k != "use_tts"} - if use_tts: - prompt = prompt + self._TTS_SUFFIX - - input_ids = torch.tensor([tokenizer.encode(prompt, **tok_kwargs)]) - mm_inputs = self.process_mm_inputs(mm_data, mm_kwargs, tok_kwargs) - - return BatchFeature( + prompt_text = self.dummy_inputs.get_dummy_text(mm_items.get_all_counts()) + input_ids = torch.tensor([tokenizer.encode(prompt_text)]) + mm_inputs = self.process_mm_inputs(mm_data, hf_processor_mm_kwargs) + processed_data = BatchFeature( { "input_ids": input_ids, **mm_inputs, } ) - - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False + processed_data.update(passthrough_data) + return processed_data def get_image_prompt_texts( self, @@ -3528,7 +3489,6 @@ def process_audios( self, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> Mapping[str, NestedTensors]: if (audios := mm_data.get("audios")) is None: return {} @@ -3540,11 +3500,10 @@ def process_audios( if isinstance(parsed_audios, MiniCPMOAudioEmbeddingItems): audio_inputs = {} else: - audio_inputs = self._base_call_hf_processor( + audio_inputs = self._call_hf_processor_on_prompts( prompts=[self.info.audio_pattern] * len(parsed_audios), mm_data={"audios": [[audio] for audio in parsed_audios]}, mm_kwargs={**mm_kwargs, "chunk_input": True}, - tok_kwargs=tok_kwargs, out_keys={"audio_features", "audio_feature_lens"}, ) @@ -3569,7 +3528,6 @@ def process_images( self, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> Mapping[str, NestedTensors]: if (images := mm_data.get("images")) is None: return {} @@ -3581,11 +3539,10 @@ def process_images( if isinstance(parsed_images, MiniCPMVImageEmbeddingItems): image_inputs = {} else: - image_inputs = self._base_call_hf_processor( + image_inputs = self._call_hf_processor_on_prompts( prompts=[self.info.image_pattern] * len(parsed_images), mm_data={"images": [[image] for image in parsed_images]}, mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, out_keys={"pixel_values", "image_sizes", "tgt_sizes"}, ) return image_inputs @@ -3594,7 +3551,6 @@ def process_videos( self, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> Mapping[str, NestedTensors]: if (videos := mm_data.get("videos")) is None: return {} @@ -3606,14 +3562,13 @@ def process_videos( if isinstance(parsed_videos, MiniCPMVVideoEmbeddingItems): video_inputs = {} else: - video_inputs = self._base_call_hf_processor( + video_inputs = self._call_hf_processor_on_prompts( prompts=[self.info.image_pattern * len(video) for video in parsed_videos], mm_data={"images": list(parsed_videos)}, mm_kwargs={ **mm_kwargs, "max_slice_nums": self.info.get_video_max_slice_num(), }, - tok_kwargs=tok_kwargs, out_keys={"pixel_values", "image_sizes", "tgt_sizes"}, ) @@ -3625,48 +3580,30 @@ def process_mm_inputs( self, mm_data: Mapping[str, object], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], ) -> Mapping[str, NestedTensors]: return { - **self.process_images(mm_data, mm_kwargs, tok_kwargs), - **self.process_videos(mm_data, mm_kwargs, tok_kwargs), - **self.process_audios(mm_data, mm_kwargs, tok_kwargs), + **self.process_images(mm_data, mm_kwargs), + **self.process_videos(mm_data, mm_kwargs), + **self.process_audios(mm_data, mm_kwargs), } - def _base_call_hf_processor( + def _call_hf_processor_on_prompts( self, prompts: list[str], mm_data: Mapping[str, Sequence[object]], mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], *, out_keys: set[str], ) -> dict[str, NestedTensors]: - mm_kwargs = {k: v for k, v in mm_kwargs.items() if k != "use_tts"} - # This processor supports zipping prompt and mm_data together - if self.info.get_model_version() in {(2, 6), (4, 0), (4, 5)}: - inputs = super()._call_hf_processor( - prompt=prompts, # type: ignore - mm_data=mm_data, - mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, - ) - else: - inputs = defaultdict[str, list[torch.Tensor]](list) - - for i, prompt in enumerate(prompts): - inputs_one = super()._call_hf_processor( - prompt=prompt, - mm_data={k: v[i] for k, v in mm_data.items()}, - mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, - ) - - for k, v in inputs_one.items(): - assert len(v) == 1, (k, len(v)) - inputs[k].append(v[0]) - - return {k: inputs[k] for k in out_keys} + from vllm.model_executor.models.minicpmv import MiniCPMVMultiModalProcessor + + return MiniCPMVMultiModalProcessor._call_hf_processor_on_prompts( + self, + prompts, + mm_data, + {key: value for key, value in mm_kwargs.items() if key != "use_tts"}, + out_keys=out_keys, + ) def _get_prompt_updates( self, @@ -3688,6 +3625,17 @@ def _get_prompt_updates( additional_placeholders.append((modality, sub_pattern)) placeholders += additional_placeholders + # vLLM 0.29 removed PromptUpdateDetails.select_text: prompt updates are + # token oriented now, so encode the replacement text and select the unk + # placeholder positions by token id instead of by text. + unk_token_id = tokenizer.convert_tokens_to_ids("") + + def _select_unk_positions(text: str) -> PromptUpdateDetails: + return PromptUpdateDetails.select_token_id( + tokenizer.encode(text, add_special_tokens=False), + unk_token_id, + ) + image_max_slice_nums = hf_processor_mm_kwargs.get("max_slice_nums") image_use_image_id = hf_processor_mm_kwargs.get("use_image_id") @@ -3696,14 +3644,13 @@ def get_image_replacement(item_idx: int): image_size = images.get_image_size(item_idx) - return PromptUpdateDetails.select_text( + return _select_unk_positions( self.get_image_prompt_texts( image_size, item_idx, max_slice_nums=None if image_max_slice_nums is None else int(image_max_slice_nums), # type: ignore[arg-type] use_image_id=None if image_use_image_id is None else bool(image_use_image_id), - ), - "", + ) ) def get_video_replacement(item_idx: int): @@ -3712,17 +3659,18 @@ def get_video_replacement(item_idx: int): frame_size = videos.get_frame_size(item_idx) num_frames = videos.get_num_frames(item_idx) - return PromptUpdateDetails.select_text( - self.get_video_prompt_texts(frame_size, num_frames), - "", - ) + return _select_unk_positions(self.get_video_prompt_texts(frame_size, num_frames)) get_replacement = { "image": get_image_replacement, "video": get_video_replacement, } base_updates = [ - PromptReplacement(modality=modality, target=pattern, replacement=get_replacement[modality]) + PromptReplacement( + modality=modality, + target=tokenizer.encode(pattern, add_special_tokens=False), + replacement=get_replacement[modality], + ) for modality, pattern in placeholders ] @@ -3738,16 +3686,16 @@ def get_audio_replacement(item_idx: int): else: audio_len = audios.get_audio_length(item_idx) - return PromptUpdateDetails.select_text( - self.get_audio_prompt_texts(audio_len), - "", - ) + return _select_unk_positions(self.get_audio_prompt_texts(audio_len)) return [ *base_updates, PromptReplacement( modality="audio", - target=audio_placeholder, + target=tokenizer.encode( + audio_placeholder, + add_special_tokens=False, + ), replacement=get_audio_replacement, ), ] diff --git a/vllm_omni/model_executor/models/minimax_h3/text_encoder.py b/vllm_omni/model_executor/models/minimax_h3/text_encoder.py index 85583f24331..b1933af95dc 100644 --- a/vllm_omni/model_executor/models/minimax_h3/text_encoder.py +++ b/vllm_omni/model_executor/models/minimax_h3/text_encoder.py @@ -12,6 +12,7 @@ import torch import torch.nn as nn from vllm.config import VllmConfig +from vllm.inputs import mm_input from vllm.model_executor.models.qwen3_vl import ( Qwen3VLDummyInputsBuilder, Qwen3VLForConditionalGeneration, @@ -26,10 +27,6 @@ MINIMAX_H3_PRESENTATION_TASK_KEY, ) from vllm_omni.model_executor.models.minimax_h3.preprocessing import ( - IMAGE_PAD, - VIDEO_PAD, - VISION_END, - VISION_START, build_minimax_h3_presentation, ) from vllm_omni.model_executor.models.output_templates import OmniOutput @@ -96,58 +93,57 @@ def _cached_apply_hf_processor(self, inputs: Any, timing_ctx: Any): def _apply_hf_processor_main( self, - prompt: str | list[int], mm_items: Any, hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - *, - enable_hf_prompt_update: bool, ): - task_value = hf_processor_mm_kwargs.get(MINIMAX_H3_PRESENTATION_TASK_KEY) - if task_value is None: - return super()._apply_hf_processor_main( - prompt, - mm_items, - hf_processor_mm_kwargs, - tokenization_kwargs, - enable_hf_prompt_update=enable_hf_prompt_update, - ) - if not isinstance(prompt, str): - raise ValueError("MiniMax H3 presentation requires the original prompt text") - - task = str(task_value) - condition_labels = self._condition_labels(hf_processor_mm_kwargs.get(MINIMAX_H3_CONDITION_LABELS_KEY)) base_kwargs = self._base_processor_kwargs(hf_processor_mm_kwargs) - image_count = sum(kind == "image" for kind, _ in condition_labels) - video_count = sum(kind == "video" for kind, _ in condition_labels) - image_slot = f"{VISION_START}{IMAGE_PAD}{VISION_END}" - hf_processor = self.info.get_hf_processor(**base_kwargs) - video_slot = ( - VIDEO_PAD if self._expands_only_video_token(hf_processor) else f"{VISION_START}{VIDEO_PAD}{VISION_END}" - ) - processing_prompt = image_slot * image_count + video_slot * video_count - if not processing_prompt: - processing_prompt = prompt - - _, processed_data, _ = super()._apply_hf_processor_main( - processing_prompt, + return super()._apply_hf_processor_main( mm_items, base_kwargs, - tokenization_kwargs, - enable_hf_prompt_update=enable_hf_prompt_update, ) + + def apply(self, inputs: Any, timing_ctx: Any): + task_value = inputs.hf_processor_mm_kwargs.get(MINIMAX_H3_PRESENTATION_TASK_KEY) + if task_value is None: + return super().apply(inputs, timing_ctx) + + prompt_ids = self._postprocess_prompt(inputs.prompt) + mm_info = self._cached_apply_hf_processor(inputs, timing_ctx) + base_kwargs = self._base_processor_kwargs(inputs.hf_processor_mm_kwargs) + tokenizer = self.info.get_tokenizer() + prompt_text = tokenizer.decode(prompt_ids, skip_special_tokens=False) + condition_labels = self._condition_labels(inputs.hf_processor_mm_kwargs.get(MINIMAX_H3_CONDITION_LABELS_KEY)) + out_mm_data = mm_info.kwargs.get_data() image_processor = self.info.get_image_processor(**base_kwargs) ids, _ = _build_minimax_h3_presentation( - self.info.get_tokenizer(), - prompt=prompt, - task=task, + tokenizer, + prompt=prompt_text, + task=str(task_value), condition_labels=condition_labels, - image_grid_thw=processed_data.get("image_grid_thw"), - video_grid_thw=processed_data.get("video_grid_thw"), - video_timestamps=processed_data.get("timestamps"), + image_grid_thw=out_mm_data.get("image_grid_thw"), + video_grid_thw=out_mm_data.get("video_grid_thw"), + video_timestamps=out_mm_data.get("timestamps"), merge_size=int(image_processor.merge_size), ) - return ids.tolist(), processed_data, True + prompt_ids = ids.tolist() + + mm_item_counts = inputs.mm_data_items.get_all_counts() + self._validate_mm_kwargs(mm_info.kwargs, mm_item_counts) + self._validate_mm_updates(mm_info.prompt_updates, mm_item_counts) + mm_placeholders = self._find_mm_placeholders( + prompt_ids, + mm_info.prompt_updates, + ) + self._validate_mm_placeholders(mm_placeholders, mm_item_counts) + mm_placeholder_ranges = { + modality: [item.to_range() for item in placeholders] for modality, placeholders in mm_placeholders.items() + } + return mm_input( + prompt_token_ids=prompt_ids, + mm_kwargs=mm_info.kwargs, + mm_hashes=mm_info.hashes, + mm_placeholders=mm_placeholder_ranges, + ) def _get_prompt_updates( self, diff --git a/vllm_omni/model_executor/models/minimax_music3/talker.py b/vllm_omni/model_executor/models/minimax_music3/talker.py index c4fb876faf4..9baa7315c1a 100644 --- a/vllm_omni/model_executor/models/minimax_music3/talker.py +++ b/vllm_omni/model_executor/models/minimax_music3/talker.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 """Stage-0 talker for MiniMax Music 3. A Qwen3 backbone predicts one audio frame per decode step. Each frame is @@ -37,7 +37,7 @@ 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 +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper from vllm.v1.outputs import SamplerOutput from vllm_omni.model_executor.models.output_templates import OmniOutput @@ -980,8 +980,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: through vLLM. The audio embedding table and the depth decoder live in a separate component of the repo and are loaded explicitly. """ - loader = AutoWeightsLoader(self, skip_prefixes=["audio_embeddings.", "rvq_decoder."]) - loaded = loader.load_weights(weights) + loader = AutoWeightsLoader(self) + loaded = loader.load_weights( + weights, mapper=WeightsMapper(orig_to_new_prefix={"audio_embeddings.": None, "rvq_decoder.": None}) + ) loaded |= load_depth_decoder_weights( self.vllm_config, audio_embeddings=self.audio_embeddings, diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice.py b/vllm_omni/model_executor/models/omnivoice/omnivoice.py index 26a6f6b5935..a0b0d725e76 100644 --- a/vllm_omni/model_executor/models/omnivoice/omnivoice.py +++ b/vllm_omni/model_executor/models/omnivoice/omnivoice.py @@ -11,6 +11,7 @@ import os from collections.abc import Iterable, Mapping, Sequence +from dataclasses import replace import numpy as np import torch @@ -28,7 +29,6 @@ from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, - BaseMultiModalProcessor, BaseProcessingInfo, ProcessorInputs, PromptIndexTargets, @@ -37,6 +37,7 @@ ) from vllm.sequence import IntermediateTensors +from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig @@ -63,13 +64,29 @@ def get_data_parser(self): ) -class OmniVoiceMultiModalProcessor(BaseMultiModalProcessor[OmniVoiceMultiModalProcessingInfo]): +class OmniVoiceMultiModalProcessor(OmniMultiModalProcessor[OmniVoiceMultiModalProcessingInfo]): """Processes text + optional reference audio for OmniVoice. For voice cloning: text + reference audio → tokenized reference For auto voice: text only """ + def apply(self, inputs: ProcessorInputs, timing_ctx): + tokenizer = self.info.get_tokenizer() + prompt_text = tokenizer.decode(inputs.prompt, skip_special_tokens=False) + # Tokenize once here; _apply_hf_processor_main below only encodes the + # reference audio, so the media path never re-tokenizes the text. + prompt_ids = self._encode_prompt_text(prompt_text, inputs.hf_processor_mm_kwargs).reshape(-1).tolist() + inputs = replace( + inputs, + prompt=prompt_ids, + hf_processor_mm_kwargs={ + **inputs.hf_processor_mm_kwargs, + self._OMNI_PROMPT_TEXT_KEY: prompt_text, + }, + ) + return super().apply(inputs, timing_ctx) + def _ensure_cached_runtime_components(self, model_dir: str, config: OmniVoiceConfig) -> None: cached_model_dir = getattr(self, "_cached_model_dir", None) if cached_model_dir == model_dir: @@ -97,22 +114,38 @@ def _ensure_cached_runtime_components(self, model_dir: str, config: OmniVoiceCon self._cached_model_dir = model_dir - def _call_hf_processor( + def _apply_hf_processor_main( self, - prompt: str, - mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], ) -> BatchFeature: - config = self.info.ctx.get_hf_config() - model_dir = self.info.ctx.model_config.model - self._ensure_cached_runtime_components(model_dir, config) + """Encode the reference audio, if any; the prompt was tokenized in apply(). + + Upstream calls this for every request, including text-only ones and + ones whose audio is already cached (empty ``mm_items``), so it must not + touch the text tokenizer. + """ + valid_items = mm_items.select({key for key, count in mm_items.get_all_counts().items() if count > 0}) + mm_data, passthrough = self._get_hf_mm_data(valid_items) + audio = self._get_reference_audio(mm_data) + if audio is None: + return BatchFeature(dict(passthrough)) + processed = self._encode_reference_audio(*audio) + processed.update(passthrough) + return processed + def _get_reference_audio(self, mm_data: Mapping[str, object]) -> tuple[object, int] | None: audio = mm_data.get("audio", None) if audio is None: - audio = mm_data.get("audios") - if audio is not None: - audio = audio[0], config.sample_rate + audios = mm_data.get("audios") + if audios is not None: + audio = audios[0], self.info.ctx.get_hf_config().sample_rate + return audio + + def _encode_prompt_text(self, prompt: str, mm_kwargs: Mapping[str, object]) -> torch.Tensor: + config = self.info.ctx.get_hf_config() + model_dir = self.info.ctx.model_config.model + self._ensure_cached_runtime_components(model_dir, config) # Build text prompt with control tokens lang = mm_kwargs.get("lang", None) @@ -136,19 +169,14 @@ def _call_hf_processor( full_text = prompt text_prompt = f"{style_text}<|text_start|>{full_text}<|text_end|>" - text_tokens = self.text_tokenizer(text_prompt, return_tensors="pt").input_ids.squeeze(0) # [N_text] + return self.text_tokenizer(text_prompt, return_tensors="pt").input_ids.squeeze(0) # [N_text] - if audio is None: - # Text-only path (auto voice mode) - return BatchFeature( - { - "input_ids": text_tokens, - "input_len": [len(text_tokens)], - } - ) + def _encode_reference_audio(self, audio_signal: object, sr: int) -> BatchFeature: + """Voice cloning: encode the reference audio to 8-codebook tokens.""" + config = self.info.ctx.get_hf_config() + model_dir = self.info.ctx.model_config.model + self._ensure_cached_runtime_components(model_dir, config) - # Voice cloning: encode reference audio to tokens - audio_signal, sr = audio if isinstance(audio_signal, np.ndarray): audio_signal = torch.from_numpy(audio_signal).float() if audio_signal.dim() == 1: @@ -169,14 +197,12 @@ def _call_hf_processor( if ref_audio_tokens.dim() == 3: ref_audio_tokens = ref_audio_tokens.squeeze(0) # [8, T_ref] - ft = BatchFeature( + return BatchFeature( { - "input_ids": text_tokens, "ref_audio_tokens": ref_audio_tokens, # [8, T_ref] "ref_audio_len": [ref_audio_tokens.shape[1]], } ) - return ft def _get_mm_fields_config( self, @@ -188,15 +214,6 @@ def _get_mm_fields_config( "ref_audio_len": MultiModalFieldConfig.batched("audio"), } - def _hf_processor_applies_updates( - self, - prompt_text: str, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> bool: - return False - def _get_prompt_updates( self, mm_items: MultiModalDataItems, diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py index 35d988895b3..b3e916f25d7 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from collections.abc import Iterable from functools import cached_property @@ -139,11 +142,11 @@ def sample( return self.language_model.sample(logits, sampling_metadata) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=["thinker.", "token2wav."], + loader = AutoWeightsLoader(self) + loaded = loader.load_weights( + weights, + mapper=(self.hf_to_vllm_mapper) | WeightsMapper(orig_to_new_prefix={"thinker.": None, "token2wav.": None}), ) - loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) # Log load summary try: total_bytes = 0 diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py index ba073aad2a0..6bbefceab7b 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Thin Omni wrapper: reuse upstream Qwen2.5-Omni thinker with minimal overrides.""" import hashlib @@ -9,6 +12,7 @@ import PIL.Image as PILImage import torch from torch import nn +from transformers.feature_extraction_utils import BatchFeature from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import ( Qwen2_5OmniThinkerConfig, ) @@ -360,25 +364,40 @@ class Qwen2_5OmniThinkerMultiModalProcessor( ): """Override to fix use_audio_in_video detection when mm cache returns None.""" - def _call_hf_processor( + def _apply_hf_processor_main( self, - prompt: str, - mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], - ): - mm_kwargs = _presampled_videos_hf_kwargs(mm_data, mm_kwargs) - mm_kwargs = _coerce_use_audio_in_video_for_hf_processor(mm_data, mm_kwargs) - hf_inputs = super()._call_hf_processor( - prompt=prompt, - mm_data=mm_data, - mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> BatchFeature: + """Apply omni-specific HF kwargs before the upstream Qwen2.5-Omni + processor runs. + + Upstream's ``_apply_hf_processor_main`` is fully custom and no longer + routes through the generic ``_preprocess_hf_mm_data`` / + ``_postprocess_hf_mm_data`` hooks, so the omni video pre-sampling and + per-video ``use_audio_in_video`` mask handling have to run here. + """ + valid_mm_items = mm_items.select({k for k, c in mm_items.get_all_counts().items() if c > 0}) + processor_data, _ = self._get_hf_mm_data(valid_mm_items) + + hf_processor_mm_kwargs = _presampled_videos_hf_kwargs( + processor_data, + hf_processor_mm_kwargs, ) - per_video_mask = mm_kwargs.get(_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY) + hf_processor_mm_kwargs = _coerce_use_audio_in_video_for_hf_processor( + processor_data, + hf_processor_mm_kwargs, + ) + + processed_data = super()._apply_hf_processor_main( + mm_items, + hf_processor_mm_kwargs, + ) + + per_video_mask = hf_processor_mm_kwargs.get(_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY) if per_video_mask is not None: - hf_inputs["use_audio_in_video"] = torch.tensor(per_video_mask) - return hf_inputs + processed_data["use_audio_in_video"] = torch.tensor(per_video_mask) + return processed_data def _get_mm_fields_config( self, @@ -530,17 +549,17 @@ def get_replacement_qwen2_video(item_idx: int): return [ PromptReplacement( modality="audio", - target=audio_token, + target=[audio_token_id], replacement=get_replacement_qwen2_audio, ), PromptReplacement( modality="image", - target=image_token, + target=[image_token_id], replacement=partial(get_replacement_qwen2_vision, modality="image"), ), PromptReplacement( modality="video", - target=video_token, + target=[video_token_id], replacement=get_replacement_qwen2_video, ), ] @@ -620,7 +639,7 @@ def _cached_apply_hf_processor( self, inputs: ProcessorInputs, timing_ctx: TimingContext, - ) -> tuple[list[int], MultiModalProcessingInfo, bool]: + ) -> MultiModalProcessingInfo: """Cache-aware: processes video/audio pairs as units, using pair-specific cache keys. Falls back to standard cache if no pairs.""" @@ -673,16 +692,9 @@ def _cached_apply_hf_processor( ) with timing_ctx.record("apply_hf_processor"): - ( - prompt_ids, - mm_missing_processed_data, - is_update_applied, - ) = self._apply_hf_processor_main( - prompt=inputs.prompt, + mm_missing_processed_data = self._apply_hf_processor_main( mm_items=mm_missing_data_items, hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=inputs.tokenization_kwargs, - enable_hf_prompt_update=False, ) mm_missing_kwargs = MultiModalKwargsItems.from_hf_inputs( @@ -714,7 +726,7 @@ def _cached_apply_hf_processor( prompt_updates=mm_prompt_updates, ) - return prompt_ids, mm_info, is_update_applied + return mm_info def _derive_audio_from_video_placeholders( self, @@ -785,7 +797,6 @@ def _maybe_apply_prompt_updates( prompt_ids: list[int], mm_kwargs: MultiModalKwargsItems, mm_prompt_updates: MultiModalPromptUpdates, - is_update_applied: bool, ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]: mm_item_counts = mm_items.get_all_counts() self._validate_mm_kwargs(mm_kwargs, mm_item_counts) @@ -794,38 +805,28 @@ def _maybe_apply_prompt_updates( video_use_audio_in_video = self._get_video_use_audio_in_video(mm_kwargs, mm_prompt_updates) use_audio_in_video = any(video_use_audio_in_video) - if is_update_applied: - mm_placeholders = self._find_mm_placeholders( + if use_audio_in_video and "audio" in mm_prompt_updates: + filtered_updates = {k: v for k, v in mm_prompt_updates.items() if k != "audio"} + prompt_ids, mm_placeholders = self._apply_prompt_updates( prompt_ids, - mm_prompt_updates, + filtered_updates, ) - self._validate_mm_placeholders( + mm_placeholders = self._derive_audio_from_video_placeholders( mm_placeholders, - mm_item_counts, + mm_prompt_updates, + video_use_audio_in_video, ) else: - if use_audio_in_video and "audio" in mm_prompt_updates: - filtered_updates = {k: v for k, v in mm_prompt_updates.items() if k != "audio"} - prompt_ids, mm_placeholders = self._apply_prompt_updates( - prompt_ids, - filtered_updates, - ) - mm_placeholders = self._derive_audio_from_video_placeholders( - mm_placeholders, - mm_prompt_updates, - video_use_audio_in_video, - ) - else: - prompt_ids, mm_placeholders = self._apply_prompt_updates( - prompt_ids, - mm_prompt_updates, - ) - - self._validate_mm_placeholders( - mm_placeholders, - mm_item_counts, + prompt_ids, mm_placeholders = self._apply_prompt_updates( + prompt_ids, + mm_prompt_updates, ) + self._validate_mm_placeholders( + mm_placeholders, + mm_item_counts, + ) + return prompt_ids, mm_placeholders @@ -1397,11 +1398,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if self.visual is None: skip_prefixes.extend(["visual."]) - loader = AutoWeightsLoader( - self, - skip_prefixes=skip_prefixes, + loader = AutoWeightsLoader(self) + loaded_weights = loader.load_weights( + weights, + mapper=(self.hf_to_vllm_mapper) + | WeightsMapper(orig_to_new_prefix={name: None for name in (skip_prefixes or ())}), ) - loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) return loaded_weights diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py index 91eafd2a876..e9e92c904e8 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from collections.abc import Iterable import torch @@ -20,6 +23,7 @@ from vllm.model_executor.models.utils import ( AutoWeightsLoader, PPMissingLayer, + WeightsMapper, is_pp_missing_parameter, make_empty_intermediate_tensors_factory, make_layers, @@ -441,8 +445,8 @@ def sample( return next_tokens def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), + loader = AutoWeightsLoader(self) + return loader.load_weights( + weights, + mapper=WeightsMapper(orig_to_new_prefix={"lm_head.": None} if self.config.tie_word_embeddings else {}), ) - return loader.load_weights(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 5c3129be9ab..67352eecb6c 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 @@ -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 # Copyright 2025 The Qwen team. """Inference-only Qwen3-Omni-Moe Code2Wav model.""" @@ -336,11 +336,11 @@ def chunked_decode_streaming( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights from HuggingFace checkpoint.""" - loader = AutoWeightsLoader( - self, - skip_prefixes=["thinker.", "talker."], # Already loaded above + loader = AutoWeightsLoader(self) + loaded = loader.load_weights( + weights, + mapper=(self.hf_to_vllm_mapper) | WeightsMapper(orig_to_new_prefix={"thinker.": None, "talker.": None}), ) - loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) # Log load summary try: diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index 8337f8dadf8..ea3e2fcfe36 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from collections.abc import Iterable import torch @@ -299,13 +302,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: to vLLM's internal structure. Code predictor weights are routed to its custom loader for vocab extension support. """ - loader = AutoWeightsLoader( - self, - skip_prefixes=["thinker.", "code2wav."], - # "code_predictor."], - ) + loader = AutoWeightsLoader(self) # Don't apply mapper again since we already did it - loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + loaded = loader.load_weights( + weights, + mapper=(self.hf_to_vllm_mapper) | WeightsMapper(orig_to_new_prefix={"thinker.": None, "code2wav.": None}), + ) # Log load summary try: diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index bbfea498ba1..b2a7d7abecc 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -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 # Copyright 2025 The Qwen team. # Copyright 2023 The vLLM team. # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. @@ -31,10 +31,7 @@ import torch import torch.nn as nn import torch.nn.functional as F -from packaging.version import Version from transformers import PretrainedConfig -from transformers import __version__ as TRANSFORMERS_VERSION -from transformers.feature_extraction_utils import BatchFeature from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( Qwen3OmniMoeConfig, Qwen3OmniMoeThinkerConfig, @@ -95,6 +92,9 @@ from vllm.model_executor.models.qwen3_omni_moe_thinker import ( Qwen3OmniMoeAudioEncoderLayer as _Qwen3OmniMoeAudioEncoderLayer, ) +from vllm.model_executor.models.qwen3_omni_moe_thinker import ( + Qwen3OmniMoeThinkerMultiModalProcessor as _UpstreamQwen3Processor, +) from vllm.model_executor.models.qwen3_omni_moe_thinker import ( _get_feat_extract_output_lengths, ) @@ -124,7 +124,6 @@ Qwen2_5OmniThinkerMultiModalProcessor, _get_request_video_use_audio_in_video, _get_video_second_per_grid_t, - _presampled_videos_hf_kwargs, ) from vllm_omni.model_executor.models.qwen3_omni.quantization import ( Qwen3OmniNestedSupportsQuant, @@ -281,7 +280,12 @@ def __init__( self.scaling = self.head_dim**-0.5 - self.qkv = QKVParallelLinear( + # vLLM 0.29 renamed the audio attention projection to ``qkv_proj`` and + # its stacked mapping rewrites checkpoint ``q_proj``/``k_proj``/ + # ``v_proj`` onto that name, so the module must match or weight loading + # fails with "no module or parameter named ...self_attn.qkv_proj". + # The inherited forward also calls ``self.qkv_proj``. + self.qkv_proj = QKVParallelLinear( hidden_size=self.embed_dim, head_size=self.head_dim, total_num_heads=self.num_heads, @@ -289,7 +293,7 @@ def __init__( bias=True, disable_tp=disable_tp, quant_config=quant_config, - prefix=f"{prefix}.qkv", + prefix=f"{prefix}.qkv_proj", ) self.out_proj = RowParallelLinear( @@ -696,151 +700,10 @@ def get_mm_max_tokens_per_item( class Qwen3OmniMoeThinkerMultiModalProcessor( Qwen2_5OmniThinkerMultiModalProcessor, + _UpstreamQwen3Processor, ): - def _call_hf_processor( - self, - prompt: str, - mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], - ) -> BatchFeature: - mm_data = dict(mm_data) - mm_kwargs = dict(_presampled_videos_hf_kwargs(mm_data, mm_kwargs)) - audios = mm_data.pop("audios", []) - - def pad_to_hop_length(x: np.ndarray, hop_length: int) -> np.ndarray: - length = x.shape[-1] - if length % hop_length != 0: - pad_length = hop_length - (length % hop_length) - x = np.pad(x, (0, pad_length), mode="constant", constant_values=0) - return x - - # NOTE: WhisperFeatureExtractor cannot handle empty list of audios - feature_extractor = self.info.get_feature_extractor(**mm_kwargs) - hop_length = feature_extractor.hop_length - if audios: - # NOTE: Qwen3-Omni processor accept "audio" - # To make sure the cache works with padding=True, we pre-padded - # the audio to multiple of hop_length. - mm_data["audio"] = [ - pad_to_hop_length(audio, hop_length) - if isinstance(audio, np.ndarray) - else (pad_to_hop_length(audio[0], hop_length), audio[1]) - for audio in audios - ] - - # TODO(Isotr0py): Remove this patch after upstream fix PR - # released and Transformers version update: - # https://github.com/huggingface/transformers/pull/41473 - mm_kwargs = dict(mm_kwargs) - tok_kwargs = dict(tok_kwargs) - mm_kwargs["audio_kwargs"] = dict(mm_kwargs.get("audio_kwargs") or {}) - mm_kwargs["text_kwargs"] = dict(mm_kwargs.get("text_kwargs") or {}) - if Version(TRANSFORMERS_VERSION) < Version("4.58.0"): - # Extract audio_sample_rate before restructuring - audio_sample_rate = mm_kwargs.pop("audio_sample_rate", None) - - # move truncation to audio_kwargs level to avoid conflict - # with tok_kwargs - mm_kwargs["audio_kwargs"].setdefault("truncation", mm_kwargs.pop("truncation", False)) - mm_kwargs["text_kwargs"].setdefault("truncation", tok_kwargs.pop("truncation", False)) - - # Validate and conditionally pass audio_sample_rate - # WhisperFeatureExtractor has a fixed sampling rate, and vLLM's - # audio loader already resamples audio to the target rate. - # Only pass the value if it matches to avoid unexpected behavior. - if audio_sample_rate is not None: - expected_sr = feature_extractor.sampling_rate - if audio_sample_rate != expected_sr: - logger.warning( - "[%s] audio_sample_rate mismatch: user provided %dHz " - "but model expects %dHz. Ignoring user value. " - "vLLM's audio loader already resampled to %dHz.", - self.__class__.__name__, - audio_sample_rate, - expected_sr, - expected_sr, - ) - else: - # Sample rate matches, safe to pass - mm_kwargs["audio_kwargs"]["audio_sample_rate"] = audio_sample_rate - - hf_inputs = super()._call_hf_processor( - prompt=prompt, - mm_data=mm_data, - mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, - ) - - if ( - "audio_feature_lengths" in hf_inputs - and "feature_attention_mask" in hf_inputs - and (audios := mm_data.get("audio", [])) - ): - audio_num_frames = [] - for _, audio in enumerate(audios): - audio_length = len(audio[0]) if isinstance(audio, tuple) else len(audio) - num_frame = ( - (audio_length // hop_length) if audio_length % hop_length == 0 else (audio_length // hop_length - 1) - ) - if mm_kwargs.get("truncation", False): - num_frame = min(num_frame, feature_extractor.n_samples // hop_length) - audio_num_frames.append(num_frame) - hf_inputs["feature_attention_mask"] = [torch.ones(num_frame) for num_frame in audio_num_frames] - hf_inputs["audio_feature_lengths"] = torch.tensor(audio_num_frames) - return hf_inputs - - def _maybe_apply_prompt_updates( - self, - mm_items: MultiModalDataItems, - prompt_ids: list[int], - mm_kwargs: MultiModalKwargsItems, - mm_prompt_updates: MultiModalPromptUpdates, - is_update_applied: bool, - ) -> tuple[list[int], str, Mapping[str, list[PlaceholderFeaturesInfo]]]: - """ - Qwen3-Omni reimplements this function to handle `use_audio_in_video`. - """ - mm_item_counts = mm_items.get_all_counts() - self._validate_mm_kwargs(mm_kwargs, mm_item_counts) - - video_use_audio_in_video = self._get_video_use_audio_in_video(mm_kwargs, mm_prompt_updates) - use_audio_in_video = any(video_use_audio_in_video) - - # normal case with `use_audio_in_video=False` - if is_update_applied: - mm_placeholders = self._find_mm_placeholders( - prompt_ids, - mm_prompt_updates, - ) - self._validate_mm_placeholders( - mm_placeholders, - mm_item_counts, - ) - else: - if use_audio_in_video and "audio" in mm_prompt_updates: - filtered_updates = {k: v for k, v in mm_prompt_updates.items() if k != "audio"} - prompt_ids, mm_placeholders = self._apply_prompt_updates( - prompt_ids, - filtered_updates, - ) - mm_placeholders = self._derive_audio_from_video_placeholders( - mm_placeholders, - mm_prompt_updates, - video_use_audio_in_video, - ) - else: - prompt_ids, mm_placeholders = self._apply_prompt_updates( - prompt_ids, - mm_prompt_updates, - ) - - self._validate_mm_placeholders( - mm_placeholders, - mm_item_counts, - ) - - return prompt_ids, mm_placeholders + # Preserve Omni's per-video overrides while inheriting upstream Qwen3's + # audio processing through the shared, cooperative Qwen2.5 base. def get_updates_use_audio_in_video( self, @@ -978,17 +841,17 @@ def get_replacement_qwen2_video(item_idx: int): return [ PromptReplacement( modality="audio", - target=audio_token, + target=[audio_token_id], replacement=get_replacement_qwen2_audio, ), PromptReplacement( modality="image", - target=image_token, + target=[image_token_id], replacement=partial(get_replacement_qwen2_vision, modality="image"), ), PromptReplacement( modality="video", - target=video_token, + target=[video_token_id], replacement=get_replacement_qwen2_video, ), ] @@ -1114,7 +977,10 @@ def _process_audio_input( audio_input: Qwen2_5OmniAudioFeatureInputs, ) -> tuple[torch.Tensor, ...]: input_features = audio_input["input_features"] - audio_feature_lengths = audio_input["audio_feature_lengths"] + # vLLM 0.29 marks audio_feature_lengths keep_on_cpu, and the audio tower + # derives device placement from feature_lens, so move it explicitly. + # Mirrors upstream Qwen3OmniMoeThinker._process_audio_input. + audio_feature_lengths = audio_input["audio_feature_lengths"].to(input_features.device, non_blocking=True) audio_output_lengths = _get_feat_extract_output_lengths(audio_feature_lengths) @@ -1540,11 +1406,11 @@ def compute_logits( return self.language_model.compute_logits(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader( - self, - skip_prefixes=["talker.", "code2wav."], + loader = AutoWeightsLoader(self) + loaded_weights = loader.load_weights( + weights, + mapper=(self.hf_to_vllm_mapper) | WeightsMapper(orig_to_new_prefix={"talker.": None, "code2wav.": None}), ) - loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) return loaded_weights diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py index 87df8788975..f1692f73de0 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import os @@ -11,7 +14,7 @@ from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.model_loader import DefaultModelLoader -from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp @@ -552,10 +555,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: subfolder="speech_tokenizer", ) subfolder_weights = model_loader._get_weights_iterator(source) - loaded = AutoWeightsLoader( - self, - skip_prefixes=["encoder."], - ).load_weights(subfolder_weights) + loaded = AutoWeightsLoader(self).load_weights( + subfolder_weights, mapper=WeightsMapper(orig_to_new_prefix={"encoder.": None}) + ) device = self.vllm_config.device_config.device self.decoder.to(device=device, dtype=torch.float32) 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 e65334ff500..dfb24516a8a 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 @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from __future__ import annotations import copy @@ -1100,10 +1103,9 @@ def _talker_and_collect_speaker(ws: Iterable[tuple[str, torch.Tensor]]): subfolder="speech_tokenizer", ) subfolder_weights = model_loader._get_weights_iterator(source) - enc_loaded = AutoWeightsLoader( - self, - skip_prefixes=["decoder."], - ).load_weights(subfolder_weights) + enc_loaded = AutoWeightsLoader(self).load_weights( + subfolder_weights, mapper=WeightsMapper(orig_to_new_prefix={"decoder.": None}) + ) loaded |= enc_loaded # AutoWeightsLoader only loads parameters; the encoder's VQ diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py index b00e2487020..6dca8b358dc 100644 --- a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """ Step-Audio2 Thinker - Stage 1 LLM for Audio Understanding @@ -19,6 +22,7 @@ import torch.nn.functional as F from torch.nn.utils.rnn import pad_sequence from transformers import AutoTokenizer +from transformers.feature_extraction_utils import BatchFeature from transformers.processing_utils import ProcessorMixin from vllm.config import VllmConfig from vllm.inputs import MultiModalDataDict @@ -503,6 +507,10 @@ def _get_prompt_updates( processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) audio_token = getattr(processor, "audio_token", "") + audio_token_ids = self.info.get_tokenizer().encode( + audio_token, + add_special_tokens=False, + ) audio_items = mm_items.get("audio", []) num_audio_items = len(audio_items) if audio_items else 0 @@ -548,27 +556,22 @@ def get_replacement_audio(item_idx: int): return [ PromptReplacement( modality="audio", - target=audio_token, + target=audio_token_ids, replacement=get_replacement_audio, ) ] - def _call_hf_processor( + def _preprocess_hf_mm_data( self, - prompt: str, mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, object] | None = None, - tok_kwargs: Mapping[str, object] | None = None, - **_: object, - ): - """Call HF processor and post-process outputs""" + hf_processor_mm_kwargs: Mapping[str, object], + ) -> tuple[Mapping[str, object], Mapping[str, object]]: + mm_data, hf_processor_mm_kwargs = super()._preprocess_hf_mm_data( + mm_data, + hf_processor_mm_kwargs, + ) mm_data = dict(mm_data) audios = mm_data.pop("audios", []) - - mm_kwargs = mm_kwargs or {} - tok_kwargs = tok_kwargs or {} - - # CRITICAL: Ensure audios is ALWAYS a list to prevent string iteration if audios: if isinstance(audios, str): mm_data["audio"] = [audios] @@ -576,20 +579,24 @@ def _call_hf_processor( mm_data["audio"] = audios else: mm_data["audio"] = [audios] + return mm_data, hf_processor_mm_kwargs - hf_inputs = super()._call_hf_processor( - prompt=prompt, - mm_data=mm_data, - mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, + def _postprocess_hf_mm_data( + self, + mm_data: Mapping[str, object], + hf_processor_mm_kwargs: Mapping[str, object], + processed_data: BatchFeature, + ) -> BatchFeature: + processed_data = super()._postprocess_hf_mm_data( + mm_data, + hf_processor_mm_kwargs, + processed_data, ) - - if "audio_mels" not in hf_inputs: - hf_inputs["audio_mels"] = torch.empty((0, 128, 0)) - if "audio_lens" not in hf_inputs: - hf_inputs["audio_lens"] = torch.tensor([], dtype=torch.int32) - - return hf_inputs + if "audio_mels" not in processed_data: + processed_data["audio_mels"] = torch.empty((0, 128, 0)) + if "audio_lens" not in processed_data: + processed_data["audio_lens"] = torch.tensor([], dtype=torch.int32) + return processed_data # ============================================================================ diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py index 4f63bbec191..211159afaeb 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.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 math from collections.abc import Iterable, Mapping, Sequence @@ -48,11 +48,9 @@ from vllm.multimodal.processing import BaseDummyInputsBuilder, BaseMultiModalProcessor from vllm.multimodal.processing.processor import ( BaseProcessingInfo, - MultiModalProcessingInfo, ProcessorInputs, PromptReplacement, PromptUpdate, - TimingContext, ) from vllm.sequence import IntermediateTensors from vllm.tokenizers import cached_tokenizer_from_config @@ -871,7 +869,7 @@ def get_replacement(item_idx: int): return [ PromptReplacement( modality="audio", - target="", # Never match the prompt (see below note) + target=[audio_id], replacement=get_replacement, ), ] @@ -899,15 +897,24 @@ def _apply_hf_processor_mm_only( ) return mm_processed_data - def _cached_apply_hf_processor( + def _maybe_apply_prompt_updates( self, - inputs: ProcessorInputs, - timing_ctx: TimingContext, - ) -> tuple[list[int], MultiModalProcessingInfo, bool]: - prompt_ids, mm_info, _ = super()._cached_apply_hf_processor(inputs, timing_ctx) - - # NOTE: The tokens are already inserted by the chat template - return prompt_ids, mm_info, True + mm_items: MultiModalDataItems, + prompt_ids: list[int], + mm_kwargs: MultiModalKwargsItems, + mm_prompt_updates, + ): + # Voxtral's Mistral chat template has already inserted the complete + # audio-token run, so locate it without applying the replacement again. + mm_item_counts = mm_items.get_all_counts() + self._validate_mm_kwargs(mm_kwargs, mm_item_counts) + self._validate_mm_updates(mm_prompt_updates, mm_item_counts) + mm_placeholders = self._find_mm_placeholders( + prompt_ids, + mm_prompt_updates, + ) + self._validate_mm_placeholders(mm_placeholders, mm_item_counts) + return prompt_ids, mm_placeholders @MULTIMODAL_REGISTRY.register_processor( diff --git a/vllm_omni/platforms/__init__.py b/vllm_omni/platforms/__init__.py index 942eb7c24bd..d5ee64353bf 100644 --- a/vllm_omni/platforms/__init__.py +++ b/vllm_omni/platforms/__init__.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 import traceback @@ -142,7 +142,11 @@ def resolve_current_omni_platform_cls_qualname() -> str: if platform_cls_qualname is not None: activated_plugins.append(name) except Exception: - pass + logger.debug( + "OmniPlatform plugin %s failed during detection.", + name, + exc_info=True, + ) activated_builtin_plugins = list(set(activated_plugins) & set(builtin_omni_platform_plugins.keys())) activated_oot_plugins = list(set(activated_plugins) & set(platform_plugins.keys())) diff --git a/vllm_omni/reasoning/step_audio_reasoning_parser.py b/vllm_omni/reasoning/step_audio_reasoning_parser.py index 81ef6349686..ddf5a434f56 100644 --- a/vllm_omni/reasoning/step_audio_reasoning_parser.py +++ b/vllm_omni/reasoning/step_audio_reasoning_parser.py @@ -5,7 +5,7 @@ from itertools import islice from typing import TYPE_CHECKING -from vllm.entrypoints.openai.engine.protocol import DeltaMessage +from vllm.entrypoints.generate.base.protocol import DeltaMessage from vllm.logger import init_logger from vllm.reasoning import ReasoningParser from vllm.tokenizers import TokenizerLike diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 811f86f179a..6625effe87a 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -2128,7 +2128,6 @@ def propose_draft_token_ids(sampled_token_ids): elif self.valid_sampled_token_count_event is not None: assert spec_decode_common_attn_metadata is not None next_token_ids, valid_sampled_tokens_count = self.drafter.prepare_next_token_ids_padded( - self.optimistic_seq_lens_cpu, sampled_token_ids, self.requests, self.input_batch, diff --git a/vllm_omni/worker/sampling_utils.py b/vllm_omni/worker/sampling_utils.py index 84588113864..39ecae469d8 100644 --- a/vllm_omni/worker/sampling_utils.py +++ b/vllm_omni/worker/sampling_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 """Sampling-state guards shared by the GPU and NPU AR model runners.""" import torch @@ -48,7 +48,7 @@ def sanitize_min_tokens_stop_ids(logitsprocs: LogitsProcessors, logits_vocab: in if not min_toks: continue needs_rebuild = False - for _, _, stop_tok_ids in min_toks.values(): + for _, _, stop_tok_ids, _ in min_toks.values(): oob = [tok for tok in stop_tok_ids if tok >= logits_vocab] if not oob: continue @@ -63,10 +63,19 @@ def sanitize_min_tokens_stop_ids(logitsprocs: LogitsProcessors, logits_vocab: in if needs_rebuild: reqs: list[int] = [] tok_ids: list[int] = [] - for index, (_, _, stop_tok_ids) in min_toks.items(): + restore_reqs: list[int] = [] + restore_tok_ids: list[int] = [] + for index, (_, _, stop_tok_ids, uses_structured_output) in min_toks.items(): reqs.extend([index] * len(stop_tok_ids)) tok_ids.extend(stop_tok_ids) + if uses_structured_output: + restore_reqs.extend([index] * len(stop_tok_ids)) + restore_tok_ids.extend(stop_tok_ids) proc.logits_slice = ( proc._device_tensor(reqs, torch.int32), proc._device_tensor(tok_ids, torch.int32), ) + proc.restore_logits_slice = ( + proc._device_tensor(restore_reqs, torch.int32), + proc._device_tensor(restore_tok_ids, torch.int32), + ) From bbac4df7d3392791ce64a574c4f78c18927e61fc Mon Sep 17 00:00:00 2001 From: herotai214 <68222888+herotai214@users.noreply.github.com> Date: Thu, 10 Sep 2026 15:03:21 +0800 Subject: [PATCH 49/93] [Refactor] P0.2: Migrate API server helpers out of api_server (#5453) Signed-off-by: herotai214 --- .../test_pipeline_lingbot_video.py | 4 +- .../openai_api/test_api_server_guards.py | 26 +- .../openai_api/test_image_server.py | 12 +- .../openai_api/test_omni_sleep_wakeup.py | 4 +- .../openai_api/test_serving_audio_generate.py | 6 +- .../openai_api/test_serving_speech.py | 10 +- .../openai_api/test_video_server.py | 24 +- vllm_omni/config/endpoint_policy.py | 19 +- vllm_omni/entrypoints/README.md | 21 + vllm_omni/entrypoints/cli/README.md | 21 + vllm_omni/entrypoints/cli/benchmark/README.md | 18 + vllm_omni/entrypoints/duplex/warmup.py | 130 ++ vllm_omni/entrypoints/openai/README.md | 58 + vllm_omni/entrypoints/openai/api_server.py | 1749 +---------------- vllm_omni/entrypoints/openai/app_state.py | 97 + .../entrypoints/openai/audio_utils_mixin.py | 20 + vllm_omni/entrypoints/openai/chat_template.py | 55 + vllm_omni/entrypoints/openai/diffusion.py | 76 + vllm_omni/entrypoints/openai/errors.py | 50 + .../entrypoints/openai/image_api_utils.py | 23 +- vllm_omni/entrypoints/openai/images/README.md | 31 + .../entrypoints/openai/images/__init__.py | 0 .../entrypoints/openai/images/helpers.py | 245 +++ vllm_omni/entrypoints/openai/lora.py | 38 + vllm_omni/entrypoints/openai/models/README.md | 19 + .../entrypoints/openai/models/__init__.py | 0 .../entrypoints/openai/models/serving.py | 59 + vllm_omni/entrypoints/openai/utils.py | 18 + vllm_omni/entrypoints/openai/video/README.md | 33 + .../entrypoints/openai/video/__init__.py | 0 .../openai/video/generation/README.md | 33 + .../openai/video/generation/__init__.py | 0 .../openai/video/generation/helpers.py | 986 ++++++++++ .../entrypoints/openai/video_api_utils.py | 21 +- vllm_omni/entrypoints/serve/README.md | 22 + vllm_omni/entrypoints/serve/__init__.py | 0 .../entrypoints/serve/omni_control/README.md | 19 + .../serve/omni_control/__init__.py | 0 .../serve/omni_control/protocol.py | 14 + vllm_omni/entrypoints/serve/profile/README.md | 19 + .../entrypoints/serve/profile/__init__.py | 0 .../entrypoints/serve/profile/protocol.py | 14 + vllm_omni/entrypoints/serve/profile/utils.py | 27 + vllm_omni/entrypoints/serve/utils/README.md | 23 + vllm_omni/entrypoints/serve/utils/__init__.py | 0 vllm_omni/entrypoints/serve/utils/errors.py | 92 + vllm_omni/entrypoints/serve/utils/routes.py | 45 + 47 files changed, 2437 insertions(+), 1744 deletions(-) create mode 100644 vllm_omni/entrypoints/README.md create mode 100644 vllm_omni/entrypoints/cli/README.md create mode 100644 vllm_omni/entrypoints/cli/benchmark/README.md create mode 100644 vllm_omni/entrypoints/duplex/warmup.py create mode 100644 vllm_omni/entrypoints/openai/README.md create mode 100644 vllm_omni/entrypoints/openai/app_state.py create mode 100644 vllm_omni/entrypoints/openai/chat_template.py create mode 100644 vllm_omni/entrypoints/openai/diffusion.py create mode 100644 vllm_omni/entrypoints/openai/images/README.md create mode 100644 vllm_omni/entrypoints/openai/images/__init__.py create mode 100644 vllm_omni/entrypoints/openai/images/helpers.py create mode 100644 vllm_omni/entrypoints/openai/lora.py create mode 100644 vllm_omni/entrypoints/openai/models/README.md create mode 100644 vllm_omni/entrypoints/openai/models/__init__.py create mode 100644 vllm_omni/entrypoints/openai/models/serving.py create mode 100644 vllm_omni/entrypoints/openai/video/README.md create mode 100644 vllm_omni/entrypoints/openai/video/__init__.py create mode 100644 vllm_omni/entrypoints/openai/video/generation/README.md create mode 100644 vllm_omni/entrypoints/openai/video/generation/__init__.py create mode 100644 vllm_omni/entrypoints/openai/video/generation/helpers.py create mode 100644 vllm_omni/entrypoints/serve/README.md create mode 100644 vllm_omni/entrypoints/serve/__init__.py create mode 100644 vllm_omni/entrypoints/serve/omni_control/README.md create mode 100644 vllm_omni/entrypoints/serve/omni_control/__init__.py create mode 100644 vllm_omni/entrypoints/serve/omni_control/protocol.py create mode 100644 vllm_omni/entrypoints/serve/profile/README.md create mode 100644 vllm_omni/entrypoints/serve/profile/__init__.py create mode 100644 vllm_omni/entrypoints/serve/profile/protocol.py create mode 100644 vllm_omni/entrypoints/serve/profile/utils.py create mode 100644 vllm_omni/entrypoints/serve/utils/README.md create mode 100644 vllm_omni/entrypoints/serve/utils/__init__.py create mode 100644 vllm_omni/entrypoints/serve/utils/errors.py create mode 100644 vllm_omni/entrypoints/serve/utils/routes.py 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 846a2642f35..3775ef665c5 100644 --- a/tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.py +++ b/tests/diffusion/models/lingbot_video/test_pipeline_lingbot_video.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 @@ -168,7 +168,7 @@ def test_postprocess_preserves_image_contract_for_serving(): format_diffusion_outputs, normalize_diffusion_postprocess_output, ) - from vllm_omni.entrypoints.openai.api_server import _extract_images_from_result + from vllm_omni.entrypoints.openai.images.helpers import _extract_images_from_result sampling = OmniDiffusionSamplingParams( num_inference_steps=1, diff --git a/tests/entrypoints/openai_api/test_api_server_guards.py b/tests/entrypoints/openai_api/test_api_server_guards.py index 71f93f6dccb..ee141182c22 100644 --- a/tests/entrypoints/openai_api/test_api_server_guards.py +++ b/tests/entrypoints/openai_api/test_api_server_guards.py @@ -55,6 +55,7 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm_omni.entrypoints.openai import api_server +from vllm_omni.entrypoints.serve.utils import errors as serve_errors pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -656,14 +657,16 @@ def test_images_generation_without_engine_preserves_service_unavailable_error() assert exc_info.value.detail == "Multi-stage engine not initialized. Start server with a multi-stage omni model." -def test_images_generation_without_multistage_chat_handler_preserves_unavailable_error(monkeypatch) -> None: +def test_images_generation_without_multistage_chat_handler_preserves_unavailable_error() -> None: """Lock images generation when chat serving was not wired. Fails if multi-stage images stop requiring ``openai_serving_chat``, or the 503 detail changes after images/chat ownership splits. """ app = FastAPI() - app.state.engine_client = SimpleNamespace(stage_configs=[object(), object()]) + # Prefer real stage_type input over monkeypatching ``get_stage_type``: + # ``_get_engine_and_model`` lives in ``app_state`` and binds utils directly. + app.state.engine_client = SimpleNamespace(stage_configs=[{"stage_type": "llm"}, {"stage_type": "diffusion"}]) app.state.stage_configs = app.state.engine_client.stage_configs app.state.openai_serving_models = SimpleNamespace(base_model_paths=[SimpleNamespace(name="demo-model")]) app.state.openai_serving_chat = None @@ -671,7 +674,6 @@ def test_images_generation_without_multistage_chat_handler_preserves_unavailable raw_request = _request_for(app, method="POST", path="/v1/images/generations") request = api_server.ImageGenerationRequest(prompt="a cat", model="demo-model") - monkeypatch.setattr(api_server, "get_stage_type", lambda _stage_cfg: "diffusion") with pytest.raises(HTTPException) as exc_info: asyncio.run(api_server.generate_images(request, raw_request)) @@ -709,9 +711,7 @@ def test_engine_error_json_response_includes_request_and_stage_fields(monkeypatc in P0.2 is not a false red. """ app = FastAPI() - # TODO(P0.2): retarget this call if ``_register_omni_exception_handlers`` - # leaves api_server (red here is a rename, not a surface regression). - api_server._register_omni_exception_handlers(app) + serve_errors._register_omni_exception_handlers(app) app.state.engine_client = SimpleNamespace(errored=True, engine=SimpleNamespace(is_alive=lambda: False)) app.state.server = object() app.state.args = SimpleNamespace(log_error_stack=False) @@ -719,7 +719,7 @@ def test_engine_error_json_response_includes_request_and_stage_fields(monkeypatc req = _request_for(app, method="POST", path="/v1/chat/completions") req.state.request_metadata = SimpleNamespace(request_id="req-123") - monkeypatch.setattr(api_server, "terminate_if_errored", lambda **_kwargs: None) + monkeypatch.setattr(serve_errors, "terminate_if_errored", lambda **_kwargs: None) exc = EngineGenerateError("boom") exc.error_stage_id = 2 # type: ignore[attr-defined] @@ -740,14 +740,12 @@ def test_engine_dead_error_handler_registered_returns_json(monkeypatch) -> None: JSON with ``request_id``. """ app = FastAPI() - # TODO(P0.2): retarget this call if ``_register_omni_exception_handlers`` - # leaves api_server (red here is a rename, not a surface regression). - api_server._register_omni_exception_handlers(app) + serve_errors._register_omni_exception_handlers(app) app.state.engine_client = SimpleNamespace(errored=True, engine=SimpleNamespace(is_alive=lambda: False)) app.state.server = object() app.state.args = SimpleNamespace(log_error_stack=False) - monkeypatch.setattr(api_server, "terminate_if_errored", lambda **_kwargs: None) + monkeypatch.setattr(serve_errors, "terminate_if_errored", lambda **_kwargs: None) handler = app.exception_handlers[EngineDeadError] req = _request_for(app) @@ -766,7 +764,7 @@ async def test_pure_diffusion_app_state_key_snapshot(monkeypatch) -> None: models), wires a handler that must stay None, or leaves a required handler as None. """ - stage = SimpleNamespace(engine_args={}) + stage = SimpleNamespace(stage_type="diffusion", engine_args={}) engine = _FakeEngineClient(stage_configs=[stage]) def _for_diffusion_factory(label: str): @@ -776,7 +774,6 @@ def _factory(cls, *args, **kwargs): return _factory - monkeypatch.setattr(api_server, "get_stage_type", lambda _cfg: "diffusion") monkeypatch.setattr(api_server.OmniOpenAIServingChat, "for_diffusion", _for_diffusion_factory("chat")) monkeypatch.setattr(api_server.OmniOpenAIServingChatBatch, "for_diffusion", _for_diffusion_factory("chat_batch")) monkeypatch.setattr( @@ -811,7 +808,7 @@ def _factory(cls, *args, **kwargs): @pytest.mark.asyncio async def test_pure_diffusion_speech_forwards_media_access_args(monkeypatch) -> None: - stage = SimpleNamespace(engine_args={}) + stage = SimpleNamespace(stage_type="diffusion", engine_args={}) engine = _FakeEngineClient(stage_configs=[stage]) speech_kwargs = {} @@ -827,7 +824,6 @@ def _speech_factory(cls, *args, **kwargs): speech_kwargs.update(kwargs) return _marker("speech") - monkeypatch.setattr(api_server, "get_stage_type", lambda _cfg: "diffusion") monkeypatch.setattr(api_server.OmniOpenAIServingChat, "for_diffusion", _for_diffusion_factory("chat")) monkeypatch.setattr(api_server.OmniOpenAIServingChatBatch, "for_diffusion", _for_diffusion_factory("chat_batch")) monkeypatch.setattr( diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index b2eb2edfc9b..c45563617c2 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Tests for async image generation API endpoints. @@ -24,11 +24,13 @@ from vllm.sampling_params import RequestOutputKind from vllm_omni.entrypoints.async_omni import AsyncOmni -from vllm_omni.entrypoints.openai.api_server import _check_max_generated_image_size, _DiffusionServingModels, router +from vllm_omni.entrypoints.openai.api_server import router from vllm_omni.entrypoints.openai.image_api_utils import ( encode_image_base64, parse_size, ) +from vllm_omni.entrypoints.openai.images.helpers import _check_max_generated_image_size +from vllm_omni.entrypoints.openai.models.serving import _DiffusionServingModels from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat from vllm_omni.errors import GuardrailViolationError from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -219,7 +221,7 @@ def test_client(mock_async_diffusion): app.state.stage_configs = [SimpleNamespace(stage_type="diffusion")] from vllm.entrypoints.openai.models.protocol import BaseModelPath - from vllm_omni.entrypoints.openai.api_server import _DiffusionServingModels + from vllm_omni.entrypoints.openai.models.serving import _DiffusionServingModels app.state.openai_serving_models = _DiffusionServingModels( [BaseModelPath(name="Qwen/Qwen-Image", model_path="Qwen/Qwen-Image")] @@ -2126,7 +2128,7 @@ def test_normalize_image(): """Test _normalize_image with various input types""" import numpy as np - from vllm_omni.entrypoints.openai.api_server import _normalize_image + from vllm_omni.entrypoints.openai.images.helpers import _normalize_image # Test PIL Image input img = Image.new("RGB", (64, 64), color="red") @@ -2163,7 +2165,7 @@ def test_extract_images_from_result(): """Test _extract_images_from_result with various result formats""" import numpy as np - from vllm_omni.entrypoints.openai.api_server import _extract_images_from_result + from vllm_omni.entrypoints.openai.images.helpers import _extract_images_from_result # Test empty result class EmptyResult: diff --git a/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py b/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py index c47eb8da6f2..615248fcfb3 100644 --- a/tests/entrypoints/openai_api/test_omni_sleep_wakeup.py +++ b/tests/entrypoints/openai_api/test_omni_sleep_wakeup.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 dataclasses @@ -264,7 +264,7 @@ def pure_diffusion_app(pure_diffusion_engine, mocker): mocker.patch( "vllm_omni.entrypoints.openai.api_server.OmniOpenAIServingSpeech" ).for_diffusion.return_value = mocker.MagicMock() - mocker.patch("vllm_omni.entrypoints.openai.api_server._DiffusionServingModels") + mocker.patch("vllm_omni.entrypoints.openai.models.serving._DiffusionServingModels") loop = asyncio.new_event_loop() try: diff --git a/tests/entrypoints/openai_api/test_serving_audio_generate.py b/tests/entrypoints/openai_api/test_serving_audio_generate.py index e1150c05bd0..977910d4e3f 100644 --- a/tests/entrypoints/openai_api/test_serving_audio_generate.py +++ b/tests/entrypoints/openai_api/test_serving_audio_generate.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import logging from inspect import Signature, signature from types import SimpleNamespace @@ -18,6 +21,7 @@ from vllm_omni.entrypoints.openai.serving_audio_generate import ( OmniOpenAIServingAudioGenerate, ) +from vllm_omni.entrypoints.serve.utils import errors as serve_errors from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -515,7 +519,7 @@ def test_api_server_engine_error_response_includes_request_and_stage_id( handler.create_audio_generate = AsyncMock(side_effect=exc) app = _make_api_server_test_app(handler) - with patch.object(api_server_module, "terminate_if_errored") as terminate_mock: + with patch.object(serve_errors, "terminate_if_errored") as terminate_mock: with TestClient(app) as client: response = client.post("/v1/audio/generate", json={"input": "Hello"}) diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 3922ab4a765..6040bede35e 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -32,6 +32,7 @@ from vllm_omni.diffusion.sched.step_scheduler import StepScheduler from vllm_omni.entrypoints.omni_base import OmniEngineDeadError from vllm_omni.entrypoints.openai import api_server as api_server_module +from vllm_omni.entrypoints.openai import errors as openai_errors from vllm_omni.entrypoints.openai import serving_speech as serving_speech_module from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin from vllm_omni.entrypoints.openai.protocol.audio import ( @@ -58,6 +59,7 @@ from vllm_omni.entrypoints.openai.tts_adapters.ming_tts import MingTTSAdapter from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSAdapter, Qwen3TTSCodecLimitError from vllm_omni.entrypoints.openai.tts_adapters.voxtral import VoxtralTTSAdapter +from vllm_omni.entrypoints.serve.utils import errors as serve_errors from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.model_executor.models.fish_speech.prompt_utils import ( FISH_TEXT_ONLY_SYSTEM_PROMPT, @@ -3783,7 +3785,7 @@ def _fake_create_error_response(message, err_type="BadRequestError", status_code fake_base = mocker.MagicMock() fake_base.create_error_response.side_effect = _fake_create_error_response - mocker.patch.object(api_server_module, "base", return_value=fake_base) + mocker.patch.object(openai_errors, "base", return_value=fake_base) return fake_base @@ -4098,7 +4100,7 @@ def test_api_server_create_speech_engine_error_response_includes_request_and_sta ) ) - terminate_mock = mocker.patch.object(api_server_module, "terminate_if_errored") + terminate_mock = mocker.patch.object(serve_errors, "terminate_if_errored") raw_request = _make_api_server_request(handler, path="/v1/audio/speech") raw_request.app.state.args = SimpleNamespace(log_error_stack=False) @@ -4130,8 +4132,8 @@ def test_omni_engine_error_handler_includes_request_and_stage_id(mocker: MockerF ) app.state.server = SimpleNamespace() - terminate_mock = mocker.patch.object(api_server_module, "terminate_if_errored") - api_server_module._register_omni_exception_handlers(app) + terminate_mock = mocker.patch.object(serve_errors, "terminate_if_errored") + serve_errors._register_omni_exception_handlers(app) @app.get("/boom") async def boom(request: Request): diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 1b18a9ae518..b65b00e59fb 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -37,6 +37,12 @@ from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo, ReferenceImage from vllm_omni.entrypoints.openai.storage import LocalStorageManager from vllm_omni.entrypoints.openai.stores import AsyncDictStore, TaskRegistry +from vllm_omni.entrypoints.openai.video.generation import helpers as video_generation_helpers +from vllm_omni.entrypoints.openai.video.generation.helpers import ( + MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES, + _read_upload_limited, + _reference_video_decode_spec, +) from vllm_omni.errors import GuardrailViolationError from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -192,6 +198,8 @@ def isolated_video_backends(tmp_path, monkeypatch): monkeypatch.setattr(api_server, "VIDEO_STORE", store) monkeypatch.setattr(api_server, "VIDEO_TASKS", tasks) monkeypatch.setattr(api_server, "STORAGE_MANAGER", storage) + monkeypatch.setattr(video_generation_helpers, "VIDEO_STORE", store) + monkeypatch.setattr(video_generation_helpers, "STORAGE_MANAGER", storage) return store, tasks, storage @@ -245,7 +253,7 @@ async def fake_init_app_state(engine_client, state, args): monkeypatch.setattr(api_server, "build_openai_app", lambda args, supported_tasks: FastAPI()) monkeypatch.setattr(api_server, "serve_http", fake_serve_http) monkeypatch.setattr(api_server.STORAGE_MANAGER, "start", fake_storage_start) - monkeypatch.setattr(api_server, "_get_vllm_config", fake_get_vllm_config) + monkeypatch.setattr(api_server.openai_app_state, "_get_vllm_config", fake_get_vllm_config) monkeypatch.setattr(api_server, "omni_init_app_state", fake_init_app_state) monkeypatch.setattr(api_server, "get_uvicorn_log_config", lambda args: None) @@ -978,7 +986,7 @@ def test_cosmos3_reference_video_limit_uses_v2v_condition_frames(): extra_params={"condition_frame_indexes_vision": [0, 2]}, ) - spec = api_server._reference_video_decode_spec(request, _cosmos3_stage_configs()) + spec = _reference_video_decode_spec(request, _cosmos3_stage_configs()) assert spec.max_frames == 9 assert spec.keep == "first" @@ -990,7 +998,7 @@ def test_cosmos3_reference_video_limit_preserves_action_frames(): extra_params={"action_mode": "inverse_dynamics", "action_chunk_size": 16}, ) - assert api_server._reference_video_decode_spec(request, _cosmos3_stage_configs()).max_frames == 17 + assert _reference_video_decode_spec(request, _cosmos3_stage_configs()).max_frames == 17 def test_cosmos3_reference_video_limit_caps_condition_frames_to_output_frames(): @@ -1000,7 +1008,7 @@ def test_cosmos3_reference_video_limit_caps_condition_frames_to_output_frames(): extra_params={"condition_frame_indexes_vision": [0, 20]}, ) - assert api_server._reference_video_decode_spec(request, _cosmos3_stage_configs()).max_frames == 5 + assert _reference_video_decode_spec(request, _cosmos3_stage_configs()).max_frames == 5 def test_s2v_video_generation_with_audio_reference_form(test_client, mocker: MockerFixture): @@ -1681,15 +1689,15 @@ def test_h3_multipart_maps_pillow_pixel_limit_error(field, test_client, monkeypa @pytest.mark.asyncio async def test_h3_upload_limit_checks_declared_size_before_read(): class OversizedUpload: - size = api_server.MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES + 1 + size = MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES + 1 async def read(self, _size): raise AssertionError("the oversized upload must be rejected before reading") with pytest.raises(HTTPException, match="size limit"): - await api_server._read_upload_limited( + await _read_upload_limited( OversizedUpload(), - max_bytes=api_server.MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES, + max_bytes=MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES, ) @@ -2438,7 +2446,7 @@ def test_cosmos3_control_upload_rejects_existing_control_source(test_client): ) def test_cosmos3_control_upload_rejects_invalid_size(control_bytes, message, test_client, monkeypatch): test_client.app.state.openai_serving_video._engine_client.model_class_name = "Cosmos3OmniDiffusersPipeline" - monkeypatch.setattr(api_server, "CONTROL_REFERENCE_MAX_BYTES", 3) + monkeypatch.setattr(video_generation_helpers, "CONTROL_REFERENCE_MAX_BYTES", 3) response = test_client.post( "/v1/videos/sync", diff --git a/vllm_omni/config/endpoint_policy.py b/vllm_omni/config/endpoint_policy.py index a6c9e4cf675..f0aca95eb9e 100644 --- a/vllm_omni/config/endpoint_policy.py +++ b/vllm_omni/config/endpoint_policy.py @@ -2,16 +2,16 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Endpoint restriction policy for omni pipelines.""" -from collections.abc import Set from dataclasses import dataclass from enum import Enum from typing import NamedTuple from fastapi import FastAPI, Request from fastapi.responses import JSONResponse -from starlette.routing import Route from vllm.entrypoints.serve.exception_handling.error_response import create_error_response +from vllm_omni.entrypoints.serve.utils.routes import remove_route_from_app + class RouteTarget(NamedTuple): """A server path & supported methods.""" @@ -53,21 +53,6 @@ async def rejection_handler(raw_request: Request): return rejection_handler -def remove_route_from_app( - app: FastAPI, - path: str, - methods: Set[str], -) -> None: - """Remove routes matching a path and one of the given HTTP methods.""" - routes_to_remove = [ - route - for route in app.routes - if isinstance(route, Route) and route.path == path and route.methods is not None and route.methods & methods - ] - for route in routes_to_remove: - app.routes.remove(route) - - def shutdown_unsupported_routes( app: FastAPI, endpoint_restrictions: tuple[EndpointRestriction, ...], diff --git a/vllm_omni/entrypoints/README.md b/vllm_omni/entrypoints/README.md new file mode 100644 index 00000000000..7af9681440d --- /dev/null +++ b/vllm_omni/entrypoints/README.md @@ -0,0 +1,21 @@ +# Entrypoints + +This package contains user-facing server, CLI, orchestration, and protocol +entrypoints for vLLM-Omni. + +## What Belongs Here + +- Runtime orchestration entrypoints such as `AsyncOmni`, `Omni`, and + request-state helpers. +- CLI integration under `cli/`. +- Shared cross-entrypoint configuration and stage utilities. + +## What Does Not Belong Here + +- OpenAI endpoint-family route bodies and request helpers; those should move + toward `openai//` packages as part of #5227. +- Generic model execution internals that belong under engine, worker, or model + executor packages. + +Some top-level utility files remain acceptable homes under Phase 0.2 (`P0.2`) of #5227. +They are not part of the OpenAI helper migration unless explicitly referenced by #5227. diff --git a/vllm_omni/entrypoints/cli/README.md b/vllm_omni/entrypoints/cli/README.md new file mode 100644 index 00000000000..37b77391686 --- /dev/null +++ b/vllm_omni/entrypoints/cli/README.md @@ -0,0 +1,21 @@ +# CLI Entrypoints + +This package owns vLLM-Omni command-line integration. + +## What Belongs Here + +- Omni CLI command registration and dispatch. +- CLI argument parsing and validation for `vllm ... --omni`. +- Headless/runtime launch wiring that ultimately calls server bootstrap code. +- CLI-only presentation helpers such as the logo. + +## What Does Not Belong Here + +- FastAPI route bodies. +- OpenAI endpoint request helpers. +- Server app utilities that are shared outside CLI; those belong under + `entrypoints/serve`. + +The CLI currently imports `openai.api_server.omni_run_server`. If server +bootstrap moves later, update this import, but do not fold route helpers into +the CLI package. Related endpoint cleanup is tracked by #5227. diff --git a/vllm_omni/entrypoints/cli/benchmark/README.md b/vllm_omni/entrypoints/cli/benchmark/README.md new file mode 100644 index 00000000000..a385aaa7b93 --- /dev/null +++ b/vllm_omni/entrypoints/cli/benchmark/README.md @@ -0,0 +1,18 @@ +# Benchmark CLI + +This package owns Omni benchmark subcommands exposed through the CLI. + +## What Belongs Here + +- CLI parser extensions for benchmark commands. +- Benchmark command dispatch. +- Benchmark-only argument groups. + +## What Does Not Belong Here + +- Runtime server route helpers. +- OpenAI endpoint implementation details. +- Benchmark execution internals that belong under `vllm_omni.benchmarks`. + +This package is not part of the OpenAI endpoint helper migration. It may need +import updates only if CLI/bootstrap ownership changes during #5227. diff --git a/vllm_omni/entrypoints/duplex/warmup.py b/vllm_omni/entrypoints/duplex/warmup.py new file mode 100644 index 00000000000..c380775ef73 --- /dev/null +++ b/vllm_omni/entrypoints/duplex/warmup.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Startup warmup probe for duplex ``/v1/realtime``. + +Moved out of ``api_server.py`` under P0.2 of #5227. The +``omni_run_server_worker`` scheduling and ``/v1/realtime`` hold-clients +preamble stay in ``api_server.py`` until P0.3. +""" + +import asyncio +import json +import time +import uuid + +import httpx +import pybase64 as base64 +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +async def _warmup_duplex_realtime(app, args, warmup_frames: int) -> None: + """Run silent frames through a throwaway realtime session at startup. + + Self-connects over the server's own websocket endpoint so the warmup + exercises the exact production path (serving adapter, duplex control + plane, every pipeline stage). One-time costs — kernel JIT compilation, + first prefill/decode paths, codec caches — are paid here instead of by + the first real client; ``/v1/realtime`` holds other connections until + the warmup finishes (see ``duplex_warmup_done``). + """ + warmup_done = getattr(app.state, "duplex_warmup_done", None) + try: + try: + import websockets + except ImportError: + logger.warning("Duplex warmup skipped: the 'websockets' package is not installed.") + return + served = getattr(args, "served_model_name", None) + if isinstance(served, list | tuple) and served: + model_name = served[0] + elif isinstance(served, str) and served: + model_name = served + else: + model_name = args.model + adapter = getattr(getattr(app.state, "openai_serving_duplex", None), "_serving_runtime_adapter", None) + frame_samples = int(getattr(adapter, "silence_continuation_samples", 16000)) + silence = base64.b64encode(bytes(frame_samples * 4)).decode("ascii") + from vllm_omni.experimental.fullduplex.client import build_realtime_url + + url = ( + build_realtime_url(f"ws://127.0.0.1:{args.port}/v1/realtime", model_name, autostart=False) + + "&vllm_omni_warmup=1" + ) + logger.info("Duplex warmup starting: %d silent frames of %d samples.", warmup_frames, frame_samples) + start = time.time() + # The socket is bound before uvicorn finishes starting, so a plain + # connect can be accepted and then reset. Wait for /health first. + async with httpx.AsyncClient() as http: + for _ in range(120): + try: + if (await http.get(f"http://127.0.0.1:{args.port}/health", timeout=2)).status_code == 200: + break + except httpx.HTTPError: + pass + await asyncio.sleep(1) + else: + logger.warning("Duplex warmup: /health never became ready; skipping warmup.") + return + async with websockets.connect(url, max_size=64 * 1024 * 1024, open_timeout=10) as ws: + await ws.send( + json.dumps( + { + "type": "session.update", + "session": { + "session_id": f"warmup-{uuid.uuid4().hex[:8]}", + "model": model_name, + "modalities": ["audio", "text"], + "input_audio_format": "pcm_f32le", + "output_audio_format": "pcm16", + "idle_timeout_s": 60, + "turn_detection": None, + "extra_body": {"auto_response": True}, + }, + } + ) + ) + saw_audio = False + sent = 0 + + async def _recv_until(predicate, timeout_s: float) -> bool: + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + try: + raw = await asyncio.wait_for(ws.recv(), timeout=max(0.05, deadline - time.monotonic())) + except (TimeoutError, asyncio.TimeoutError): + return False + event = json.loads(raw) + if predicate(event): + return True + return False + + if not await _recv_until(lambda e: e.get("type") == "session.created", 30): + logger.warning("Duplex warmup: no session.created within 30 s; aborting warmup.") + return + while sent < warmup_frames: + await ws.send(json.dumps({"type": "input_audio_buffer.append", "audio": silence})) + sent += 1 + await asyncio.sleep(0.08) + # Wait for the pipeline's first audio output so every stage ran. + saw_audio = await _recv_until(lambda e: e.get("type") == "response.audio.delta", 30) + # Close the session EXPLICITLY: a bare websocket close parks the + # session in its disconnect grace window, where it would keep + # occupying the max_sessions=1 slot and block the first client. + await ws.send(json.dumps({"type": "session.close"})) + await _recv_until(lambda e: e.get("type") == "session.closed", 10) + logger.info( + "Duplex warmup finished in %.1f s (%d silent frames, first audio %s).", + time.time() - start, + sent, + "received" if saw_audio else "NOT received", + ) + except Exception: + logger.exception("Duplex warmup failed; continuing to serve (first session pays cold-start costs).") + finally: + if warmup_done is not None: + warmup_done.set() + + +__all__ = ["_warmup_duplex_realtime"] diff --git a/vllm_omni/entrypoints/openai/README.md b/vllm_omni/entrypoints/openai/README.md new file mode 100644 index 00000000000..0bdadb85fd9 --- /dev/null +++ b/vllm_omni/entrypoints/openai/README.md @@ -0,0 +1,58 @@ +# OpenAI Entrypoint Helpers + +This package contains OpenAI-compatible serving code for vLLM-Omni. + +## Hard Rule: `*_utils*` vs package `helpers.py` + +These are **different layers**. During the #5227 refactor they coexist on purpose. +**Do not mix them. No ownership overlap.** + +| Layer | Files | Longevity | PUT HERE | DO NOT PUT HERE | +| --- | --- | --- | --- | --- | +| Shared media / serving utils | root `image_api_utils.py`, `video_api_utils.py`, `audio_utils_mixin.py`, `utils.py` | **Temporary.** TODO(#5227, Phase 1): tidy up / absorb into endpoint-family packages in Phase 1 modality PRs (P1.1-P1.3) | Codec / bytes / format / pure media conversion; cross-cutting stage/pipeline helpers reused by protocol + multiple serving paths | HTTP form parsing, FastAPI `Request` / app-state orchestration, job runners, endpoint response factories from `api_server.py` | +| Endpoint-family helpers | `images/helpers.py`, `video/generation/helpers.py`, … | **Longer home** through helper (P0.2) + router (P0.3) stages; Phase 1 PRs may split further, not delete ownership | Request/job/runtime helpers peeled out of `api_server.py` for one endpoint family | Encode/decode/base64/media backends that already belong in root `*_utils*` | + +### Allowed dependency direction + +```text +api_server / serving / protocol + │ + ├──► package helpers.py (endpoint orchestration; longer home) + │ │ + │ └──► root *_utils* (shared media / format; temporary) + │ + └──► root *_utils* (direct use OK) +``` + +- `helpers` **may import** `*_utils*` +- `*_utils*` **must not import** package `helpers` +- Never copy a symbol into both layers; pick one owner + +### Where to put new code (until family PRs finish) + +1. Pure image/video/audio bytes or format conversion with **no** FastAPI `Request` / job store / route semantics → root `*_utils*` only if it must stay shared; prefer not growing these files. +2. Endpoint request parsing, validation against engine/app state, job lifecycle, or response shaping that used to sit beside routes in `api_server.py` → owning package `helpers.py` (**prefer this**). +3. Cross-family OpenAI app-state / bootstrap → `app_state.py` / similar top-level OpenAI helpers — **not** into modality `*_utils*`. +4. If unsure: **do not grow root `*_utils*`**. Put it in the owning package `helpers.py` and leave a TODO(#5227) for the modality PR. + +**TODO(#5227, Phase 1):** root `image_api_utils.py`, `video_api_utils.py`, `audio_utils_mixin.py`, and likely pieces of `utils.py` should be tidied, moved, or absorbed during Phase 1 audio / images / video family PRs (P1.1-P1.3). Package `helpers.py` files stay as the endpoint ownership surface until those PRs further split them. + +## What Belongs Here (top-level OpenAI helpers) + +- OpenAI protocol or response-shape helpers shared by multiple OpenAI endpoint + families. +- App-state accessors for OpenAI serving objects, for example helpers that read + `request.app.state.openai_serving_*`. +- Bootstrap helpers that are specific to OpenAI-compatible serving and are not + route bodies. + +## What Does Not Belong Here + +- FastAPI app/server mechanics that are not OpenAI-specific. Put those under + `vllm_omni.entrypoints.serve.utils`. +- Modality-specific request parsing, validation, or response formatting. Put + those under the owning endpoint package, such as `images/` or + `video/generation/`. +- Broad cleanup/rename of existing `image_api_utils.py`, `video_api_utils.py`, + `audio_utils_mixin.py`, and `utils.py` in this helper-only stage (P0.2) is out + of scope; tracked by #5227 Phase 1 modality PRs instead. diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index e020be7bf87..adb7d6e9daf 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -1,9 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""OpenAI-compatible API server bootstrap for vLLM-Omni. + +This module owns app construction, server startup, app-state initialization, +and route bodies that have not yet moved to endpoint-owned modules.""" + import asyncio -import base64 import dataclasses -import io import json import multiprocessing import multiprocessing.forkserver as forkserver @@ -11,24 +14,17 @@ # Image generation API imports import random -import tempfile import time -import uuid from argparse import Namespace -from collections.abc import AsyncIterator, Mapping, Sequence -from contextlib import asynccontextmanager, suppress +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager from http import HTTPStatus -from numbers import Integral -from pathlib import Path -from typing import Annotated, Any, Literal, cast +from typing import Annotated, Any, Literal -import httpx -import numpy as np import vllm.envs as envs from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, Request, UploadFile, WebSocket from fastapi.responses import FileResponse, JSONResponse, Response, StreamingResponse from PIL import Image -from pydantic import BaseModel, Field from starlette.datastructures import State from starlette.types import ASGIApp, Receive, Scope, Send from vllm.engine.protocol import EngineClient @@ -36,7 +32,7 @@ from vllm.entrypoints.chat_utils import ChatTemplateConfig, load_chat_template from vllm.entrypoints.generate.base.protocol import RequestResponseMetadata from vllm.entrypoints.launchers.cli_args import make_arg_parser -from vllm.entrypoints.launchers.launcher import serve_http, terminate_if_errored +from vllm.entrypoints.launchers.launcher import serve_http from vllm.entrypoints.launchers.utils.server_utils import get_uvicorn_log_config from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer from vllm.entrypoints.openai.api_server import build_app as build_openai_app @@ -59,17 +55,10 @@ from vllm.entrypoints.pooling.pooling.serving import ServingPooling from vllm.entrypoints.pooling.scoring.serving import ServingScores from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens -from vllm.entrypoints.serve import create_error_response -from vllm.entrypoints.serve.engine.protocol import ( - ErrorResponse, - ModelCard, - ModelList, - ModelPermission, -) +from vllm.entrypoints.serve.engine.protocol import ErrorResponse # vLLM moved `base` from openai.basic.api_router to serve.instrumentator.basic. # Keep a fallback for older/newer upstream layouts during rebase windows. -from vllm.entrypoints.serve.instrumentator.basic import base from vllm.entrypoints.serve.tokenize.serving import ServingTokenization from vllm.entrypoints.serve.utils.api_utils import ( load_aware_call, @@ -95,21 +84,49 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm_omni.config.endpoint_policy import ( - remove_route_from_app, shutdown_unsupported_routes, ) -from vllm_omni.diffusion.models.interface import ReferenceVideoDecodeSpec from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.duplex.capability import should_enable_duplex_endpoint from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler +from vllm_omni.entrypoints.duplex.warmup import _warmup_duplex_realtime +from vllm_omni.entrypoints.openai import app_state as openai_app_state +from vllm_omni.entrypoints.openai.app_state import ( + ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL, + OmniAudioGenerate, + OmniBatchChat, + Omnichat, + Omnispeech, + _get_engine_and_model, +) from vllm_omni.entrypoints.openai.batch_serving import OmniOpenAIServingChatBatch -from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError +from vllm_omni.entrypoints.openai.chat_template import _load_model_chat_template_json +from vllm_omni.entrypoints.openai.diffusion import ( + MAX_UINT32_SEED, + _generate_with_async_omni, + apply_stage_default_sampling_params, +) +from vllm_omni.entrypoints.openai.errors import ( + _create_speech_error_json_response, + _error_response_to_json_response, +) from vllm_omni.entrypoints.openai.image_api_utils import ( SUPPORTED_LAYERED_RESOLUTIONS, encode_image_base64_with_compression, parse_size, validate_layered_layers, ) +from vllm_omni.entrypoints.openai.images.helpers import ( + _build_hunyuan_edit_extra_args, + _check_max_generated_image_size, + _choose_output_format, + _extract_images_from_result, + _get_max_edit_input_images, + _load_input_images, + _update_if_not_none, +) +from vllm_omni.entrypoints.openai.lora import _get_lora_from_json_str, _parse_lora_request +from vllm_omni.entrypoints.openai.models import serving as openai_models_serving from vllm_omni.entrypoints.openai.protocol.audio import ( BatchSpeechRequest, OpenAICreateAudioGenerateRequest, @@ -122,11 +139,7 @@ ResponseFormat, ) from vllm_omni.entrypoints.openai.protocol.videos import ( - SecondStr, - SizeStr, - VideoAction, VideoDeleteResponse, - VideoError, VideoGenerationRequest, VideoGenerationStatus, VideoListResponse, @@ -145,21 +158,30 @@ ) from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler from vllm_omni.entrypoints.openai.serving_video_stream import create_streaming_video_handler -from vllm_omni.entrypoints.openai.stage_params import ( - build_stage_sampling_params_list, - get_default_sampling_params_list, -) from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER, FileStorageHandle from vllm_omni.entrypoints.openai.stores import VIDEO_STORE, VIDEO_TASKS -from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request -from vllm_omni.entrypoints.openai.video_api_utils import ( - VideoFrames, - _decode_image_bytes, - _validate_image_pixel_limit, - decode_audio_url, - decode_input_reference, +from vllm_omni.entrypoints.openai.utils import get_stage_type +from vllm_omni.entrypoints.openai.video.generation.helpers import ( + VIDEO_SYNC_TIMEOUT_S, + _cleanup_video_references, + _parse_video_form, + _run_video_generation_job, + _status_code_for_video_failure, + _unpack_video_generation_result, + video_response_from_request, ) from vllm_omni.entrypoints.openpi.serving import ServingRealtimeRobotOpenPI +from vllm_omni.entrypoints.serve.omni_control.protocol import OmniSleepRequest, OmniWakeupRequest +from vllm_omni.entrypoints.serve.profile.protocol import ProfileRequest +from vllm_omni.entrypoints.serve.profile.utils import _should_enable_profiler_endpoints +from vllm_omni.entrypoints.serve.utils.errors import ( + _create_engine_error_json_response, + _register_omni_exception_handlers, +) +from vllm_omni.entrypoints.serve.utils.routes import ( + _remove_route_from_router, + remove_route_from_app, +) from vllm_omni.entrypoints.utils import PureDiffusionLauncherAdapter from vllm_omni.errors import OmniClientError from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt @@ -169,393 +191,12 @@ logger = init_logger(__name__) router = APIRouter() -MAX_UINT32_SEED = 2**32 - 1 -MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES = 30 * 1024 * 1024 -MINIMAX_H3_MAX_REFERENCE_VIDEO_BYTES = 50 * 1024 * 1024 -MINIMAX_H3_MAX_REFERENCE_AUDIO_BYTES = 15 * 1024 * 1024 -MINIMAX_H3_MAX_REFERENCE_COUNT = 12 -MINIMAX_H3_REFERENCE_IMAGE_FORMATS = frozenset({"jpeg", "png", "webp", "heic", "heif"}) -MINIMAX_H3_REFERENCE_VIDEO_SUFFIXES = frozenset({".mp4", ".mov"}) -MINIMAX_H3_REFERENCE_AUDIO_SUFFIXES = frozenset({".wav", ".mp3"}) -CONTROL_REFERENCE_IMAGE_SUFFIXES = frozenset({".bmp", ".gif", ".jpg", ".jpeg", ".png", ".tif", ".tiff", ".webp"}) -CONTROL_REFERENCE_VIDEO_SUFFIXES = frozenset({".mkv", ".mov", ".mp4", ".webm"}) -CONTROL_REFERENCE_MAX_BYTES = 512 * 1024 * 1024 profiler_router = APIRouter() -def _load_model_chat_template_json(model: str) -> str | None: - """Load a model-level chat_template.json from a local path or HF cache. - - Some multimodal HF repos, including Qwen3-Omni, ship the chat template as a - separate file instead of embedding it in tokenizer_config.json. Transformers - 4.44+ no longer supplies a default template, so serving must pass that model - template explicitly when the user did not provide --chat-template. - """ - candidate = Path(model) / "chat_template.json" - template_path: str | None = str(candidate) if candidate.is_file() else None - - if template_path is None: - try: - from vllm_omni.transformers_utils.repo_utils import hf_api - - template_path = hf_api().hf_hub_download( - repo_id=model, - filename="chat_template.json", - local_files_only=True, - ) - except Exception: - return None - - try: - with open(template_path, encoding="utf-8") as f: - payload = json.load(f) - except Exception as exc: - logger.warning("Failed to load chat template from %s: %s", template_path, exc) - return None - - if isinstance(payload, dict): - template = payload.get("chat_template") - elif isinstance(payload, str): - template = payload - else: - template = None - - if not isinstance(template, str) or not template.strip(): - logger.warning("Ignoring malformed chat template payload in %s", template_path) - return None - - logger.info("Loaded chat template from %s", template_path) - return template - - -def _should_enable_profiler_endpoints(stage_configs: list | None) -> bool: - """Check if any stage has profiler_config set in its engine_args.""" - if not stage_configs: - return False - for stage in stage_configs: - engine_args = stage.get("engine_args") if isinstance(stage, dict) else getattr(stage, "engine_args", None) - if engine_args is None: - continue - profiler_config = ( - engine_args.get("profiler_config") - if isinstance(engine_args, dict) - else getattr(engine_args, "profiler_config", None) - ) - if profiler_config is not None: - profiler = ( - profiler_config.get("profiler") - if isinstance(profiler_config, dict) - else getattr(profiler_config, "profiler", None) - ) - if profiler is not None: - return True - return False - - -class ProfileRequest(BaseModel): - """Request model for profiling endpoints.""" - - stages: list[int] | None = Field( - default=None, - description="List of stage IDs to profile. If None, profiles all stages.", - ) - - -def _remove_route_from_router( - router: APIRouter, - path: str, - methods: set[str] | None = None, -) -> None: - methods_set = {method.upper() for method in methods} if methods else None - for route in list(router.routes): - if getattr(route, "path", None) != path: - continue - if methods_set is not None: - route_methods = {method.upper() for method in (getattr(route, "methods", None) or set())} - if not (route_methods & methods_set): - continue - router.routes.remove(route) - - -ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL = "endpoint-load-metrics-format" - - -async def _get_vllm_config(engine_client: EngineClient) -> Any: - if hasattr(engine_client, "get_vllm_config"): - return await engine_client.get_vllm_config() - return getattr(engine_client, "vllm_config", None) - - -def _register_omni_exception_handlers(app) -> None: - """Override upstream vLLM exception handlers with Omni-aware versions. - - The upstream ``engine_error_handler`` is designed for ``AsyncLLM`` (single - EngineCore process). Omni uses a multi-stage orchestrator with different - health semantics, so we register our own handlers that: - - - Log multi-stage diagnostic info (orchestrator liveness, per-stage health) - when an ``EngineDeadError`` is caught. - - Call ``terminate_if_errored`` - - Return an OpenAI-compatible error JSON response. - """ - - async def omni_engine_error_handler( - req: Request, - exc: EngineDeadError | EngineGenerateError, - ): - request_id = _get_request_id_from_request(req) - - if req.app.state.args.log_error_stack: - logger.exception("Engine Exception caught. Request id: %s", request_id) - - return _create_engine_error_json_response(req, exc) - - app.exception_handler(EngineGenerateError)(omni_engine_error_handler) - app.exception_handler(EngineDeadError)(omni_engine_error_handler) - - -def _get_request_id_from_request(req: Request) -> str | None: - return req.state.request_metadata.request_id if hasattr(req.state, "request_metadata") else None - - -def _build_engine_error_payload( - exc: EngineDeadError | EngineGenerateError, - *, - request_id: str | None, -) -> tuple[dict[str, Any], int]: - err = create_error_response(exc) - payload = err.model_dump() - error_body = payload.get("error", {}) - - error_body["request_id"] = request_id - error_body["error_stage_id"] = getattr(exc, "error_stage_id", None) - - return payload, err.error.code - - -def _create_engine_error_json_response( - req: Request, - exc: EngineDeadError | EngineGenerateError, -) -> JSONResponse: - request_id = _get_request_id_from_request(req) - error_stage_id = getattr(exc, "error_stage_id", None) - engine = req.app.state.engine_client - - if isinstance(exc, EngineDeadError): - # Log Omni-specific diagnostic information for dead engines. - orchestrator_alive = engine.engine.is_alive() if hasattr(engine, "engine") else "N/A" - logger.error( - "EngineDeadError: orchestrator_alive=%s, errored=%s, request_id=%s, error_stage_id=%s", - orchestrator_alive, - engine.errored, - request_id, - error_stage_id, - ) - - terminate_if_errored( - server=req.app.state.server, - engine=engine, - ) - - payload, status_code = _build_engine_error_payload(exc, request_id=request_id) - return JSONResponse(content=payload, status_code=status_code) - - -def _error_response_to_json_response( - err: ErrorResponse, - *, - status_code: HTTPStatus | int | None = None, - default_status_code: HTTPStatus | int = HTTPStatus.BAD_REQUEST, -) -> JSONResponse: - resolved_status = int( - status_code - if status_code is not None - else (err.error.code if err.error and err.error.code is not None else default_status_code) - ) - payload = err.model_dump() - if err.error: - payload["error"]["code"] = resolved_status - return JSONResponse(content=payload, status_code=resolved_status) - - -def _create_speech_error_json_response( - raw_request: Request, - message: str, - *, - err_type: str = "BadRequestError", - status_code: HTTPStatus = HTTPStatus.BAD_REQUEST, -) -> JSONResponse: - err = base(raw_request).create_error_response( - message=message, - err_type=err_type, - status_code=status_code, - ) - return _error_response_to_json_response(err, status_code=status_code) - - -class _DiffusionServingModels: - """Minimal OpenAIServingModels implementation for diffusion-only servers. - - vLLM's /v1/models route expects `app.state.openai_serving_models` to expose - `show_available_models()`. In pure diffusion mode we don't initialize the - full OpenAIServingModels (it depends on LLM-specific processors), so we - provide a lightweight fallback. - """ - - class _NullModelConfig: - def __getattr__(self, name): - return None - - class _Unsupported: - def __init__(self, name: str): - self.name = name - - def __call__(self, *args, **kwargs): - raise NotImplementedError(f"{self.name} is not supported in diffusion mode") - - def __getattr__(self, attr): - raise NotImplementedError(f"{self.name}.{attr} is not supported in diffusion mode") - - def __init__(self, base_model_paths: list[BaseModelPath]) -> None: - self._base_model_paths = base_model_paths - self.model_config = self._NullModelConfig() - - @property - def base_model_paths(self) -> list[BaseModelPath]: - return self._base_model_paths - - def is_base_model(self, model_name: str) -> bool: - return any(p.name == model_name for p in self._base_model_paths) - - def __getattr__(self, name): - return self._Unsupported(name) - - async def show_available_models(self) -> ModelList: - return ModelList( - data=[ - ModelCard( - id=base_model.name, - root=base_model.model_path, - permission=[ModelPermission()], - ) - for base_model in self._base_model_paths - ] - ) - - # Server entry points -async def _warmup_duplex_realtime(app, args, warmup_frames: int) -> None: - """Run silent frames through a throwaway realtime session at startup. - - Self-connects over the server's own websocket endpoint so the warmup - exercises the exact production path (serving adapter, duplex control - plane, every pipeline stage). One-time costs — kernel JIT compilation, - first prefill/decode paths, codec caches — are paid here instead of by - the first real client; ``/v1/realtime`` holds other connections until - the warmup finishes (see ``duplex_warmup_done``). - """ - import base64 - - warmup_done = getattr(app.state, "duplex_warmup_done", None) - try: - try: - import websockets - except ImportError: - logger.warning("Duplex warmup skipped: the 'websockets' package is not installed.") - return - served = getattr(args, "served_model_name", None) - if isinstance(served, list | tuple) and served: - model_name = served[0] - elif isinstance(served, str) and served: - model_name = served - else: - model_name = args.model - adapter = getattr(getattr(app.state, "openai_serving_duplex", None), "_serving_runtime_adapter", None) - frame_samples = int(getattr(adapter, "silence_continuation_samples", 16000)) - silence = base64.b64encode(bytes(frame_samples * 4)).decode("ascii") - from vllm_omni.experimental.fullduplex.client import build_realtime_url - - url = ( - build_realtime_url(f"ws://127.0.0.1:{args.port}/v1/realtime", model_name, autostart=False) - + "&vllm_omni_warmup=1" - ) - logger.info("Duplex warmup starting: %d silent frames of %d samples.", warmup_frames, frame_samples) - start = time.time() - # The socket is bound before uvicorn finishes starting, so a plain - # connect can be accepted and then reset. Wait for /health first. - async with httpx.AsyncClient() as http: - for _ in range(120): - try: - if (await http.get(f"http://127.0.0.1:{args.port}/health", timeout=2)).status_code == 200: - break - except httpx.HTTPError: - pass - await asyncio.sleep(1) - else: - logger.warning("Duplex warmup: /health never became ready; skipping warmup.") - return - async with websockets.connect(url, max_size=64 * 1024 * 1024, open_timeout=10) as ws: - await ws.send( - json.dumps( - { - "type": "session.update", - "session": { - "session_id": f"warmup-{uuid.uuid4().hex[:8]}", - "model": model_name, - "modalities": ["audio", "text"], - "input_audio_format": "pcm_f32le", - "output_audio_format": "pcm16", - "idle_timeout_s": 60, - "turn_detection": None, - "extra_body": {"auto_response": True}, - }, - } - ) - ) - saw_audio = False - sent = 0 - - async def _recv_until(predicate, timeout_s: float) -> bool: - deadline = time.monotonic() + timeout_s - while time.monotonic() < deadline: - try: - raw = await asyncio.wait_for(ws.recv(), timeout=max(0.05, deadline - time.monotonic())) - except (TimeoutError, asyncio.TimeoutError): - return False - event = json.loads(raw) - if predicate(event): - return True - return False - - if not await _recv_until(lambda e: e.get("type") == "session.created", 30): - logger.warning("Duplex warmup: no session.created within 30 s; aborting warmup.") - return - while sent < warmup_frames: - await ws.send(json.dumps({"type": "input_audio_buffer.append", "audio": silence})) - sent += 1 - await asyncio.sleep(0.08) - # Wait for the pipeline's first audio output so every stage ran. - saw_audio = await _recv_until(lambda e: e.get("type") == "response.audio.delta", 30) - # Close the session EXPLICITLY: a bare websocket close parks the - # session in its disconnect grace window, where it would keep - # occupying the max_sessions=1 slot and block the first client. - await ws.send(json.dumps({"type": "session.close"})) - await _recv_until(lambda e: e.get("type") == "session.closed", 10) - logger.info( - "Duplex warmup finished in %.1f s (%d silent frames, first audio %s).", - time.time() - start, - sent, - "received" if saw_audio else "NOT received", - ) - except Exception: - logger.exception("Duplex warmup failed; continuing to serve (first session pays cold-start costs).") - finally: - if warmup_done is not None: - warmup_done.set() - - async def omni_run_server(args, **uvicorn_kwargs) -> None: """Run a single-worker API server. @@ -650,7 +291,7 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, remove_route_from_app(app, "/stop_profile", frozenset({"POST"})) app.include_router(profiler_router) - vllm_config = await _get_vllm_config(engine_client) + vllm_config = await openai_app_state._get_vllm_config(engine_client) # Check if pure diffusion mode (vllm_config will be None) is_pure_diffusion = vllm_config is None @@ -875,7 +516,7 @@ async def omni_init_app_state( args: Parsed command-line arguments """ # Get vllm_config from engine_client (following 0.14.0 pattern) - vllm_config = await _get_vllm_config(engine_client) + vllm_config = await openai_app_state._get_vllm_config(engine_client) # Detect if it's pure Diffusion mode (single stage and is Diffusion) is_pure_diffusion = False @@ -911,7 +552,7 @@ async def omni_init_app_state( if is_pure_diffusion: state.vllm_config = None state.diffusion_engine = engine_client - state.openai_serving_models = _DiffusionServingModels(base_model_paths) + state.openai_serving_models = openai_models_serving._DiffusionServingModels(base_model_paths) # OMNI: tokenization endpoints are not supported in pure diffusion mode. state.serving_tokenization = None @@ -970,7 +611,7 @@ async def omni_init_app_state( # LLM or multi-stage mode: use standard initialization logic if vllm_config is None: # Try to get vllm_config from engine_client - vllm_config = await _get_vllm_config(engine_client) + vllm_config = await openai_app_state._get_vllm_config(engine_client) if vllm_config is None: logger.warning("vllm_config is None, some features may not work correctly") @@ -1311,26 +952,6 @@ async def omni_init_app_state( state.server_load_metrics = 0 -def Omnivideo(request: Request) -> OmniOpenAIServingVideo | None: - return request.app.state.openai_serving_video - - -def Omnichat(request: Request) -> OmniOpenAIServingChat | None: - return request.app.state.openai_serving_chat - - -def OmniBatchChat(request: Request) -> OmniOpenAIServingChatBatch | None: - return request.app.state.openai_serving_chat_batch - - -def Omnispeech(request: Request) -> OmniOpenAIServingSpeech | None: - return request.app.state.openai_serving_speech - - -def OmniAudioGenerate(request: Request) -> OmniOpenAIServingAudioGenerate | None: - return getattr(request.app.state, "openai_serving_audio_generate", None) - - @router.post( "/v1/chat/completions", dependencies=[Depends(validate_json_request)], @@ -2584,1237 +2205,6 @@ async def edit_images( raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Image edit failed: {str(e)}") -def _get_engine_and_model(raw_request: Request): - # Get engine client (AsyncOmni) from app state - engine_client: EngineClient | AsyncOmni | None = getattr(raw_request.app.state, "engine_client", None) - if engine_client is None: - raise HTTPException( - status_code=HTTPStatus.SERVICE_UNAVAILABLE.value, - detail="Multi-stage engine not initialized. Start server with a multi-stage omni model.", - ) - - # Check if there's a diffusion stage. - # Prefer app state (compat layer populated at startup), then fall back to - # the engine client's stage configs for refactored AsyncOmni paths. - stage_configs = getattr(raw_request.app.state, "stage_configs", None) - if not stage_configs: - stage_configs = getattr(engine_client, "stage_configs", None) - if not stage_configs: - raise HTTPException( - status_code=HTTPStatus.SERVICE_UNAVAILABLE.value, - detail="Stage configs not found. Start server with a multi-stage omni model.", - ) - - normalized_stage_configs = list(stage_configs) - has_diffusion_stage = any(get_stage_type(stage_cfg) == "diffusion" for stage_cfg in normalized_stage_configs) - - if not has_diffusion_stage: - raise HTTPException( - status_code=HTTPStatus.SERVICE_UNAVAILABLE.value, - detail="No diffusion stage found in multi-stage pipeline.", - ) - - # Get server's loaded model name - serving_models = getattr(raw_request.app.state, "openai_serving_models", None) - base_model_paths = getattr(serving_models, "base_model_paths", None) if serving_models else None - if base_model_paths: - model_name = base_model_paths[0].name - else: - model_name = "unknown" - - return engine_client, model_name, normalized_stage_configs - - -def _get_diffusion_od_config(raw_request: Request, engine_client: Any) -> Any: - diffusion_engine = getattr(raw_request.app.state, "diffusion_engine", None) or engine_client - get_diffusion_od_config = getattr(diffusion_engine, "get_diffusion_od_config", None) - return ( - get_diffusion_od_config() if callable(get_diffusion_od_config) else getattr(diffusion_engine, "od_config", None) - ) - - -def _get_max_edit_input_images(raw_request: Request, engine_client: Any) -> int | None: - od_config = _get_diffusion_od_config(raw_request, engine_client) - if od_config is None: - # Preserve the existing compatibility behavior when the diffusion - # config is not exposed on the serving surface. - return None - - supports_multimodal_inputs = getattr(od_config, "supports_multimodal_inputs", None) - if not isinstance(supports_multimodal_inputs, bool): - # Older serving surfaces and mocked engines may expose a placeholder - # object instead of a real diffusion config. Treat that as "unknown" - # so existing single-image flows keep working. - return None - - if not supports_multimodal_inputs: - return 1 - - max_input_images = getattr(od_config, "max_multimodal_image_inputs", None) - if max_input_images is None: - return None - if isinstance(max_input_images, bool) or not isinstance(max_input_images, Integral): - return None - if max_input_images < 1: - return None - return int(max_input_images) - - -def _get_lora_from_json_str(lora_body): - if lora_body is None: - return None - try: - lora_dict = json.loads(lora_body) - except json.JSONDecodeError: - raise HTTPException(status_code=400, detail="Invalid LoRA JSON string") - - if not isinstance(lora_dict, dict): - raise HTTPException(status_code=400, detail="LoRA must be a JSON object") - - return lora_dict - - -def _parse_lora_request(lora_body: dict[str, Any]): - try: - return parse_lora_request(lora_body) - except ValueError as e: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail=str(e), - ) from e - - -async def _generate_with_async_omni( - engine_client: AsyncOmni | Any, - gen_params: Any, - stage_configs: list[Any], - **kwargs, -): - engine_client = cast(AsyncOmni, engine_client) - result = None - normalized_stage_configs = list(stage_configs) - if not normalized_stage_configs: - raise HTTPException( - status_code=HTTPStatus.SERVICE_UNAVAILABLE.value, - detail="Stage configs not found. Start server with a multi-stage omni model.", - ) - sampling_params_list = build_stage_sampling_params_list( - normalized_stage_configs, - get_default_sampling_params_list(engine_client), - diffusion_params=gen_params, - replace_diffusion_params=True, - ) - - async for output in engine_client.generate( - sampling_params_list=sampling_params_list, - **kwargs, - ): - result = output - - if result is None: - raise HTTPException( - status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, - detail="No output generated from multi-stage pipeline.", - ) - return result - - -def _check_max_generated_image_size( - app_state_args: Any, - width: int | None, - height: int | None, - resolution: int | None = None, -) -> None: - """Raise 400 if the requested image size exceeds --max-generated-image-size.""" - max_generated_image_size = getattr(app_state_args, "max_generated_image_size", None) - # Check max_generated_image_size - if max_generated_image_size is None: - return - if width is not None and height is not None: - if width * height > max_generated_image_size: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail=f"Requested image size {width}x{height} exceeds the maximum allowed " - f"size of {max_generated_image_size} pixels. You can reduce the requested size " - f"or increase the server's --max-generated-image-size limit.", - ) - elif resolution is not None: - # When resolution is set, the output size is resolution * resolution - if resolution * resolution > max_generated_image_size: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail=f"Requested resolution {resolution} (max {resolution}x{resolution} pixels) " - f"exceeds the maximum allowed size of {max_generated_image_size} pixels. " - f"You can reduce the requested size or increase the server's --max-generated-image-size limit.", - ) - - -def _build_hunyuan_edit_extra_args( - *, - bot_task: str | None, - sys_type: str | None, - system_prompt: str | None, -) -> dict[str, Any]: - """Map Hunyuan /v1/images/edits form fields to DiT ``extra_args``.""" - extra_args: dict[str, Any] = {} - effective_use_system_prompt = sys_type - if effective_use_system_prompt is None and bot_task is not None: - from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import resolve_sys_type - - effective_use_system_prompt = resolve_sys_type(bot_task) - if effective_use_system_prompt is not None: - extra_args["use_system_prompt"] = effective_use_system_prompt - if system_prompt is not None: - extra_args["system_prompt"] = system_prompt - if bot_task is not None: - extra_args["bot_task"] = bot_task - return extra_args - - -def _update_if_not_none(object: Any, key: str, val: Any) -> None: - if val is not None: - setattr(object, key, val) - - -def _normalize_image(image: Any) -> Any: - """Normalize a single image output to a PIL-compatible format.""" - if isinstance(image, Image.Image): - return image - if not isinstance(image, np.ndarray): - raise ValueError(f"Unsupported image type: {type(image)}") - if not np.issubdtype(image.dtype, np.integer) and not np.issubdtype(image.dtype, np.floating): - raise ValueError(f"Unsupported dtype: {image.dtype}") - if isinstance(image, np.ndarray): - while image.ndim > 3: - image = image[0] - if image.min() < 0: - if image.min() < -1.01 or image.max() > 1.01: - logger.warning( - f"Image float range [{image.min():.2f}, {image.max():.2f}] outside expected [-1, 1]. " - f"Clipping to [-1, 1] before normalization." - ) - image = np.clip(image, -1.0, 1.0) * 0.5 + 0.5 - elif image.max() > 1.01: - logger.warning( - f"Image float range [{image.min():.2f}, {image.max():.2f}] outside expected [0, 1]. " - f"Clipping to [0, 1] before normalization." - ) - image = (np.clip(image, 0.0, 1.0) * 255).astype(np.uint8) - image = Image.fromarray(image) - return image - - -def _extract_images_from_result(result: Any) -> list[Any]: - images = [] - if hasattr(result, "images") and result.images: - images = result.images - # Handle when generate more than one image - if images and isinstance(images[0], np.ndarray) and images[0].shape[0] > 1 and images[0].ndim == 5: - # Unwrap batch: (N, T, H, W, C) -> [img1, img2, ...] - images = list(images[0]) - # Flatten nested lists (e.g., from layered models like Qwen-Image-Layered). - # Note: This only flattens one level deep. Deeper nesting is not supported. - flattened = [] - for img in images: - if isinstance(img, list): - flattened.extend(img) - else: - flattened.append(img) - return [_normalize_image(img) for img in flattened] - - -async def _load_input_images( - inputs: list[str], - *, - normalize_rgb: bool = True, -) -> list[Image.Image]: - """ - convert to PIL.Image.Image list - """ - if isinstance(inputs, str): - inputs = [inputs] - - images: list[Image.Image] = [] - - for inp in inputs: - # 1. URL + base64 - if isinstance(inp, str) and inp.startswith("data:image"): - try: - _, b64_data = inp.split(",", 1) - image_bytes = base64.b64decode(b64_data) - img = Image.open(io.BytesIO(image_bytes)) - images.append(img) - except Exception as e: - raise ValueError(f"Invalid base64 image: {e}") - - # 2. URL - elif isinstance(inp, str) and inp.startswith("http"): - async with httpx.AsyncClient(timeout=60) as client: - try: - resp = await client.get(inp) - resp.raise_for_status() - img = Image.open(io.BytesIO(resp.content)) - images.append(img) - except Exception as e: - raise ValueError(f"Failed to download image from URL {inp}: {e}") - - # 3. UploadFile - elif hasattr(inp, "file"): - try: - img_data = await inp.read() - img = Image.open(io.BytesIO(img_data)) - images.append(img) - except Exception as e: - raise ValueError(f"Failed to open uploaded file: {e}") - - else: - raise ValueError(f"Unsupported input: {inp}") - - if not images: - raise ValueError("No valid input images found") - - if not normalize_rgb: - return images - - # Match the offline HunyuanImage3 image-edit example path, which eagerly - # normalizes input files with ``Image.open(...).convert("RGB")`` before - # they reach the AR stage. Keeping uploads as RGBA/P PIL objects makes - # online IT2I observe a different visual input than offline (for example - # transparent-logo uploads alpha-composited over white instead of black), - # which is enough for HunyuanImage3 AR recaption to diverge before DiT - # sees the request -- root cause of the "online 3 magnets vs offline 1 - # magnet" systematic semantic mismatch. - return [img.convert("RGB") for img in images] - - -def _choose_output_format(output_format: str | None, background: str | None) -> str: - # Normalize and choose extension - fmt = (output_format or "").lower() - if fmt in {"jpg", "png", "webp", "jpeg"}: - return fmt - # If transparency requested, prefer png - if (background or "auto").lower() == "transparent": - return "png" - # Default - return "jpeg" - - -def apply_stage_default_sampling_params( - default_params_json: str | None, - sampling_params: Any, - stage_key: str, -) -> None: - """ - Update a stage's sampling parameters with vLLM-Omni defaults. - - Args: - default_params_json: JSON string of stage-keyed default parameters - sampling_params: The sampling parameters object to update - stage_key: The stage ID/key in the pipeline - """ - if default_params_json is not None: - default_params_dict = json.loads(default_params_json) - if stage_key in default_params_dict: - stage_defaults = default_params_dict[stage_key] - for param_name, param_value in stage_defaults.items(): - if hasattr(sampling_params, param_name): - setattr(sampling_params, param_name, param_value) - - -def _resolve_video_runtime_context(raw_request: Request) -> tuple[str | None, list[Any] | None]: - app_model_name = None - serving_models = getattr(raw_request.app.state, "openai_serving_models", None) - if serving_models and getattr(serving_models, "base_model_paths", None): - base_paths = serving_models.base_model_paths - if base_paths: - app_model_name = base_paths[0].name - - app_stage_configs = getattr(raw_request.app.state, "stage_configs", None) - return app_model_name, app_stage_configs - - -def _parse_form_json(value: str | None, expected_type: type | None = None) -> Any: - if value is None or value == "": - return None - try: - parsed = json.loads(value) - except json.JSONDecodeError as exc: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail="Invalid JSON in form field.", - ) from exc - if expected_type is not None and not isinstance(parsed, expected_type): - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail=f"Invalid JSON in form field: expected {expected_type.__name__}, got {type(parsed).__name__}.", - ) - return parsed - - -def _config_get(config: Any, key: str, default: Any = None) -> Any: - if isinstance(config, dict): - return config.get(key, default) - if hasattr(config, "get"): - try: - return config.get(key, default) - except Exception: - pass - return getattr(config, key, default) - - -def _stage_engine_args(stage_cfg: Any) -> Any: - return _config_get(stage_cfg, "engine_args", {}) or {} - - -def _diffusion_model_classes(stage_configs: list[Any] | None) -> list[type]: - if not stage_configs: - return [] - - from vllm_omni.diffusion.registry import DiffusionModelRegistry - - model_classes: list[type] = [] - for stage_cfg in stage_configs: - if get_stage_type(stage_cfg) != "diffusion": - continue - model_class_name = _config_get(_stage_engine_args(stage_cfg), "model_class_name") - if not model_class_name: - continue - model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name) - if model_cls is not None: - model_classes.append(model_cls) - return model_classes - - -def _normalize_reference_video_decode_spec(spec: ReferenceVideoDecodeSpec) -> ReferenceVideoDecodeSpec: - max_frames = spec.max_frames - if max_frames is not None: - try: - max_frames = int(max_frames) - except (TypeError, ValueError) as exc: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail="Invalid reference video decode spec: max_frames must be an integer.", - ) from exc - if max_frames <= 0: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail="Invalid reference video decode spec: max_frames must be positive.", - ) - - keep = str(spec.keep or "first").strip().lower() - if keep not in {"first", "last"}: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail="Invalid reference video decode spec: keep must be either 'first' or 'last'.", - ) - return ReferenceVideoDecodeSpec(max_frames=max_frames, keep=cast(Literal["first", "last"], keep)) - - -def _reference_video_decode_spec( - req: VideoGenerationRequest, - stage_configs: list[Any] | None, -) -> ReferenceVideoDecodeSpec: - video_params = req.resolve_video_params() - extra_params = req.extra_params if isinstance(req.extra_params, dict) else {} - for model_cls in _diffusion_model_classes(stage_configs): - resolver = getattr(model_cls, "reference_video_decode_spec", None) - if resolver is None: - continue - try: - spec = resolver(num_frames=video_params.num_frames, extra_args=extra_params) - except (TypeError, ValueError) as exc: - raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(exc)) from exc - if spec is not None: - return _normalize_reference_video_decode_spec(spec) - return ReferenceVideoDecodeSpec(max_frames=video_params.num_frames, keep="first") - - -def video_response_from_request(model_name: str, req: VideoGenerationRequest) -> VideoResponse: - video_params = req.resolve_video_params() - duration_s = None - if req.seconds is not None: - duration_s = float(req.seconds) - elif video_params.num_frames is not None and video_params.fps is not None: - duration_s = video_params.num_frames / video_params.fps - - resp = VideoResponse( - model=model_name, - status=VideoGenerationStatus.QUEUED, - size=req.size, - prompt=req.prompt, - quality=req.quality or "default", - fps=video_params.fps, - num_frames=video_params.num_frames, - duration_s=duration_s, - ) - resp.seconds = str(req.seconds or resp.seconds) - return resp - - -def _status_code_for_video_failure(error: VideoError | None) -> int: - if error is None: - return HTTPStatus.INTERNAL_SERVER_ERROR.value - - if isinstance(error.code, int): - if 400 <= error.code < 600: - return error.code - return HTTPStatus.INTERNAL_SERVER_ERROR.value - - if error.code == "HTTPException": - status_text, _, _ = error.message.partition(":") - try: - status_code = int(status_text) - except ValueError: - return HTTPStatus.INTERNAL_SERVER_ERROR.value - if 400 <= status_code < 600: - return status_code - return HTTPStatus.INTERNAL_SERVER_ERROR.value - - if error.code == "EngineDeadError": - return HTTPStatus.INTERNAL_SERVER_ERROR.value - if error.code == "EngineGenerateError": - return HTTPStatus.INTERNAL_SERVER_ERROR.value - - return HTTPStatus.INTERNAL_SERVER_ERROR.value - - -def _video_error_from_exception(exc: Exception) -> VideoError: - if isinstance(exc, HTTPException): - message = str(exc.detail) if exc.detail else str(exc) - return VideoError(code=exc.status_code, message=message) - - if isinstance(exc, OmniClientError): - return VideoError(code=exc.status_code, message=exc.message) - - if isinstance(exc, (EngineGenerateError, EngineDeadError)): - err = create_error_response(exc) - return VideoError(code=err.error.code, message=err.error.message) - - return VideoError( - code=HTTPStatus.INTERNAL_SERVER_ERROR.value, - message=str(exc), - ) - - -async def _cleanup_video(video_id: str): - try: - await STORAGE_MANAGER.delete(video_id) - except Exception: - logger.warning("Failed to cleanup partial video file '%s'", video_id) - - -def _cleanup_video_references( - reference_video: ReferenceVideo | None, - reference_audio: ReferenceAudio | None, - control_path: str | None = None, -) -> None: - if reference_video is not None: - for path in reference_video.cleanup_paths: - if os.path.exists(path): - os.unlink(path) - if reference_audio is not None: - cleanup_paths = reference_audio.cleanup_paths or tuple(_reference_list(reference_audio.path)) - for path in cleanup_paths: - if os.path.exists(path): - os.unlink(path) - if control_path is not None and os.path.exists(control_path): - os.unlink(control_path) - - -def _unpack_video_generation_result( - result: Sequence[object], -) -> tuple[bytes, dict[str, float], float, VideoAction | None, dict[str, object]]: - video_metadata: dict[str, object] = {} - if len(result) == 5: - video_bytes, stage_durations, peak_memory_mb, action, raw_metadata = result - if isinstance(raw_metadata, dict): - video_metadata = {str(key): value for key, value in raw_metadata.items()} - else: - video_bytes, stage_durations, peak_memory_mb, action = result - return ( - cast(bytes, video_bytes), - cast(dict[str, float], stage_durations), - float(cast(float, peak_memory_mb)), - cast(VideoAction | None, action), - video_metadata, - ) - - -async def _run_video_generation_job( - handler: OmniOpenAIServingVideo, - request: VideoGenerationRequest, - video_id: str, - reference_image: ReferenceImage | None = None, - reference_video: ReferenceVideo | None = None, - reference_audio: ReferenceAudio | None = None, - control_path: str | None = None, - app_state: Any | None = None, -) -> None: - job = await VIDEO_STORE.get(video_id) - if job is None: - logger.warning("Video job %s missing before generation task started; skipping", video_id) - _cleanup_video_references(reference_video, reference_audio, control_path) - return - - await VIDEO_STORE.update_fields(video_id, {"status": VideoGenerationStatus.IN_PROGRESS}) - started_at = time.perf_counter() - try: - video_bytes, stage_durations, peak_memory_mb, action, video_metadata = _unpack_video_generation_result( - await handler.generate_video_bytes( - request, - video_id, - reference_image=reference_image, - reference_video=reference_video, - reference_audio=reference_audio, - ) - ) - - save_context = await STORAGE_MANAGER.save(video_bytes, video_id) - logger.info("Video request %s persisted %s output file.", video_id, save_context.key) - - updated_fields = { - "status": VideoGenerationStatus.COMPLETED, - "progress": 100, - "file_name": f"{video_id}.{job.file_extension}", - "completed_at": save_context.created_at, - "inference_time_s": time.perf_counter() - started_at, - "stage_durations": stage_durations, - "peak_memory_mb": peak_memory_mb, - "action": action, - } - updated_fields.update({key: value for key, value in video_metadata.items() if value is not None}) - if save_context.expires_at is not None: - updated_fields["expires_at"] = save_context.expires_at - - await VIDEO_STORE.update_fields(video_id, updated_fields) - except (EngineGenerateError, EngineDeadError) as exc: - logger.exception("Video generation failed (engine error) for id=%s", video_id) - - await _cleanup_video(video_id) - await VIDEO_STORE.update_fields( - video_id, - { - "status": VideoGenerationStatus.FAILED, - "completed_at": int(time.time()), - "error": _video_error_from_exception(exc), - "inference_time_s": time.perf_counter() - started_at, - }, - ) - # Background tasks can't propagate exceptions to FastAPI handlers. - # Actively signal shutdown when the engine is dead. - if app_state is not None and isinstance(exc, EngineDeadError): - terminate_if_errored( - server=app_state.server, - engine=app_state.engine_client, - ) - except Exception as exc: - logger.exception("Video generation failed for id=%s", video_id) - - await _cleanup_video(video_id) - await VIDEO_STORE.update_fields( - video_id, - { - "status": VideoGenerationStatus.FAILED, - "completed_at": int(time.time()), - "error": _video_error_from_exception(exc), - "inference_time_s": time.perf_counter() - started_at, - }, - ) - except asyncio.CancelledError: - await _cleanup_video(video_id) - await VIDEO_STORE.pop(video_id) - raise - finally: - _cleanup_video_references(reference_video, reference_audio, control_path) - - -VIDEO_SYNC_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_VIDEO_SYNC_TIMEOUT", 600.0)) - - -async def _persist_uploaded_video_references(uploads: list[UploadFile]) -> list[str]: - paths: list[str] = [] - try: - for upload in uploads: - suffix = Path(upload.filename or "").suffix.lower() - if suffix not in {".mkv", ".mov", ".mp4", ".webm"}: - suffix = ".mp4" - fd, path = tempfile.mkstemp(prefix="vllm_omni_video_reference_", suffix=suffix) - paths.append(path) - with os.fdopen(fd, "wb") as output: - while chunk := await upload.read(1024 * 1024): - output.write(chunk) - except Exception: - for path in paths: - if os.path.exists(path): - os.unlink(path) - raise - return paths - - -async def _persist_uploaded_control_reference( - upload: UploadFile, - *, - max_bytes: int = CONTROL_REFERENCE_MAX_BYTES, -) -> str: - """Stream one model control upload to request-scoped local storage.""" - kind = _uploaded_media_kind(upload) - suffix = Path(upload.filename or "").suffix.lower() - supported_suffixes = CONTROL_REFERENCE_IMAGE_SUFFIXES | CONTROL_REFERENCE_VIDEO_SUFFIXES - if kind == "audio" or (suffix and suffix not in supported_suffixes): - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail="control_reference must be an image or video file.", - ) - if not suffix: - suffix = ".png" if kind == "image" else ".mp4" - - declared_size = getattr(upload, "size", None) - if isinstance(declared_size, Integral) and int(declared_size) > max_bytes: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail=f"control_reference exceeds the {max_bytes // (1024 * 1024)} MiB size limit.", - ) - - fd, path = tempfile.mkstemp(prefix="vllm_omni_control_reference_", suffix=suffix) - size = 0 - persisted = False - try: - with os.fdopen(fd, "wb") as output: - while chunk := await upload.read(min(1024 * 1024, max_bytes - size + 1)): - size += len(chunk) - if size > max_bytes: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail=f"control_reference exceeds the {max_bytes // (1024 * 1024)} MiB size limit.", - ) - output.write(chunk) - if size == 0: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail="control_reference must not be empty.", - ) - persisted = True - return path - finally: - if not persisted: - with suppress(OSError): - os.unlink(path) - - -def _validate_control_upload( - handler: OmniOpenAIServingVideo, - request: VideoGenerationRequest, - control_reference: UploadFile | None, - control_type: str | None, -) -> str | None: - """Validate a generic uploaded control against model capabilities.""" - if control_reference is None: - if control_type is not None: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail="control_type requires a control_reference upload.", - ) - return None - if control_type is None or not control_type.strip(): - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail="control_reference requires control_type.", - ) - - normalized_type = control_type.strip().lower() - supported_types = frozenset(getattr(handler, "supported_control_upload_types", ())) - if normalized_type not in supported_types: - supported = ", ".join(sorted(supported_types)) or "none" - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail=( - f"Control upload type '{normalized_type}' is not supported by this model. " - f"Supported control upload types: {supported}." - ), - ) - - existing = (request.extra_params or {}).get(normalized_type) - if existing is not None: - if not isinstance(existing, Mapping) and existing is not True: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail=f"extra_params.{normalized_type} must be an object when control_reference is uploaded.", - ) - if isinstance(existing, Mapping) and any(existing.get(key) is not None for key in ("control", "control_path")): - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail=( - "Provide either control_reference or " - f"extra_params.{normalized_type}.control/control_path, not both." - ), - ) - return normalized_type - - -def _attach_control_upload( - request: VideoGenerationRequest, - control_type: str, - control_path: str | None = None, -) -> None: - """Declare an uploaded control and attach its persisted path when available.""" - extra_params = dict(request.extra_params or {}) - existing = extra_params.get(control_type) - control_params = dict(existing) if isinstance(existing, Mapping) else {} - if control_path is not None: - control_params["control_path"] = control_path - extra_params[control_type] = control_params - request.extra_params = extra_params - - -def _reference_list(value: Any) -> list[Any]: - if value is None: - return [] - return list(value) if isinstance(value, list) else [value] - - -def _uploaded_media_kind(upload: UploadFile) -> str: - content_type = (upload.content_type or "").lower() - if content_type.startswith("image/"): - return "image" - if content_type.startswith("audio/"): - return "audio" - suffix = Path(upload.filename or "").suffix.lower() - if suffix in {".png", ".jpg", ".jpeg", ".webp", ".heic", ".heif"}: - return "image" - if suffix in {".wav", ".mp3", ".m4a", ".aac", ".flac", ".ogg"}: - return "audio" - return "video" - - -def _minimax_h3_upload_limit(upload: UploadFile) -> int: - kind = _uploaded_media_kind(upload) - if kind == "image": - return MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES - if kind == "audio": - return MINIMAX_H3_MAX_REFERENCE_AUDIO_BYTES - return MINIMAX_H3_MAX_REFERENCE_VIDEO_BYTES - - -async def _read_upload_limited(upload: UploadFile, *, max_bytes: int | None = None) -> bytes: - """Read an upload with an optional hard byte limit.""" - if max_bytes is None: - return await upload.read() - - declared_size = getattr(upload, "size", None) - if isinstance(declared_size, Integral) and int(declared_size) > max_bytes: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail=f"Uploaded reference exceeds the {max_bytes // (1024 * 1024)} MiB size limit.", - ) - - chunks: list[bytes] = [] - total = 0 - while True: - chunk = await upload.read(min(1024 * 1024, max_bytes - total + 1)) - if not chunk: - break - total += len(chunk) - if total > max_bytes: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail=f"Uploaded reference exceeds the {max_bytes // (1024 * 1024)} MiB size limit.", - ) - chunks.append(chunk) - return b"".join(chunks) - - -def _validate_minimax_h3_image_payload( - payload: bytes, - *, - filename: str | None, - allow_non_image: bool = False, -) -> None: - """Validate a H3 image before converting it to a format-less PIL image.""" - try: - with Image.open(io.BytesIO(payload)) as image: - image_format = str(image.format or "").lower() - if image_format in MINIMAX_H3_REFERENCE_IMAGE_FORMATS: - _validate_image_pixel_limit(image) - except InvalidInputReferenceError as exc: - raise HTTPException(HTTPStatus.BAD_REQUEST.value, detail=str(exc)) from exc - except Image.DecompressionBombError as exc: - raise HTTPException( - HTTPStatus.BAD_REQUEST.value, - detail=f"Invalid uploaded image reference: {filename}; image exceeds the decoder pixel limit.", - ) from exc - except (OSError, ValueError) as exc: - if allow_non_image: - return - raise HTTPException(400, detail=f"Invalid uploaded image reference: {filename}") from exc - - if image_format not in MINIMAX_H3_REFERENCE_IMAGE_FORMATS: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail=( - "MiniMax H3 reference images must use JPG, JPEG, PNG, WEBP, HEIC, or HEIF; " - f"got {image_format or 'unknown'}." - ), - ) - - -async def _persist_uploaded_media_references( - uploads: list[UploadFile], -) -> tuple[list[Image.Image], list[str], list[str]]: - """Persist a mixed MiniMax H3 multipart reference list. - - Images are decoded in memory; videos and audio remain files because H3's - reference encoders need the original container streams (including video - soundtracks). - """ - images: list[Image.Image] = [] - videos: list[str] = [] - audios: list[str] = [] - paths: list[str] = [] - if len(uploads) > MINIMAX_H3_MAX_REFERENCE_COUNT: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail=f"MiniMax H3 accepts at most {MINIMAX_H3_MAX_REFERENCE_COUNT} total references.", - ) - try: - for upload in uploads: - kind = _uploaded_media_kind(upload) - payload = await _read_upload_limited(upload, max_bytes=_minimax_h3_upload_limit(upload)) - if kind == "image": - try: - _validate_minimax_h3_image_payload(payload, filename=upload.filename) - images.append( - _decode_image_bytes( - payload, - source=f"uploaded image reference: {upload.filename}", - ) - ) - except InvalidInputReferenceError as exc: - raise HTTPException(HTTPStatus.BAD_REQUEST.value, detail=str(exc)) from exc - continue - suffix = Path(upload.filename or "").suffix.lower() - if kind == "video" and suffix and suffix not in MINIMAX_H3_REFERENCE_VIDEO_SUFFIXES: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail="MiniMax H3 reference videos must use an MP4 or MOV file.", - ) - if kind == "audio" and suffix and suffix not in MINIMAX_H3_REFERENCE_AUDIO_SUFFIXES: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail="MiniMax H3 reference audio must use a WAV or MP3 file.", - ) - if not suffix or len(suffix) > 8: - suffix = ".mp3" if kind == "audio" else ".mp4" - fd, path = tempfile.mkstemp(prefix="vllm_omni_reference_", suffix=suffix) - paths.append(path) - with os.fdopen(fd, "wb") as output: - output.write(payload) - if kind == "audio": - audios.append(path) - else: - videos.append(path) - except Exception: - for path in paths: - if os.path.exists(path): - os.unlink(path) - raise - return images, videos, audios - - -async def _parse_video_form( - raw_request: Request, - prompt: str = Form(...), - input_reference: UploadFile | None = File(default=None), - input_references: list[UploadFile] | None = File(default=None), - control_reference: UploadFile | None = File(default=None), - control_type: str | None = Form(default=None), - image_reference: str | None = Form(default=None), - video_reference: str | None = Form(default=None), - audio_reference: str | None = Form(default=None), - model: str | None = Form(default=None), - seconds: SecondStr | None = Form(default=None), - size: SizeStr | None = Form(default=None), - user: str | None = Form(default=None), - width: int | None = Form(default=None), - height: int | None = Form(default=None), - num_frames: int | None = Form(default=None), - fps: float | None = Form(default=None, ge=1, allow_inf_nan=False), - aspect_ratio: str | None = Form(default=None), - short_edge: int | None = Form(default=None, ge=1), - num_outputs_per_prompt: int = Form(default=1, ge=1, le=10), - start_time_seconds: float | None = Form(default=None, ge=0.0), - quality: str | None = Form(default=None), - num_inference_steps: int | None = Form(default=None), - guidance_scale: float | None = Form(default=None), - guidance_scale_2: float | None = Form(default=None), - boundary_ratio: float | None = Form(default=None), - flow_shift: float | None = Form(default=None), - true_cfg_scale: float | None = Form(default=None), - seed: int | None = Form(default=None), - generate_sound: bool | None = Form(default=None), - sound_duration: float | None = Form(default=None, gt=0.0), - negative_prompt: str | None = Form(default=None), - enable_frame_interpolation: bool | None = Form(default=None), - frame_interpolation_exp: int | None = Form(default=None, ge=1), - frame_interpolation_scale: float | None = Form(default=None, gt=0.0), - frame_interpolation_model_path: str | None = Form(default=None), - lora: str | None = Form(default=None), - extra_params: str | None = Form(default=None), - return_stage_metrics: bool | None = Form(default=None), -) -> tuple[ - VideoGenerationRequest, - "OmniOpenAIServingVideo", - str, - ReferenceImage | None, - ReferenceVideo | None, - ReferenceAudio | None, - str | None, -]: - """FastAPI dependency that parses video form data, validates inputs, - resolves the handler, and decodes any reference image. - - Used by both ``POST /v1/videos`` (async) and ``POST /v1/videos/sync``. - """ - input_references = input_references or [] - input_reference_bytes: bytes | None = None - parsed_image_reference = _parse_form_json(image_reference) - parsed_video_reference = _parse_form_json(video_reference) - parsed_audio_reference = _parse_form_json(audio_reference) - - if input_references and any( - item is not None for item in (parsed_image_reference, parsed_video_reference, input_reference) - ): - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail="Provide input_references alone, without input_reference, image_reference, or video_reference.", - ) - if input_reference is not None and (parsed_image_reference is not None or parsed_video_reference is not None): - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail=( - "Provide only one of input_reference, image_reference, or video_reference when using " - "input_reference; image_reference and video_reference may be combined." - ), - ) - - request_data: dict[str, Any] = { - "prompt": prompt, - "model": model, - "seconds": seconds, - "size": size, - "image_reference": parsed_image_reference, - "video_reference": parsed_video_reference, - "audio_reference": parsed_audio_reference, - "user": user, - "width": width, - "height": height, - "num_frames": num_frames, - "fps": fps, - "aspect_ratio": aspect_ratio, - "short_edge": short_edge, - "num_outputs_per_prompt": num_outputs_per_prompt, - "start_time_seconds": start_time_seconds, - "quality": quality, - "num_inference_steps": num_inference_steps, - "guidance_scale": guidance_scale, - "guidance_scale_2": guidance_scale_2, - "boundary_ratio": boundary_ratio, - "flow_shift": flow_shift, - "true_cfg_scale": true_cfg_scale, - "seed": seed, - "generate_sound": generate_sound, - "sound_duration": sound_duration, - "negative_prompt": negative_prompt, - "enable_frame_interpolation": enable_frame_interpolation, - "frame_interpolation_exp": frame_interpolation_exp, - "frame_interpolation_scale": frame_interpolation_scale, - "frame_interpolation_model_path": frame_interpolation_model_path, - "lora": _parse_form_json(lora, expected_type=dict), - "extra_params": _parse_form_json(extra_params, expected_type=dict), - "return_stage_metrics": return_stage_metrics, - } - request_data = {k: v for k, v in request_data.items() if v is not None} - request = VideoGenerationRequest(**request_data) - - handler = Omnivideo(raw_request) - if handler is None: - raise HTTPException( - status_code=HTTPStatus.SERVICE_UNAVAILABLE.value, - detail="Video generation handler not initialized.", - ) - logger.info("Video generation handler: %s", type(handler).__name__) - try: - app_model_name, app_stage_configs = _resolve_video_runtime_context(raw_request) - effective_model_name = handler.model_name or app_model_name or request.model or "unknown" - if request.model is not None and effective_model_name is not None and request.model != effective_model_name: - raise HTTPException( - status_code=HTTPStatus.BAD_REQUEST.value, - detail=( - f"Model mismatch: request specifies '{request.model}' but server is running " - f"'{effective_model_name}'." - ), - ) - handler.set_stage_configs_if_missing(app_stage_configs) - except HTTPException: - raise - except Exception as e: - logger.exception("Video generation setup failed: %s", e) - raise HTTPException( - status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, - detail=f"Video generation setup failed: {str(e)}", - ) - - normalized_control_type = _validate_control_upload(handler, request, control_reference, control_type) - if normalized_control_type is not None: - # Make the selected transfer mode visible while choosing the model's - # reference-video decode policy. The upload itself stays unpersisted - # until reference parsing succeeds, preserving failure cleanup. - _attach_control_upload(request, normalized_control_type) - - supports_mixed_reference_inputs = bool(getattr(handler, "supports_mixed_reference_inputs", False)) - if input_reference is not None: - input_reference_bytes = await _read_upload_limited( - input_reference, - max_bytes=_minimax_h3_upload_limit(input_reference) if supports_mixed_reference_inputs else None, - ) - if supports_mixed_reference_inputs: - input_reference_kind = _uploaded_media_kind(input_reference) - _validate_minimax_h3_image_payload( - input_reference_bytes, - filename=input_reference.filename, - allow_non_image=input_reference_kind != "image", - ) - - decode_spec = ReferenceVideoDecodeSpec() - if not input_references and (parsed_video_reference is not None or input_reference_bytes is not None): - stage_configs = ( - handler.stage_configs - or app_stage_configs - or getattr(getattr(handler, "_engine_client", None), "stage_configs", None) - ) - decode_spec = _reference_video_decode_spec(request, stage_configs) - reference_image = None - reference_video = None - reference_audio: ReferenceAudio | None = None - if input_references: - if not supports_mixed_reference_inputs: - video_paths = await _persist_uploaded_video_references(input_references) - reference_video = ReferenceVideo(data=video_paths, cleanup_paths=tuple(video_paths)) - images, audio_paths = [], [] - else: - images, video_paths, audio_paths = await _persist_uploaded_media_references(input_references) - if images: - reference_image = ReferenceImage(data=images if len(images) > 1 else images[0]) - if video_paths: - reference_video = ReferenceVideo(data=video_paths, cleanup_paths=tuple(video_paths)) - if audio_paths: - reference_audio = ReferenceAudio(path=audio_paths, cleanup_paths=tuple(audio_paths)) - else: - video_paths: list[str] = [] - try: - image_items = _reference_list(request.image_reference) - video_items = _reference_list(request.video_reference) - image_data = [] - for item in image_items: - media_data = await decode_input_reference(item, None, None) - if not isinstance(media_data, Image.Image): - raise InvalidInputReferenceError("image_reference did not decode to an image") - image_data.append(media_data) - - video_frames: list[Image.Image] | None = None - for item in video_items: - media_data = await decode_input_reference( - None, - item, - None, - max_video_frames=decode_spec.max_frames, - video_keep=decode_spec.keep, - ) - if not isinstance(media_data, VideoFrames): - raise InvalidInputReferenceError("video_reference did not decode to a video") - if media_data.source_path is not None: - video_paths.append(media_data.source_path) - else: - if len(video_items) != 1: - raise InvalidInputReferenceError( - "multiple video URL references must be downloadable source videos" - ) - video_frames = list(media_data) - - if input_reference_bytes is not None: - media_data = await decode_input_reference( - None, - None, - input_reference_bytes, - max_video_frames=decode_spec.max_frames, - video_keep=decode_spec.keep, - ) - if isinstance(media_data, Image.Image): - image_data.append(media_data) - elif isinstance(media_data, VideoFrames): - if media_data.source_path is not None: - video_paths.append(media_data.source_path) - else: - video_frames = list(media_data) - - if image_data: - reference_image = ReferenceImage(data=image_data if len(image_data) > 1 else image_data[0]) - if video_paths: - reference_video = ReferenceVideo(data=video_paths, cleanup_paths=tuple(video_paths)) - elif video_frames is not None: - reference_video = ReferenceVideo(data=video_frames) - except InvalidInputReferenceError as exc: - for path in video_paths: - if os.path.exists(path): - os.unlink(path) - raise HTTPException(400, detail=str(exc) or "Invalid input reference.") from exc - - audio_paths = [] if reference_audio is None else list(_reference_list(reference_audio.path)) - if request.audio_reference is not None: - try: - for audio_reference in _reference_list(request.audio_reference): - audio_paths.append(await decode_audio_url(audio_reference.audio_url)) - except InvalidInputReferenceError as exc: - _cleanup_video_references(reference_video, reference_audio) - cleanup_paths = set(() if reference_audio is None else reference_audio.cleanup_paths) - for path in audio_paths: - if path not in cleanup_paths and os.path.exists(path): - os.unlink(path) - raise HTTPException(400, detail=str(exc)) from exc - if audio_paths: - cleanup_paths = ( - tuple(audio_paths) - if reference_audio is None - else reference_audio.cleanup_paths - + tuple(path for path in audio_paths if path not in reference_audio.cleanup_paths) - ) - reference_audio = ReferenceAudio( - path=audio_paths if len(audio_paths) > 1 else audio_paths[0], - cleanup_paths=cleanup_paths, - ) - - control_path: str | None = None - if control_reference is not None and normalized_control_type is not None: - try: - control_path = await _persist_uploaded_control_reference( - control_reference, - max_bytes=CONTROL_REFERENCE_MAX_BYTES, - ) - _attach_control_upload(request, normalized_control_type, control_path) - except (asyncio.CancelledError, HTTPException, OSError, TypeError, ValueError): - _cleanup_video_references(reference_video, reference_audio, control_path) - raise - - return ( - request, - handler, - effective_model_name, - reference_image, - reference_video, - reference_audio, - control_path, - ) - - @router.post( "/v1/videos", responses={ @@ -4166,15 +2556,6 @@ async def stop_profile(raw_request: Request, request: ProfileRequest | None = No ) -class OmniSleepRequest(BaseModel): - stage_ids: list[int] - level: int = 2 - - -class OmniWakeupRequest(BaseModel): - stage_ids: list[int] - - @router.post("/v1/omni/sleep") async def omni_sleep(request: OmniSleepRequest, raw_request: Request): engine_client = raw_request.app.state.engine_client diff --git a/vllm_omni/entrypoints/openai/app_state.py b/vllm_omni/entrypoints/openai/app_state.py new file mode 100644 index 00000000000..e56659a4a0b --- /dev/null +++ b/vllm_omni/entrypoints/openai/app_state.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""OpenAI app-state accessors. + +This module owns small helpers that read initialized serving objects or engine +configuration from ``request.app.state``.""" + +from http import HTTPStatus +from typing import Any + +from fastapi import HTTPException, Request +from vllm.engine.protocol import EngineClient + +from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.entrypoints.openai.batch_serving import OmniOpenAIServingChatBatch +from vllm_omni.entrypoints.openai.serving_audio_generate import OmniOpenAIServingAudioGenerate +from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat +from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech +from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo +from vllm_omni.entrypoints.openai.utils import get_stage_type + +ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL = "endpoint-load-metrics-format" + + +async def _get_vllm_config(engine_client: EngineClient) -> Any: + if hasattr(engine_client, "get_vllm_config"): + return await engine_client.get_vllm_config() + return getattr(engine_client, "vllm_config", None) + + +def Omnivideo(request: Request) -> OmniOpenAIServingVideo | None: + return request.app.state.openai_serving_video + + +def Omnichat(request: Request) -> OmniOpenAIServingChat | None: + return request.app.state.openai_serving_chat + + +def OmniBatchChat(request: Request) -> OmniOpenAIServingChatBatch | None: + return request.app.state.openai_serving_chat_batch + + +def Omnispeech(request: Request) -> OmniOpenAIServingSpeech | None: + return request.app.state.openai_serving_speech + + +def OmniAudioGenerate(request: Request) -> OmniOpenAIServingAudioGenerate | None: + return getattr(request.app.state, "openai_serving_audio_generate", None) + + +def _get_engine_and_model(raw_request: Request): + # Get engine client (AsyncOmni) from app state + engine_client: EngineClient | AsyncOmni | None = getattr(raw_request.app.state, "engine_client", None) + if engine_client is None: + raise HTTPException( + status_code=HTTPStatus.SERVICE_UNAVAILABLE.value, + detail="Multi-stage engine not initialized. Start server with a multi-stage omni model.", + ) + + # Check if there's a diffusion stage. + # Prefer app state (compat layer populated at startup), then fall back to + # the engine client's stage configs for refactored AsyncOmni paths. + stage_configs = getattr(raw_request.app.state, "stage_configs", None) + if not stage_configs: + stage_configs = getattr(engine_client, "stage_configs", None) + if not stage_configs: + raise HTTPException( + status_code=HTTPStatus.SERVICE_UNAVAILABLE.value, + detail="Stage configs not found. Start server with a multi-stage omni model.", + ) + + normalized_stage_configs = list(stage_configs) + has_diffusion_stage = any(get_stage_type(stage_cfg) == "diffusion" for stage_cfg in normalized_stage_configs) + + if not has_diffusion_stage: + raise HTTPException( + status_code=HTTPStatus.SERVICE_UNAVAILABLE.value, + detail="No diffusion stage found in multi-stage pipeline.", + ) + + # Get server's loaded model name + serving_models = getattr(raw_request.app.state, "openai_serving_models", None) + base_model_paths = getattr(serving_models, "base_model_paths", None) if serving_models else None + if base_model_paths: + model_name = base_model_paths[0].name + else: + model_name = "unknown" + + return engine_client, model_name, normalized_stage_configs + + +def _get_diffusion_od_config(raw_request: Request, engine_client: Any) -> Any: + diffusion_engine = getattr(raw_request.app.state, "diffusion_engine", None) or engine_client + get_diffusion_od_config = getattr(diffusion_engine, "get_diffusion_od_config", None) + return ( + get_diffusion_od_config() if callable(get_diffusion_od_config) else getattr(diffusion_engine, "od_config", None) + ) diff --git a/vllm_omni/entrypoints/openai/audio_utils_mixin.py b/vllm_omni/entrypoints/openai/audio_utils_mixin.py index e3bce1a67ff..4382afc57b1 100644 --- a/vllm_omni/entrypoints/openai/audio_utils_mixin.py +++ b/vllm_omni/entrypoints/openai/audio_utils_mixin.py @@ -1,5 +1,25 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Audio serving utility mixin. + +PUT HERE: + - AudioMixin helpers that convert audio tensors/bytes/formats for speech + and audio serving classes (shared by multiple serving paths). + +DO NOT PUT HERE: + - FastAPI route / form / job orchestration peeled from ``api_server.py``. + Put those under an audio package ``helpers.py`` (or audio serving modules) + when extracted. + +LONGEVITY: + - This root mixin is a **temporary shared home**. + - TODO(#5227, P1.1): tidy up / move with the audio/speech family split + in the Phase 1 audio PR; do not treat this file as the long-term owner. + - Endpoint-family helpers under an audio package are the longer home for + route-adjacent logic. + +See ``openai/README.md`` (utils vs helpers, no overlap). +""" import math from io import BytesIO diff --git a/vllm_omni/entrypoints/openai/chat_template.py b/vllm_omni/entrypoints/openai/chat_template.py new file mode 100644 index 00000000000..4cdbcb156ed --- /dev/null +++ b/vllm_omni/entrypoints/openai/chat_template.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Chat-template loading helpers for OpenAI server bootstrap.""" + +import json +from pathlib import Path + +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +def _load_model_chat_template_json(model: str) -> str | None: + """Load a model-level chat_template.json from a local path or HF cache. + + Some multimodal HF repos, including Qwen3-Omni, ship the chat template as a + separate file instead of embedding it in tokenizer_config.json. Transformers + 4.44+ no longer supplies a default template, so serving must pass that model + template explicitly when the user did not provide --chat-template. + """ + candidate = Path(model) / "chat_template.json" + template_path: str | None = str(candidate) if candidate.is_file() else None + + if template_path is None: + try: + from vllm_omni.transformers_utils.repo_utils import hf_api + + template_path = hf_api().hf_hub_download( + repo_id=model, + filename="chat_template.json", + local_files_only=True, + ) + except Exception: + return None + + try: + with open(template_path, encoding="utf-8") as f: + payload = json.load(f) + except Exception as exc: + logger.warning("Failed to load chat template from %s: %s", template_path, exc) + return None + + if isinstance(payload, dict): + template = payload.get("chat_template") + elif isinstance(payload, str): + template = payload + else: + template = None + + if not isinstance(template, str) or not template.strip(): + logger.warning("Ignoring malformed chat template payload in %s", template_path) + return None + + logger.info("Loaded chat template from %s", template_path) + return template diff --git a/vllm_omni/entrypoints/openai/diffusion.py b/vllm_omni/entrypoints/openai/diffusion.py new file mode 100644 index 00000000000..be00908a5b1 --- /dev/null +++ b/vllm_omni/entrypoints/openai/diffusion.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""OpenAI diffusion-stage request helpers. + +These helpers are stage-based rather than image-format-specific.""" + +import json +from http import HTTPStatus +from typing import Any, cast + +from fastapi import HTTPException + +from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.entrypoints.openai.stage_params import ( + build_stage_sampling_params_list, + get_default_sampling_params_list, +) + +MAX_UINT32_SEED = 2**32 - 1 + + +async def _generate_with_async_omni( + engine_client: AsyncOmni | Any, + gen_params: Any, + stage_configs: list[Any], + **kwargs, +): + engine_client = cast(AsyncOmni, engine_client) + result = None + normalized_stage_configs = list(stage_configs) + if not normalized_stage_configs: + raise HTTPException( + status_code=HTTPStatus.SERVICE_UNAVAILABLE.value, + detail="Stage configs not found. Start server with a multi-stage omni model.", + ) + sampling_params_list = build_stage_sampling_params_list( + normalized_stage_configs, + get_default_sampling_params_list(engine_client), + diffusion_params=gen_params, + replace_diffusion_params=True, + ) + + async for output in engine_client.generate( + sampling_params_list=sampling_params_list, + **kwargs, + ): + result = output + + if result is None: + raise HTTPException( + status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, + detail="No output generated from multi-stage pipeline.", + ) + return result + + +def apply_stage_default_sampling_params( + default_params_json: str | None, + sampling_params: Any, + stage_key: str, +) -> None: + """ + Update a stage's sampling parameters with vLLM-Omni defaults. + + Args: + default_params_json: JSON string of stage-keyed default parameters + sampling_params: The sampling parameters object to update + stage_key: The stage ID/key in the pipeline + """ + if default_params_json is not None: + default_params_dict = json.loads(default_params_json) + if stage_key in default_params_dict: + stage_defaults = default_params_dict[stage_key] + for param_name, param_value in stage_defaults.items(): + if hasattr(sampling_params, param_name): + setattr(sampling_params, param_name, param_value) diff --git a/vllm_omni/entrypoints/openai/errors.py b/vllm_omni/entrypoints/openai/errors.py index fda9b64361a..efa43bff1b7 100644 --- a/vllm_omni/entrypoints/openai/errors.py +++ b/vllm_omni/entrypoints/openai/errors.py @@ -1,3 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""OpenAI endpoint error helpers. + +This module owns helpers that are specific to OpenAI-compatible response +shapes. Server-level exception handling belongs in +``vllm_omni.entrypoints.serve.utils.errors``.""" + +from http import HTTPStatus + +from fastapi import Request +from fastapi.responses import JSONResponse +from vllm.entrypoints.serve.engine.protocol import ErrorResponse +from vllm.entrypoints.serve.instrumentator.basic import base + + class InvalidInputReferenceError(ValueError): def __init__(self, message: str = "Invalid input reference.") -> None: super().__init__(message) + + +def _error_response_to_json_response( + err: ErrorResponse, + *, + status_code: HTTPStatus | int | None = None, + default_status_code: HTTPStatus | int = HTTPStatus.BAD_REQUEST, +) -> JSONResponse: + resolved_status = int( + status_code + if status_code is not None + else (err.error.code if err.error and err.error.code is not None else default_status_code) + ) + payload = err.model_dump() + if err.error: + payload["error"]["code"] = resolved_status + return JSONResponse(content=payload, status_code=resolved_status) + + +# TODO(#5227, P1.1): Relocate to audio/speech helpers when speech routes move out of api_server.py. +def _create_speech_error_json_response( + raw_request: Request, + message: str, + *, + err_type: str = "BadRequestError", + status_code: HTTPStatus = HTTPStatus.BAD_REQUEST, +) -> JSONResponse: + err = base(raw_request).create_error_response( + message=message, + err_type=err_type, + status_code=status_code, + ) + return _error_response_to_json_response(err, status_code=status_code) diff --git a/vllm_omni/entrypoints/openai/image_api_utils.py b/vllm_omni/entrypoints/openai/image_api_utils.py index ed759a08017..f5654b335ec 100644 --- a/vllm_omni/entrypoints/openai/image_api_utils.py +++ b/vllm_omni/entrypoints/openai/image_api_utils.py @@ -1,11 +1,24 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ -Shared helper utilities for OpenAI-compatible image generation API. +Shared media utilities for OpenAI-compatible image APIs. -This module provides common helper functions for the image generation endpoint. -All functions work with plain Python types to maintain separation from the -FastAPI HTTP layer. +PUT HERE: + - Pure media helpers with no FastAPI Request / app-state / job semantics: + size parsing, image encode/base64, layered-layer validation. + +DO NOT PUT HERE: + - ``/v1/images*`` request/engine/job helpers — those go in ``images.helpers``. + +LONGEVITY: + - This root utils file is a **temporary shared home**. + - TODO(#5227, P1.2): tidy up / absorb into the images family (e.g. ``images/media.py`` + or equivalent) in the P1.2 image modality PR; do not treat this file as the + long-term owner. + - ``images.helpers`` is the longer home for endpoint-specific logic through + P0.2/P0.3 until P1.2 further splits it. + +See ``openai/README.md`` and ``images/README.md`` (utils vs helpers, no overlap). """ import base64 diff --git a/vllm_omni/entrypoints/openai/images/README.md b/vllm_omni/entrypoints/openai/images/README.md new file mode 100644 index 00000000000..e7f32753b96 --- /dev/null +++ b/vllm_omni/entrypoints/openai/images/README.md @@ -0,0 +1,31 @@ +# Image Endpoints + +This package owns OpenAI-compatible image generation and image edit endpoint +logic. + +## `image_api_utils.py` vs `images/helpers.py` (no overlap) + +| | `openai/image_api_utils.py` | `openai/images/helpers.py` | +| --- | --- | --- | +| Layer | Shared media utils | Image endpoint helpers | +| Longevity | **Temporary.** TODO(#5227, P1.2): absorb/tidy into this package (e.g. `media.py`) | **Longer home** for `/v1/images*` logic through P0.2/P0.3; P1.2 may split further | +| PUT HERE | `parse_size`, encode/base64, layered-layer validation — plain Python, no FastAPI `Request` | Engine/app-state limits, input image loading from the edit request, Hunyuan edit extras, result extraction/normalization | +| DO NOT PUT HERE | Request/engine/job helpers | Shared encode/`parse_size`/layered validation | +| Typical importers | `protocol/images.py`, `serving_chat.py`, image routes | `api_server.py` image routes (later `images/api_router.py`) | +| May import the other? | **No** — must not import `images.helpers` | **Yes** — may call `image_api_utils` | + +**Prefer growing `helpers.py`.** Do not grow root `image_api_utils.py` unless the symbol is clearly shared media with no endpoint semantics. + +## Put Here + +- Request parsing, validation, and response helpers used only by image + endpoints (`helpers.py`). +- Image endpoint route bodies after route extraction. +- Image-specific bridge/adapters introduced during the image refactor. + +## Do Not Put Here + +- Generic server utilities. +- App-state accessors shared by multiple endpoint families. +- Shared encode/size/layered validation already in `image_api_utils.py`. +- Video or audio endpoint behavior. diff --git a/vllm_omni/entrypoints/openai/images/__init__.py b/vllm_omni/entrypoints/openai/images/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm_omni/entrypoints/openai/images/helpers.py b/vllm_omni/entrypoints/openai/images/helpers.py new file mode 100644 index 00000000000..dd790ca3d19 --- /dev/null +++ b/vllm_omni/entrypoints/openai/images/helpers.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Image endpoint helpers peeled from ``api_server.py``. + +PUT HERE: + - ``/v1/images*`` request/engine helpers: limits, input loading, edit extras, + result extraction/normalization, and other route-adjacent logic that used + to live in ``api_server.py``. + - May import ``image_api_utils`` for encode/validate (helpers → utils OK). + +DO NOT PUT HERE: + - Shared encode/size/layered validation — keep those in ``image_api_utils``. + +LONGEVITY: + - This package ``helpers.py`` is the **longer home** for image endpoint + logic through the helper (P0.2) + router (P0.3) stages. + - The later P1.2 image modality PR (#5227) may further split this file (e.g. + validation / editing / responses) and absorb ``image_api_utils`` into + this package; until then keep growing here, not in root utils. + +See ``images/README.md`` (utils vs helpers, no overlap). +""" + +import io +from http import HTTPStatus +from numbers import Integral +from typing import Any + +import httpx +import numpy as np +import pybase64 as base64 +from fastapi import HTTPException, Request +from PIL import Image +from vllm.logger import init_logger + +from vllm_omni.entrypoints.openai.app_state import _get_diffusion_od_config + +logger = init_logger(__name__) + + +def _get_max_edit_input_images(raw_request: Request, engine_client: Any) -> int | None: + od_config = _get_diffusion_od_config(raw_request, engine_client) + if od_config is None: + # Preserve the existing compatibility behavior when the diffusion + # config is not exposed on the serving surface. + return None + + supports_multimodal_inputs = getattr(od_config, "supports_multimodal_inputs", None) + if not isinstance(supports_multimodal_inputs, bool): + # Older serving surfaces and mocked engines may expose a placeholder + # object instead of a real diffusion config. Treat that as "unknown" + # so existing single-image flows keep working. + return None + + if not supports_multimodal_inputs: + return 1 + + max_input_images = getattr(od_config, "max_multimodal_image_inputs", None) + if max_input_images is None: + return None + if isinstance(max_input_images, bool) or not isinstance(max_input_images, Integral): + return None + if max_input_images < 1: + return None + return int(max_input_images) + + +def _check_max_generated_image_size( + app_state_args: Any, + width: int | None, + height: int | None, + resolution: int | None = None, +) -> None: + """Raise 400 if the requested image size exceeds --max-generated-image-size.""" + max_generated_image_size = getattr(app_state_args, "max_generated_image_size", None) + # Check max_generated_image_size + if max_generated_image_size is None: + return + if width is not None and height is not None: + if width * height > max_generated_image_size: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"Requested image size {width}x{height} exceeds the maximum allowed " + f"size of {max_generated_image_size} pixels. You can reduce the requested size " + f"or increase the server's --max-generated-image-size limit.", + ) + elif resolution is not None: + # When resolution is set, the output size is resolution * resolution + if resolution * resolution > max_generated_image_size: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"Requested resolution {resolution} (max {resolution}x{resolution} pixels) " + f"exceeds the maximum allowed size of {max_generated_image_size} pixels. " + f"You can reduce the requested size or increase the server's --max-generated-image-size limit.", + ) + + +def _build_hunyuan_edit_extra_args( + *, + bot_task: str | None, + sys_type: str | None, + system_prompt: str | None, +) -> dict[str, Any]: + """Map Hunyuan /v1/images/edits form fields to DiT ``extra_args``.""" + extra_args: dict[str, Any] = {} + effective_use_system_prompt = sys_type + if effective_use_system_prompt is None and bot_task is not None: + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import resolve_sys_type + + effective_use_system_prompt = resolve_sys_type(bot_task) + if effective_use_system_prompt is not None: + extra_args["use_system_prompt"] = effective_use_system_prompt + if system_prompt is not None: + extra_args["system_prompt"] = system_prompt + if bot_task is not None: + extra_args["bot_task"] = bot_task + return extra_args + + +def _update_if_not_none(object: Any, key: str, val: Any) -> None: + if val is not None: + setattr(object, key, val) + + +def _normalize_image(image: Any) -> Any: + """Normalize a single image output to a PIL-compatible format.""" + if isinstance(image, Image.Image): + return image + if not isinstance(image, np.ndarray): + raise ValueError(f"Unsupported image type: {type(image)}") + if not np.issubdtype(image.dtype, np.integer) and not np.issubdtype(image.dtype, np.floating): + raise ValueError(f"Unsupported dtype: {image.dtype}") + if isinstance(image, np.ndarray): + while image.ndim > 3: + image = image[0] + if image.min() < 0: + if image.min() < -1.01 or image.max() > 1.01: + logger.warning( + f"Image float range [{image.min():.2f}, {image.max():.2f}] outside expected [-1, 1]. " + f"Clipping to [-1, 1] before normalization." + ) + image = np.clip(image, -1.0, 1.0) * 0.5 + 0.5 + elif image.max() > 1.01: + logger.warning( + f"Image float range [{image.min():.2f}, {image.max():.2f}] outside expected [0, 1]. " + f"Clipping to [0, 1] before normalization." + ) + image = (np.clip(image, 0.0, 1.0) * 255).astype(np.uint8) + image = Image.fromarray(image) + return image + + +def _extract_images_from_result(result: Any) -> list[Any]: + images = [] + if hasattr(result, "images") and result.images: + images = result.images + # Handle when generate more than one image + if images and isinstance(images[0], np.ndarray) and images[0].shape[0] > 1 and images[0].ndim == 5: + # Unwrap batch: (N, T, H, W, C) -> [img1, img2, ...] + images = list(images[0]) + # Flatten nested lists (e.g., from layered models like Qwen-Image-Layered). + # Note: This only flattens one level deep. Deeper nesting is not supported. + flattened = [] + for img in images: + if isinstance(img, list): + flattened.extend(img) + else: + flattened.append(img) + return [_normalize_image(img) for img in flattened] + + +async def _load_input_images( + inputs: list[str], + *, + normalize_rgb: bool = True, +) -> list[Image.Image]: + """ + convert to PIL.Image.Image list + """ + if isinstance(inputs, str): + inputs = [inputs] + + images: list[Image.Image] = [] + + for inp in inputs: + # 1. URL + base64 + if isinstance(inp, str) and inp.startswith("data:image"): + try: + _, b64_data = inp.split(",", 1) + image_bytes = base64.b64decode(b64_data) + img = Image.open(io.BytesIO(image_bytes)) + images.append(img) + except Exception as e: + raise ValueError(f"Invalid base64 image: {e}") + + # 2. URL + elif isinstance(inp, str) and inp.startswith("http"): + async with httpx.AsyncClient(timeout=60) as client: + try: + resp = await client.get(inp) + resp.raise_for_status() + img = Image.open(io.BytesIO(resp.content)) + images.append(img) + except Exception as e: + raise ValueError(f"Failed to download image from URL {inp}: {e}") + + # 3. UploadFile + elif hasattr(inp, "file"): + try: + img_data = await inp.read() + img = Image.open(io.BytesIO(img_data)) + images.append(img) + except Exception as e: + raise ValueError(f"Failed to open uploaded file: {e}") + + else: + raise ValueError(f"Unsupported input: {inp}") + + if not images: + raise ValueError("No valid input images found") + + if not normalize_rgb: + return images + + # Match the offline HunyuanImage3 image-edit example path, which eagerly + # normalizes input files with ``Image.open(...).convert("RGB")`` before + # they reach the AR stage. Keeping uploads as RGBA/P PIL objects makes + # online IT2I observe a different visual input than offline (for example + # transparent-logo uploads alpha-composited over white instead of black), + # which is enough for HunyuanImage3 AR recaption to diverge before DiT + # sees the request -- root cause of the "online 3 magnets vs offline 1 + # magnet" systematic semantic mismatch. + return [img.convert("RGB") for img in images] + + +def _choose_output_format(output_format: str | None, background: str | None) -> str: + # Normalize and choose extension + fmt = (output_format or "").lower() + if fmt in {"jpg", "png", "webp", "jpeg"}: + return fmt + # If transparency requested, prefer png + if (background or "auto").lower() == "transparent": + return "png" + # Default + return "jpeg" diff --git a/vllm_omni/entrypoints/openai/lora.py b/vllm_omni/entrypoints/openai/lora.py new file mode 100644 index 00000000000..b44121aee88 --- /dev/null +++ b/vllm_omni/entrypoints/openai/lora.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""OpenAI LoRA request helpers. + +Use this module for HTTP-facing LoRA parsing shared by OpenAI-compatible +endpoint families.""" + +import json +from http import HTTPStatus +from typing import Any + +from fastapi import HTTPException + +from vllm_omni.entrypoints.openai.utils import parse_lora_request + + +def _get_lora_from_json_str(lora_body): + if lora_body is None: + return None + try: + lora_dict = json.loads(lora_body) + except json.JSONDecodeError: + raise HTTPException(status_code=400, detail="Invalid LoRA JSON string") + + if not isinstance(lora_dict, dict): + raise HTTPException(status_code=400, detail="LoRA must be a JSON object") + + return lora_dict + + +def _parse_lora_request(lora_body: dict[str, Any]): + try: + return parse_lora_request(lora_body) + except ValueError as e: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=str(e), + ) from e diff --git a/vllm_omni/entrypoints/openai/models/README.md b/vllm_omni/entrypoints/openai/models/README.md new file mode 100644 index 00000000000..42974938042 --- /dev/null +++ b/vllm_omni/entrypoints/openai/models/README.md @@ -0,0 +1,19 @@ +# OpenAI Models + +This package mirrors upstream OpenAI model-list ownership. + +## Put Here + +- Helpers and serving shims that back `/v1/models`. +- Route bodies for `/v1/models` after endpoint extraction. +- Model-list protocol adapters that are OpenAI-specific. + +## Do Not Put Here + +- Actual model execution logic. +- Image/video/audio endpoint helpers. +- Generic app-state helpers. + +`_DiffusionServingModels` lives here because pure diffusion servers still need +to satisfy the OpenAI `/v1/models` contract. The route body remains in +`openai/api_server.py` until the route extraction tracked by #5227. diff --git a/vllm_omni/entrypoints/openai/models/__init__.py b/vllm_omni/entrypoints/openai/models/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm_omni/entrypoints/openai/models/serving.py b/vllm_omni/entrypoints/openai/models/serving.py new file mode 100644 index 00000000000..9917098f8a4 --- /dev/null +++ b/vllm_omni/entrypoints/openai/models/serving.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""OpenAI models serving helpers. + +This module mirrors upstream's OpenAI models area for helpers that back +``/v1/models`` behavior.""" + +from vllm.entrypoints.openai.models.protocol import BaseModelPath +from vllm.entrypoints.serve.engine.protocol import ModelCard, ModelList, ModelPermission + + +class _DiffusionServingModels: + """Minimal OpenAIServingModels implementation for diffusion-only servers. + + vLLM's /v1/models route expects `app.state.openai_serving_models` to expose + `show_available_models()`. In pure diffusion mode we don't initialize the + full OpenAIServingModels (it depends on LLM-specific processors), so we + provide a lightweight fallback. + """ + + class _NullModelConfig: + def __getattr__(self, name): + return None + + class _Unsupported: + def __init__(self, name: str): + self.name = name + + def __call__(self, *args, **kwargs): + raise NotImplementedError(f"{self.name} is not supported in diffusion mode") + + def __getattr__(self, attr): + raise NotImplementedError(f"{self.name}.{attr} is not supported in diffusion mode") + + def __init__(self, base_model_paths: list[BaseModelPath]) -> None: + self._base_model_paths = base_model_paths + self.model_config = self._NullModelConfig() + + @property + def base_model_paths(self) -> list[BaseModelPath]: + return self._base_model_paths + + def is_base_model(self, model_name: str) -> bool: + return any(p.name == model_name for p in self._base_model_paths) + + def __getattr__(self, name): + return self._Unsupported(name) + + async def show_available_models(self) -> ModelList: + return ModelList( + data=[ + ModelCard( + id=base_model.name, + root=base_model.model_path, + permission=[ModelPermission()], + ) + for base_model in self._base_model_paths + ] + ) diff --git a/vllm_omni/entrypoints/openai/utils.py b/vllm_omni/entrypoints/openai/utils.py index 4aac073d4b2..9ef01003dec 100644 --- a/vllm_omni/entrypoints/openai/utils.py +++ b/vllm_omni/entrypoints/openai/utils.py @@ -1,4 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Shared OpenAI serving utilities. + +PUT HERE: + - Cross-cutting stage-config / LoRA / pipeline helpers reused by multiple + OpenAI serving paths (not owned by a single endpoint family). + +DO NOT PUT HERE: + - Modality endpoint request/job helpers — those go in package ``helpers.py`` + (e.g. ``images.helpers``, ``video.generation.helpers``). + - Pure image/video/audio media encode/decode — those stay in the matching + root ``*_utils*`` / mixin until family PRs absorb them. + +LONGEVITY: + - TODO(#5227, Phase 1): revisit and tidy this file during later OpenAI + entrypoint refactor stages; prefer growing endpoint-family helpers over this + root grab-bag when ownership is clear. +""" from __future__ import annotations diff --git a/vllm_omni/entrypoints/openai/video/README.md b/vllm_omni/entrypoints/openai/video/README.md new file mode 100644 index 00000000000..dd73b61dcb5 --- /dev/null +++ b/vllm_omni/entrypoints/openai/video/README.md @@ -0,0 +1,33 @@ +# Video Endpoints + +This package owns OpenAI-compatible video endpoint families. + +## `video_api_utils.py` vs `video/generation/helpers.py` (no overlap) + +| | `openai/video_api_utils.py` | `openai/video/generation/helpers.py` | +| --- | --- | --- | +| Layer | Shared video media utils | `/v1/videos*` endpoint helpers | +| Longevity | **Temporary.** TODO(#5227, P1.3): move/tidy into family `media.py` | **Longer home** for generation resource logic through P0.2/P0.3; P1.3 may split further | +| PUT HERE | Decode/encode URLs and bytes, frame/audio coercion, streaming encoders | Multipart form parsing, MiniMax upload limits, job runner/cleanup, `VideoResponse` factory, runtime context from app state | +| DO NOT PUT HERE | Form/job/store orchestration | Reimplemented encode/decode | +| Typical importers | `serving_video.py`, streaming output, generation helpers | `api_server.py` video routes (later `generation/api_router.py`) | +| May import the other? | **No** — must not import `video.generation.helpers` | **Yes** — may call `video_api_utils` | + +**Prefer growing `generation/helpers.py`.** Do not grow root `video_api_utils.py` unless the symbol is clearly shared media. + +## Put Here + +- Video generation/resource endpoint code under `generation/`. +- Generated video streaming-output code under a dedicated subpackage when that + route moves. +- Video-specific request parsing, response formatting, and storage/job helpers + for the owning subpackage. + +## Do Not Put Here + +- Live streaming-input chat sessions such as `/v1/video/chat/stream`; those + should get their own route-owner package in the router split rather than + living under video generation. +- Generic server utilities. +- Shared media encode/decode that already lives in root `video_api_utils.py`. +- Image or audio endpoint behavior. diff --git a/vllm_omni/entrypoints/openai/video/__init__.py b/vllm_omni/entrypoints/openai/video/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm_omni/entrypoints/openai/video/generation/README.md b/vllm_omni/entrypoints/openai/video/generation/README.md new file mode 100644 index 00000000000..476db23c76e --- /dev/null +++ b/vllm_omni/entrypoints/openai/video/generation/README.md @@ -0,0 +1,33 @@ +# Video Generation + +This package owns the `/v1/videos` resource family. + +## Helpers vs root video utils + +See also `../README.md` and `../../README.md`. + +| | PUT HERE | DO NOT PUT HERE | Longevity | +| --- | --- | --- | --- | +| `helpers.py` | Form parse, upload validation, job run/cleanup, response factories, app-state runtime context for `/v1/videos*` | Shared encode/decode backends | **Longer home** through P0.2/P0.3; later P1.3 PR may split (form/jobs) | +| `openai/video_api_utils.py` | Shared encode/decode/media backends only | Form/job/store orchestration | **Temporary.** TODO(#5227, P1.3): move into e.g. `media.py` | + +- Helpers call utils; utils never grow endpoint/job logic. +- Prefer growing `helpers.py` over root utils. + +## Put Here + +- Multipart form parsing for video generation requests. +- Async/sync video generation job helpers. +- Video resource status, delete, retrieve, and download helpers. +- The `api_router.py` for `/v1/videos*` once route bodies move out of + `api_server.py`. + +## Do Not Put Here + +- Live video-chat streaming input. +- Generated video streaming-output sessions. +- Generic video encoding/decoding helpers used outside this resource family + (keep those in `video_api_utils.py` until the P1.3 video family PR moves them). + +The P0.2 helper migration places non-route video generation code here first. Full +route extraction (P0.3) and `video_api_utils.py` cleanup (P1.3) are tracked by #5227. diff --git a/vllm_omni/entrypoints/openai/video/generation/__init__.py b/vllm_omni/entrypoints/openai/video/generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm_omni/entrypoints/openai/video/generation/helpers.py b/vllm_omni/entrypoints/openai/video/generation/helpers.py new file mode 100644 index 00000000000..e23d876e5c1 --- /dev/null +++ b/vllm_omni/entrypoints/openai/video/generation/helpers.py @@ -0,0 +1,986 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""``/v1/videos*`` endpoint helpers peeled from ``api_server.py``. + +PUT HERE: + - Multipart form parsing, upload limits, job run/cleanup, response factories, + app-state runtime context, and other route-adjacent ``/v1/videos*`` logic + that used to live in ``api_server.py``. + - May import ``video_api_utils`` for decode/encode (helpers → utils OK). + +DO NOT PUT HERE: + - Shared media encode/decode backends — keep those in ``video_api_utils``. + +LONGEVITY: + - This package ``helpers.py`` is the **longer home** for video generation + endpoint logic through the helper (P0.2) + router (P0.3) stages. + - The later P1.3 video modality PR (#5227) may further split this file (e.g. form / + jobs) and move ``video_api_utils`` into ``media.py`` here; until then keep + growing here, not in root utils. + +See ``video/generation/README.md`` (utils vs helpers, no overlap). +""" + +import asyncio +import io +import json +import os +import tempfile +import time +from collections.abc import Mapping, Sequence +from contextlib import suppress +from http import HTTPStatus +from numbers import Integral +from pathlib import Path +from typing import Any, Literal, cast + +from fastapi import File, Form, HTTPException, Request, UploadFile +from PIL import Image +from vllm.entrypoints.launchers.launcher import terminate_if_errored +from vllm.entrypoints.serve import create_error_response +from vllm.logger import init_logger +from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError + +from vllm_omni.diffusion.models.interface import ReferenceVideoDecodeSpec +from vllm_omni.entrypoints.openai.app_state import Omnivideo +from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError +from vllm_omni.entrypoints.openai.protocol.videos import ( + SecondStr, + SizeStr, + VideoAction, + VideoError, + VideoGenerationRequest, + VideoGenerationStatus, + VideoResponse, +) +from vllm_omni.entrypoints.openai.serving_video import ( + OmniOpenAIServingVideo, + ReferenceAudio, + ReferenceImage, + ReferenceVideo, +) +from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER +from vllm_omni.entrypoints.openai.stores import VIDEO_STORE +from vllm_omni.entrypoints.openai.utils import get_stage_type +from vllm_omni.entrypoints.openai.video_api_utils import ( + VideoFrames, + _decode_image_bytes, + _validate_image_pixel_limit, + decode_audio_url, + decode_input_reference, +) +from vllm_omni.errors import OmniClientError + +logger = init_logger(__name__) + + +MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES = 30 * 1024 * 1024 + +MINIMAX_H3_MAX_REFERENCE_VIDEO_BYTES = 50 * 1024 * 1024 + +MINIMAX_H3_MAX_REFERENCE_AUDIO_BYTES = 15 * 1024 * 1024 + +MINIMAX_H3_MAX_REFERENCE_COUNT = 12 + +MINIMAX_H3_REFERENCE_IMAGE_FORMATS = frozenset({"jpeg", "png", "webp", "heic", "heif"}) + +MINIMAX_H3_REFERENCE_VIDEO_SUFFIXES = frozenset({".mp4", ".mov"}) + +MINIMAX_H3_REFERENCE_AUDIO_SUFFIXES = frozenset({".wav", ".mp3"}) + +CONTROL_REFERENCE_IMAGE_SUFFIXES = frozenset({".bmp", ".gif", ".jpg", ".jpeg", ".png", ".tif", ".tiff", ".webp"}) +CONTROL_REFERENCE_VIDEO_SUFFIXES = frozenset({".mkv", ".mov", ".mp4", ".webm"}) +CONTROL_REFERENCE_MAX_BYTES = 512 * 1024 * 1024 + +VIDEO_SYNC_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_VIDEO_SYNC_TIMEOUT", 600.0)) + + +def _resolve_video_runtime_context(raw_request: Request) -> tuple[str | None, list[Any] | None]: + app_model_name = None + serving_models = getattr(raw_request.app.state, "openai_serving_models", None) + if serving_models and getattr(serving_models, "base_model_paths", None): + base_paths = serving_models.base_model_paths + if base_paths: + app_model_name = base_paths[0].name + + app_stage_configs = getattr(raw_request.app.state, "stage_configs", None) + return app_model_name, app_stage_configs + + +def _parse_form_json(value: str | None, expected_type: type | None = None) -> Any: + if value is None or value == "": + return None + try: + parsed = json.loads(value) + except json.JSONDecodeError as exc: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="Invalid JSON in form field.", + ) from exc + if expected_type is not None and not isinstance(parsed, expected_type): + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"Invalid JSON in form field: expected {expected_type.__name__}, got {type(parsed).__name__}.", + ) + return parsed + + +def _config_get(config: Any, key: str, default: Any = None) -> Any: + if isinstance(config, dict): + return config.get(key, default) + if hasattr(config, "get"): + try: + return config.get(key, default) + except Exception: + pass + return getattr(config, key, default) + + +def _stage_engine_args(stage_cfg: Any) -> Any: + return _config_get(stage_cfg, "engine_args", {}) or {} + + +def _diffusion_model_classes(stage_configs: list[Any] | None) -> list[type]: + if not stage_configs: + return [] + + from vllm_omni.diffusion.registry import DiffusionModelRegistry + + model_classes: list[type] = [] + for stage_cfg in stage_configs: + if get_stage_type(stage_cfg) != "diffusion": + continue + model_class_name = _config_get(_stage_engine_args(stage_cfg), "model_class_name") + if not model_class_name: + continue + model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name) + if model_cls is not None: + model_classes.append(model_cls) + return model_classes + + +def _normalize_reference_video_decode_spec(spec: ReferenceVideoDecodeSpec) -> ReferenceVideoDecodeSpec: + max_frames = spec.max_frames + if max_frames is not None: + try: + max_frames = int(max_frames) + except (TypeError, ValueError) as exc: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="Invalid reference video decode spec: max_frames must be an integer.", + ) from exc + if max_frames <= 0: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="Invalid reference video decode spec: max_frames must be positive.", + ) + + keep = str(spec.keep or "first").strip().lower() + if keep not in {"first", "last"}: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="Invalid reference video decode spec: keep must be either 'first' or 'last'.", + ) + return ReferenceVideoDecodeSpec(max_frames=max_frames, keep=cast(Literal["first", "last"], keep)) + + +def _reference_video_decode_spec( + req: VideoGenerationRequest, + stage_configs: list[Any] | None, +) -> ReferenceVideoDecodeSpec: + video_params = req.resolve_video_params() + extra_params = req.extra_params if isinstance(req.extra_params, dict) else {} + for model_cls in _diffusion_model_classes(stage_configs): + resolver = getattr(model_cls, "reference_video_decode_spec", None) + if resolver is None: + continue + try: + spec = resolver(num_frames=video_params.num_frames, extra_args=extra_params) + except (TypeError, ValueError) as exc: + raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(exc)) from exc + if spec is not None: + return _normalize_reference_video_decode_spec(spec) + return ReferenceVideoDecodeSpec(max_frames=video_params.num_frames, keep="first") + + +def video_response_from_request(model_name: str, req: VideoGenerationRequest) -> VideoResponse: + video_params = req.resolve_video_params() + duration_s = None + if req.seconds is not None: + duration_s = float(req.seconds) + elif video_params.num_frames is not None and video_params.fps is not None: + duration_s = video_params.num_frames / video_params.fps + + resp = VideoResponse( + model=model_name, + status=VideoGenerationStatus.QUEUED, + size=req.size, + prompt=req.prompt, + quality=req.quality or "default", + fps=video_params.fps, + num_frames=video_params.num_frames, + duration_s=duration_s, + ) + resp.seconds = str(req.seconds or resp.seconds) + return resp + + +def _status_code_for_video_failure(error: VideoError | None) -> int: + if error is None: + return HTTPStatus.INTERNAL_SERVER_ERROR.value + + if isinstance(error.code, int): + if 400 <= error.code < 600: + return error.code + return HTTPStatus.INTERNAL_SERVER_ERROR.value + + if error.code == "HTTPException": + status_text, _, _ = error.message.partition(":") + try: + status_code = int(status_text) + except ValueError: + return HTTPStatus.INTERNAL_SERVER_ERROR.value + if 400 <= status_code < 600: + return status_code + return HTTPStatus.INTERNAL_SERVER_ERROR.value + + if error.code == "EngineDeadError": + return HTTPStatus.INTERNAL_SERVER_ERROR.value + if error.code == "EngineGenerateError": + return HTTPStatus.INTERNAL_SERVER_ERROR.value + + return HTTPStatus.INTERNAL_SERVER_ERROR.value + + +def _video_error_from_exception(exc: Exception) -> VideoError: + if isinstance(exc, HTTPException): + message = str(exc.detail) if exc.detail else str(exc) + return VideoError(code=exc.status_code, message=message) + + if isinstance(exc, OmniClientError): + return VideoError(code=exc.status_code, message=exc.message) + + if isinstance(exc, (EngineGenerateError, EngineDeadError)): + err = create_error_response(exc) + return VideoError(code=err.error.code, message=err.error.message) + + return VideoError( + code=HTTPStatus.INTERNAL_SERVER_ERROR.value, + message=str(exc), + ) + + +async def _cleanup_video(video_id: str): + try: + await STORAGE_MANAGER.delete(video_id) + except Exception: + logger.warning("Failed to cleanup partial video file '%s'", video_id) + + +def _cleanup_video_references( + reference_video: ReferenceVideo | None, + reference_audio: ReferenceAudio | None, + control_path: str | None = None, +) -> None: + if reference_video is not None: + for path in reference_video.cleanup_paths: + if os.path.exists(path): + os.unlink(path) + if reference_audio is not None: + cleanup_paths = reference_audio.cleanup_paths or tuple(_reference_list(reference_audio.path)) + for path in cleanup_paths: + if os.path.exists(path): + os.unlink(path) + if control_path is not None and os.path.exists(control_path): + os.unlink(control_path) + + +def _unpack_video_generation_result( + result: Sequence[object], +) -> tuple[bytes, dict[str, float], float, VideoAction | None, dict[str, object]]: + video_metadata: dict[str, object] = {} + if len(result) == 5: + video_bytes, stage_durations, peak_memory_mb, action, raw_metadata = result + if isinstance(raw_metadata, dict): + video_metadata = {str(key): value for key, value in raw_metadata.items()} + else: + video_bytes, stage_durations, peak_memory_mb, action = result + return ( + cast(bytes, video_bytes), + cast(dict[str, float], stage_durations), + float(cast(float, peak_memory_mb)), + cast(VideoAction | None, action), + video_metadata, + ) + + +async def _run_video_generation_job( + handler: OmniOpenAIServingVideo, + request: VideoGenerationRequest, + video_id: str, + reference_image: ReferenceImage | None = None, + reference_video: ReferenceVideo | None = None, + reference_audio: ReferenceAudio | None = None, + control_path: str | None = None, + app_state: Any | None = None, +) -> None: + job = await VIDEO_STORE.get(video_id) + if job is None: + logger.warning("Video job %s missing before generation task started; skipping", video_id) + _cleanup_video_references(reference_video, reference_audio, control_path) + return + + await VIDEO_STORE.update_fields(video_id, {"status": VideoGenerationStatus.IN_PROGRESS}) + started_at = time.perf_counter() + try: + video_bytes, stage_durations, peak_memory_mb, action, video_metadata = _unpack_video_generation_result( + await handler.generate_video_bytes( + request, + video_id, + reference_image=reference_image, + reference_video=reference_video, + reference_audio=reference_audio, + ) + ) + + save_context = await STORAGE_MANAGER.save(video_bytes, video_id) + logger.info("Video request %s persisted %s output file.", video_id, save_context.key) + + updated_fields = { + "status": VideoGenerationStatus.COMPLETED, + "progress": 100, + "file_name": f"{video_id}.{job.file_extension}", + "completed_at": save_context.created_at, + "inference_time_s": time.perf_counter() - started_at, + "stage_durations": stage_durations, + "peak_memory_mb": peak_memory_mb, + "action": action, + } + updated_fields.update({key: value for key, value in video_metadata.items() if value is not None}) + if save_context.expires_at is not None: + updated_fields["expires_at"] = save_context.expires_at + + await VIDEO_STORE.update_fields(video_id, updated_fields) + except (EngineGenerateError, EngineDeadError) as exc: + logger.exception("Video generation failed (engine error) for id=%s", video_id) + + await _cleanup_video(video_id) + await VIDEO_STORE.update_fields( + video_id, + { + "status": VideoGenerationStatus.FAILED, + "completed_at": int(time.time()), + "error": _video_error_from_exception(exc), + "inference_time_s": time.perf_counter() - started_at, + }, + ) + # Background tasks can't propagate exceptions to FastAPI handlers. + # Actively signal shutdown when the engine is dead. + if app_state is not None and isinstance(exc, EngineDeadError): + terminate_if_errored( + server=app_state.server, + engine=app_state.engine_client, + ) + except Exception as exc: + logger.exception("Video generation failed for id=%s", video_id) + + await _cleanup_video(video_id) + await VIDEO_STORE.update_fields( + video_id, + { + "status": VideoGenerationStatus.FAILED, + "completed_at": int(time.time()), + "error": _video_error_from_exception(exc), + "inference_time_s": time.perf_counter() - started_at, + }, + ) + except asyncio.CancelledError: + await _cleanup_video(video_id) + await VIDEO_STORE.pop(video_id) + raise + finally: + _cleanup_video_references(reference_video, reference_audio, control_path) + + +async def _persist_uploaded_video_references(uploads: list[UploadFile]) -> list[str]: + paths: list[str] = [] + try: + for upload in uploads: + suffix = Path(upload.filename or "").suffix.lower() + if suffix not in {".mkv", ".mov", ".mp4", ".webm"}: + suffix = ".mp4" + fd, path = tempfile.mkstemp(prefix="vllm_omni_video_reference_", suffix=suffix) + paths.append(path) + with os.fdopen(fd, "wb") as output: + while chunk := await upload.read(1024 * 1024): + output.write(chunk) + except Exception: + for path in paths: + if os.path.exists(path): + os.unlink(path) + raise + return paths + + +async def _persist_uploaded_control_reference( + upload: UploadFile, + *, + max_bytes: int = CONTROL_REFERENCE_MAX_BYTES, +) -> str: + """Stream one model control upload to request-scoped local storage.""" + kind = _uploaded_media_kind(upload) + suffix = Path(upload.filename or "").suffix.lower() + supported_suffixes = CONTROL_REFERENCE_IMAGE_SUFFIXES | CONTROL_REFERENCE_VIDEO_SUFFIXES + if kind == "audio" or (suffix and suffix not in supported_suffixes): + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="control_reference must be an image or video file.", + ) + if not suffix: + suffix = ".png" if kind == "image" else ".mp4" + + declared_size = getattr(upload, "size", None) + if isinstance(declared_size, Integral) and int(declared_size) > max_bytes: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"control_reference exceeds the {max_bytes // (1024 * 1024)} MiB size limit.", + ) + + fd, path = tempfile.mkstemp(prefix="vllm_omni_control_reference_", suffix=suffix) + size = 0 + persisted = False + try: + with os.fdopen(fd, "wb") as output: + while chunk := await upload.read(min(1024 * 1024, max_bytes - size + 1)): + size += len(chunk) + if size > max_bytes: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"control_reference exceeds the {max_bytes // (1024 * 1024)} MiB size limit.", + ) + output.write(chunk) + if size == 0: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="control_reference must not be empty.", + ) + persisted = True + return path + finally: + if not persisted: + with suppress(OSError): + os.unlink(path) + + +def _validate_control_upload( + handler: OmniOpenAIServingVideo, + request: VideoGenerationRequest, + control_reference: UploadFile | None, + control_type: str | None, +) -> str | None: + """Validate a generic uploaded control against model capabilities.""" + if control_reference is None: + if control_type is not None: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="control_type requires a control_reference upload.", + ) + return None + if control_type is None or not control_type.strip(): + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="control_reference requires control_type.", + ) + + normalized_type = control_type.strip().lower() + supported_types = frozenset(getattr(handler, "supported_control_upload_types", ())) + if normalized_type not in supported_types: + supported = ", ".join(sorted(supported_types)) or "none" + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=( + f"Control upload type '{normalized_type}' is not supported by this model. " + f"Supported control upload types: {supported}." + ), + ) + + existing = (request.extra_params or {}).get(normalized_type) + if existing is not None: + if not isinstance(existing, Mapping) and existing is not True: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"extra_params.{normalized_type} must be an object when control_reference is uploaded.", + ) + if isinstance(existing, Mapping) and any(existing.get(key) is not None for key in ("control", "control_path")): + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=( + "Provide either control_reference or " + f"extra_params.{normalized_type}.control/control_path, not both." + ), + ) + return normalized_type + + +def _attach_control_upload( + request: VideoGenerationRequest, + control_type: str, + control_path: str | None = None, +) -> None: + """Declare an uploaded control and attach its persisted path when available.""" + extra_params = dict(request.extra_params or {}) + existing = extra_params.get(control_type) + control_params = dict(existing) if isinstance(existing, Mapping) else {} + if control_path is not None: + control_params["control_path"] = control_path + extra_params[control_type] = control_params + request.extra_params = extra_params + + +def _reference_list(value: Any) -> list[Any]: + if value is None: + return [] + return list(value) if isinstance(value, list) else [value] + + +def _uploaded_media_kind(upload: UploadFile) -> str: + content_type = (upload.content_type or "").lower() + if content_type.startswith("image/"): + return "image" + if content_type.startswith("audio/"): + return "audio" + suffix = Path(upload.filename or "").suffix.lower() + if suffix in {".png", ".jpg", ".jpeg", ".webp", ".heic", ".heif"}: + return "image" + if suffix in {".wav", ".mp3", ".m4a", ".aac", ".flac", ".ogg"}: + return "audio" + return "video" + + +def _minimax_h3_upload_limit(upload: UploadFile) -> int: + kind = _uploaded_media_kind(upload) + if kind == "image": + return MINIMAX_H3_MAX_REFERENCE_IMAGE_BYTES + if kind == "audio": + return MINIMAX_H3_MAX_REFERENCE_AUDIO_BYTES + return MINIMAX_H3_MAX_REFERENCE_VIDEO_BYTES + + +async def _read_upload_limited(upload: UploadFile, *, max_bytes: int | None = None) -> bytes: + """Read an upload with an optional hard byte limit.""" + if max_bytes is None: + return await upload.read() + + declared_size = getattr(upload, "size", None) + if isinstance(declared_size, Integral) and int(declared_size) > max_bytes: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"Uploaded reference exceeds the {max_bytes // (1024 * 1024)} MiB size limit.", + ) + + chunks: list[bytes] = [] + total = 0 + while True: + chunk = await upload.read(min(1024 * 1024, max_bytes - total + 1)) + if not chunk: + break + total += len(chunk) + if total > max_bytes: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"Uploaded reference exceeds the {max_bytes // (1024 * 1024)} MiB size limit.", + ) + chunks.append(chunk) + return b"".join(chunks) + + +def _validate_minimax_h3_image_payload( + payload: bytes, + *, + filename: str | None, + allow_non_image: bool = False, +) -> None: + """Validate a H3 image before converting it to a format-less PIL image.""" + try: + with Image.open(io.BytesIO(payload)) as image: + image_format = str(image.format or "").lower() + if image_format in MINIMAX_H3_REFERENCE_IMAGE_FORMATS: + _validate_image_pixel_limit(image) + except InvalidInputReferenceError as exc: + raise HTTPException(HTTPStatus.BAD_REQUEST.value, detail=str(exc)) from exc + except Image.DecompressionBombError as exc: + raise HTTPException( + HTTPStatus.BAD_REQUEST.value, + detail=f"Invalid uploaded image reference: {filename}; image exceeds the decoder pixel limit.", + ) from exc + except (OSError, ValueError) as exc: + if allow_non_image: + return + raise HTTPException(400, detail=f"Invalid uploaded image reference: {filename}") from exc + + if image_format not in MINIMAX_H3_REFERENCE_IMAGE_FORMATS: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=( + "MiniMax H3 reference images must use JPG, JPEG, PNG, WEBP, HEIC, or HEIF; " + f"got {image_format or 'unknown'}." + ), + ) + + +async def _persist_uploaded_media_references( + uploads: list[UploadFile], +) -> tuple[list[Image.Image], list[str], list[str]]: + """Persist a mixed MiniMax H3 multipart reference list. + + Images are decoded in memory; videos and audio remain files because H3's + reference encoders need the original container streams (including video + soundtracks). + """ + images: list[Image.Image] = [] + videos: list[str] = [] + audios: list[str] = [] + paths: list[str] = [] + if len(uploads) > MINIMAX_H3_MAX_REFERENCE_COUNT: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=f"MiniMax H3 accepts at most {MINIMAX_H3_MAX_REFERENCE_COUNT} total references.", + ) + try: + for upload in uploads: + kind = _uploaded_media_kind(upload) + payload = await _read_upload_limited(upload, max_bytes=_minimax_h3_upload_limit(upload)) + if kind == "image": + try: + _validate_minimax_h3_image_payload(payload, filename=upload.filename) + images.append( + _decode_image_bytes( + payload, + source=f"uploaded image reference: {upload.filename}", + ) + ) + except InvalidInputReferenceError as exc: + raise HTTPException(HTTPStatus.BAD_REQUEST.value, detail=str(exc)) from exc + continue + suffix = Path(upload.filename or "").suffix.lower() + if kind == "video" and suffix and suffix not in MINIMAX_H3_REFERENCE_VIDEO_SUFFIXES: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="MiniMax H3 reference videos must use an MP4 or MOV file.", + ) + if kind == "audio" and suffix and suffix not in MINIMAX_H3_REFERENCE_AUDIO_SUFFIXES: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="MiniMax H3 reference audio must use a WAV or MP3 file.", + ) + if not suffix or len(suffix) > 8: + suffix = ".mp3" if kind == "audio" else ".mp4" + fd, path = tempfile.mkstemp(prefix="vllm_omni_reference_", suffix=suffix) + paths.append(path) + with os.fdopen(fd, "wb") as output: + output.write(payload) + if kind == "audio": + audios.append(path) + else: + videos.append(path) + except Exception: + for path in paths: + if os.path.exists(path): + os.unlink(path) + raise + return images, videos, audios + + +async def _parse_video_form( + raw_request: Request, + prompt: str = Form(...), + input_reference: UploadFile | None = File(default=None), + input_references: list[UploadFile] | None = File(default=None), + control_reference: UploadFile | None = File(default=None), + control_type: str | None = Form(default=None), + image_reference: str | None = Form(default=None), + video_reference: str | None = Form(default=None), + audio_reference: str | None = Form(default=None), + model: str | None = Form(default=None), + seconds: SecondStr | None = Form(default=None), + size: SizeStr | None = Form(default=None), + user: str | None = Form(default=None), + width: int | None = Form(default=None), + height: int | None = Form(default=None), + num_frames: int | None = Form(default=None), + fps: float | None = Form(default=None, ge=1, allow_inf_nan=False), + aspect_ratio: str | None = Form(default=None), + short_edge: int | None = Form(default=None, ge=1), + num_outputs_per_prompt: int = Form(default=1, ge=1, le=10), + start_time_seconds: float | None = Form(default=None, ge=0.0), + quality: str | None = Form(default=None), + num_inference_steps: int | None = Form(default=None), + guidance_scale: float | None = Form(default=None), + guidance_scale_2: float | None = Form(default=None), + boundary_ratio: float | None = Form(default=None), + flow_shift: float | None = Form(default=None), + true_cfg_scale: float | None = Form(default=None), + seed: int | None = Form(default=None), + generate_sound: bool | None = Form(default=None), + sound_duration: float | None = Form(default=None, gt=0.0), + negative_prompt: str | None = Form(default=None), + enable_frame_interpolation: bool | None = Form(default=None), + frame_interpolation_exp: int | None = Form(default=None, ge=1), + frame_interpolation_scale: float | None = Form(default=None, gt=0.0), + frame_interpolation_model_path: str | None = Form(default=None), + lora: str | None = Form(default=None), + extra_params: str | None = Form(default=None), + return_stage_metrics: bool | None = Form(default=None), +) -> tuple[ + VideoGenerationRequest, + "OmniOpenAIServingVideo", + str, + ReferenceImage | None, + ReferenceVideo | None, + ReferenceAudio | None, + str | None, +]: + """FastAPI dependency that parses video form data, validates inputs, + resolves the handler, and decodes any reference image. + + Used by both ``POST /v1/videos`` (async) and ``POST /v1/videos/sync``. + """ + input_references = input_references or [] + input_reference_bytes: bytes | None = None + parsed_image_reference = _parse_form_json(image_reference) + parsed_video_reference = _parse_form_json(video_reference) + parsed_audio_reference = _parse_form_json(audio_reference) + + if input_references and any( + item is not None for item in (parsed_image_reference, parsed_video_reference, input_reference) + ): + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="Provide input_references alone, without input_reference, image_reference, or video_reference.", + ) + if input_reference is not None and (parsed_image_reference is not None or parsed_video_reference is not None): + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=( + "Provide only one of input_reference, image_reference, or video_reference when using " + "input_reference; image_reference and video_reference may be combined." + ), + ) + + request_data: dict[str, Any] = { + "prompt": prompt, + "model": model, + "seconds": seconds, + "size": size, + "image_reference": parsed_image_reference, + "video_reference": parsed_video_reference, + "audio_reference": parsed_audio_reference, + "user": user, + "width": width, + "height": height, + "num_frames": num_frames, + "fps": fps, + "aspect_ratio": aspect_ratio, + "short_edge": short_edge, + "num_outputs_per_prompt": num_outputs_per_prompt, + "start_time_seconds": start_time_seconds, + "quality": quality, + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + "guidance_scale_2": guidance_scale_2, + "boundary_ratio": boundary_ratio, + "flow_shift": flow_shift, + "true_cfg_scale": true_cfg_scale, + "seed": seed, + "generate_sound": generate_sound, + "sound_duration": sound_duration, + "negative_prompt": negative_prompt, + "enable_frame_interpolation": enable_frame_interpolation, + "frame_interpolation_exp": frame_interpolation_exp, + "frame_interpolation_scale": frame_interpolation_scale, + "frame_interpolation_model_path": frame_interpolation_model_path, + "lora": _parse_form_json(lora, expected_type=dict), + "extra_params": _parse_form_json(extra_params, expected_type=dict), + "return_stage_metrics": return_stage_metrics, + } + request_data = {k: v for k, v in request_data.items() if v is not None} + request = VideoGenerationRequest(**request_data) + + handler = Omnivideo(raw_request) + if handler is None: + raise HTTPException( + status_code=HTTPStatus.SERVICE_UNAVAILABLE.value, + detail="Video generation handler not initialized.", + ) + logger.info("Video generation handler: %s", type(handler).__name__) + try: + app_model_name, app_stage_configs = _resolve_video_runtime_context(raw_request) + effective_model_name = handler.model_name or app_model_name or request.model or "unknown" + if request.model is not None and effective_model_name is not None and request.model != effective_model_name: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=( + f"Model mismatch: request specifies '{request.model}' but server is running " + f"'{effective_model_name}'." + ), + ) + handler.set_stage_configs_if_missing(app_stage_configs) + except HTTPException: + raise + except Exception as e: + logger.exception("Video generation setup failed: %s", e) + raise HTTPException( + status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, + detail=f"Video generation setup failed: {str(e)}", + ) + + normalized_control_type = _validate_control_upload(handler, request, control_reference, control_type) + if normalized_control_type is not None: + # Make the selected transfer mode visible while choosing the model's + # reference-video decode policy. The upload itself stays unpersisted + # until reference parsing succeeds, preserving failure cleanup. + _attach_control_upload(request, normalized_control_type) + + supports_mixed_reference_inputs = bool(getattr(handler, "supports_mixed_reference_inputs", False)) + if input_reference is not None: + input_reference_bytes = await _read_upload_limited( + input_reference, + max_bytes=_minimax_h3_upload_limit(input_reference) if supports_mixed_reference_inputs else None, + ) + if supports_mixed_reference_inputs: + input_reference_kind = _uploaded_media_kind(input_reference) + _validate_minimax_h3_image_payload( + input_reference_bytes, + filename=input_reference.filename, + allow_non_image=input_reference_kind != "image", + ) + + decode_spec = ReferenceVideoDecodeSpec() + if not input_references and (parsed_video_reference is not None or input_reference_bytes is not None): + stage_configs = ( + handler.stage_configs + or app_stage_configs + or getattr(getattr(handler, "_engine_client", None), "stage_configs", None) + ) + decode_spec = _reference_video_decode_spec(request, stage_configs) + reference_image = None + reference_video = None + reference_audio: ReferenceAudio | None = None + if input_references: + if not supports_mixed_reference_inputs: + video_paths = await _persist_uploaded_video_references(input_references) + reference_video = ReferenceVideo(data=video_paths, cleanup_paths=tuple(video_paths)) + images, audio_paths = [], [] + else: + images, video_paths, audio_paths = await _persist_uploaded_media_references(input_references) + if images: + reference_image = ReferenceImage(data=images if len(images) > 1 else images[0]) + if video_paths: + reference_video = ReferenceVideo(data=video_paths, cleanup_paths=tuple(video_paths)) + if audio_paths: + reference_audio = ReferenceAudio(path=audio_paths, cleanup_paths=tuple(audio_paths)) + else: + video_paths: list[str] = [] + try: + image_items = _reference_list(request.image_reference) + video_items = _reference_list(request.video_reference) + image_data = [] + for item in image_items: + media_data = await decode_input_reference(item, None, None) + if not isinstance(media_data, Image.Image): + raise InvalidInputReferenceError("image_reference did not decode to an image") + image_data.append(media_data) + + video_frames: list[Image.Image] | None = None + for item in video_items: + media_data = await decode_input_reference( + None, + item, + None, + max_video_frames=decode_spec.max_frames, + video_keep=decode_spec.keep, + ) + if not isinstance(media_data, VideoFrames): + raise InvalidInputReferenceError("video_reference did not decode to a video") + if media_data.source_path is not None: + video_paths.append(media_data.source_path) + else: + if len(video_items) != 1: + raise InvalidInputReferenceError( + "multiple video URL references must be downloadable source videos" + ) + video_frames = list(media_data) + + if input_reference_bytes is not None: + media_data = await decode_input_reference( + None, + None, + input_reference_bytes, + max_video_frames=decode_spec.max_frames, + video_keep=decode_spec.keep, + ) + if isinstance(media_data, Image.Image): + image_data.append(media_data) + elif isinstance(media_data, VideoFrames): + if media_data.source_path is not None: + video_paths.append(media_data.source_path) + else: + video_frames = list(media_data) + + if image_data: + reference_image = ReferenceImage(data=image_data if len(image_data) > 1 else image_data[0]) + if video_paths: + reference_video = ReferenceVideo(data=video_paths, cleanup_paths=tuple(video_paths)) + elif video_frames is not None: + reference_video = ReferenceVideo(data=video_frames) + except InvalidInputReferenceError as exc: + for path in video_paths: + if os.path.exists(path): + os.unlink(path) + raise HTTPException(400, detail=str(exc) or "Invalid input reference.") from exc + + audio_paths = [] if reference_audio is None else list(_reference_list(reference_audio.path)) + if request.audio_reference is not None: + try: + for audio_reference in _reference_list(request.audio_reference): + audio_paths.append(await decode_audio_url(audio_reference.audio_url)) + except InvalidInputReferenceError as exc: + _cleanup_video_references(reference_video, reference_audio) + cleanup_paths = set(() if reference_audio is None else reference_audio.cleanup_paths) + for path in audio_paths: + if path not in cleanup_paths and os.path.exists(path): + os.unlink(path) + raise HTTPException(400, detail=str(exc)) from exc + if audio_paths: + cleanup_paths = ( + tuple(audio_paths) + if reference_audio is None + else reference_audio.cleanup_paths + + tuple(path for path in audio_paths if path not in reference_audio.cleanup_paths) + ) + reference_audio = ReferenceAudio( + path=audio_paths if len(audio_paths) > 1 else audio_paths[0], + cleanup_paths=cleanup_paths, + ) + + control_path: str | None = None + if control_reference is not None and normalized_control_type is not None: + try: + control_path = await _persist_uploaded_control_reference( + control_reference, + max_bytes=CONTROL_REFERENCE_MAX_BYTES, + ) + _attach_control_upload(request, normalized_control_type, control_path) + except (asyncio.CancelledError, HTTPException, OSError, TypeError, ValueError): + _cleanup_video_references(reference_video, reference_audio, control_path) + raise + + return ( + request, + handler, + effective_model_name, + reference_image, + reference_video, + reference_audio, + control_path, + ) diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index 0f01fb28f91..67e8bcd36f5 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -1,7 +1,26 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ -Shared helper utilities for OpenAI-compatible video generation API. +Shared media utilities for OpenAI-compatible video APIs. + +PUT HERE: + - Shared media backends reused by generation and streaming serving: + decode/encode of image/video/audio references and frames, streaming + encoders, frame/audio coercion. No FastAPI Request / job-store orchestration. + +DO NOT PUT HERE: + - ``/v1/videos*`` multipart form parsing, upload limits, job runners, + cleanup, or response factories — those go in ``video.generation.helpers``. + +LONGEVITY: + - This root utils file is a **temporary shared home**. + - TODO(#5227, P1.3): tidy up / move into the video family (e.g. + ``video/generation/media.py``) in the P1.3 video modality PR; do not treat this + file as the long-term owner. + - ``video.generation.helpers`` is the longer home for ``/v1/videos*`` + endpoint logic through P0.2/P0.3 until P1.3 further splits it. + +See ``openai/README.md`` and ``video/README.md`` (utils vs helpers, no overlap). """ from __future__ import annotations diff --git a/vllm_omni/entrypoints/serve/README.md b/vllm_omni/entrypoints/serve/README.md new file mode 100644 index 00000000000..6fe6bf3823b --- /dev/null +++ b/vllm_omni/entrypoints/serve/README.md @@ -0,0 +1,22 @@ +# Serve Entrypoint Helpers + +This package contains server-level helpers for assembling and operating the +HTTP/FastAPI serving process. + +## What Belongs Here + +- Helpers that are about the server app rather than a public OpenAI endpoint + family. +- Operational route packages such as profiling and Omni control. +- Shared server utilities that mirror upstream `vllm.entrypoints.serve.utils` + concepts. + +## What Does Not Belong Here + +- OpenAI protocol schemas or response bodies. +- Image, video, audio, or streaming-input request parsing. +- Serving model implementations. + +Endpoint-family route bodies are still staged in `openai/api_server.py` during +the P0.2 helper migration. Route ownership (P0.3) and final folder cleanup (Phase 1) +are tracked by #5227. diff --git a/vllm_omni/entrypoints/serve/__init__.py b/vllm_omni/entrypoints/serve/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm_omni/entrypoints/serve/omni_control/README.md b/vllm_omni/entrypoints/serve/omni_control/README.md new file mode 100644 index 00000000000..713e2aac594 --- /dev/null +++ b/vllm_omni/entrypoints/serve/omni_control/README.md @@ -0,0 +1,19 @@ +# Omni Control Routes + +This package owns operational Omni server-control routes such as sleep and +wakeup. + +## Put Here + +- Request/response models for `/v1/omni/*` operational controls. +- Route bodies for Omni control endpoints after endpoint extraction. +- Small helpers that are specific to these operational controls. + +## Do Not Put Here + +- OpenAI modality endpoints. +- Generic server utilities that are shared outside Omni control. +- Engine implementation details for sleep or wakeup. + +The current route bodies still live in `openai/api_server.py` and should move +here as part of #5227. diff --git a/vllm_omni/entrypoints/serve/omni_control/__init__.py b/vllm_omni/entrypoints/serve/omni_control/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm_omni/entrypoints/serve/omni_control/protocol.py b/vllm_omni/entrypoints/serve/omni_control/protocol.py new file mode 100644 index 00000000000..f7ce7249195 --- /dev/null +++ b/vllm_omni/entrypoints/serve/omni_control/protocol.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Request models for Omni server control routes.""" + +from pydantic import BaseModel + + +class OmniSleepRequest(BaseModel): + stage_ids: list[int] + level: int = 2 + + +class OmniWakeupRequest(BaseModel): + stage_ids: list[int] diff --git a/vllm_omni/entrypoints/serve/profile/README.md b/vllm_omni/entrypoints/serve/profile/README.md new file mode 100644 index 00000000000..63d510580e7 --- /dev/null +++ b/vllm_omni/entrypoints/serve/profile/README.md @@ -0,0 +1,19 @@ +# Profile Routes + +This package owns profiling-control schemas and, after route extraction, should +own the profiling route bodies. + +## Put Here + +- Request/response models for `/start_profile` and `/stop_profile`. +- Profile-specific helper functions. +- The profile `api_router.py` once route bodies move out of `api_server.py`. + +## Do Not Put Here + +- General server error handling. +- Engine profiling implementation internals. +- OpenAI modality endpoint logic. + +The route bodies still live in `openai/api_server.py` under Phase 0.2 (`P0.2`) of #5227. +TODO(#5227, P0.3): move route bodies here during router extraction. diff --git a/vllm_omni/entrypoints/serve/profile/__init__.py b/vllm_omni/entrypoints/serve/profile/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm_omni/entrypoints/serve/profile/protocol.py b/vllm_omni/entrypoints/serve/profile/protocol.py new file mode 100644 index 00000000000..ba424e924a7 --- /dev/null +++ b/vllm_omni/entrypoints/serve/profile/protocol.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Request models for Omni profile control routes.""" + +from pydantic import BaseModel, Field + + +class ProfileRequest(BaseModel): + """Request model for profiling endpoints.""" + + stages: list[int] | None = Field( + default=None, + description="List of stage IDs to profile. If None, profiles all stages.", + ) diff --git a/vllm_omni/entrypoints/serve/profile/utils.py b/vllm_omni/entrypoints/serve/profile/utils.py new file mode 100644 index 00000000000..f444a35341f --- /dev/null +++ b/vllm_omni/entrypoints/serve/profile/utils.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Utilities for Omni profile control routes.""" + + +def _should_enable_profiler_endpoints(stage_configs: list | None) -> bool: + """Check if any stage has profiler_config set in its engine_args.""" + if not stage_configs: + return False + for stage in stage_configs: + engine_args = stage.get("engine_args") if isinstance(stage, dict) else getattr(stage, "engine_args", None) + if engine_args is None: + continue + profiler_config = ( + engine_args.get("profiler_config") + if isinstance(engine_args, dict) + else getattr(engine_args, "profiler_config", None) + ) + if profiler_config is not None: + profiler = ( + profiler_config.get("profiler") + if isinstance(profiler_config, dict) + else getattr(profiler_config, "profiler", None) + ) + if profiler is not None: + return True + return False diff --git a/vllm_omni/entrypoints/serve/utils/README.md b/vllm_omni/entrypoints/serve/utils/README.md new file mode 100644 index 00000000000..1babf1c8107 --- /dev/null +++ b/vllm_omni/entrypoints/serve/utils/README.md @@ -0,0 +1,23 @@ +# Serve Utilities + +This directory is for server/app utilities shared by Omni entrypoints. + +## Put Here + +- FastAPI app/router mutation helpers (`routes.remove_route_from_app`, + `routes._remove_route_from_router`). +- Server exception-handler registration. +- Engine-failure helpers that are not owned by one endpoint family. + +## Do Not Put Here + +- Endpoint restriction *policy* (which routes a pipeline shuts down, and why). + That stays in `vllm_omni.config.endpoint_policy`; the policy module should + *call* `remove_route_from_app` from here, not redefine it. +- OpenAI response-shape adapters. Those belong under + `vllm_omni.entrypoints.openai`. +- Modality validation or payload parsing. +- Generic helpers without a clear server/app responsibility. + +Keep this directory aligned with the spirit of upstream +`vllm.entrypoints.serve.utils`. Broader utility cleanup is tracked by #5227. diff --git a/vllm_omni/entrypoints/serve/utils/__init__.py b/vllm_omni/entrypoints/serve/utils/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm_omni/entrypoints/serve/utils/errors.py b/vllm_omni/entrypoints/serve/utils/errors.py new file mode 100644 index 00000000000..fc4c4438b23 --- /dev/null +++ b/vllm_omni/entrypoints/serve/utils/errors.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Server-level error helpers for vLLM-Omni entrypoints. + +Use this module for FastAPI/app exception handling and engine-failure response +construction that is not owned by a single OpenAI endpoint family.""" + +from typing import Any + +from fastapi import Request +from fastapi.responses import JSONResponse +from vllm.entrypoints.launchers.launcher import terminate_if_errored +from vllm.entrypoints.serve import create_error_response +from vllm.logger import init_logger +from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError + +logger = init_logger(__name__) + + +def _get_request_id_from_request(req: Request) -> str | None: + return req.state.request_metadata.request_id if hasattr(req.state, "request_metadata") else None + + +def _build_engine_error_payload( + exc: EngineDeadError | EngineGenerateError, + *, + request_id: str | None, +) -> tuple[dict[str, Any], int]: + err = create_error_response(exc) + payload = err.model_dump() + error_body = payload.get("error", {}) + + error_body["request_id"] = request_id + error_body["error_stage_id"] = getattr(exc, "error_stage_id", None) + + return payload, err.error.code + + +def _create_engine_error_json_response( + req: Request, + exc: EngineDeadError | EngineGenerateError, +) -> JSONResponse: + request_id = _get_request_id_from_request(req) + error_stage_id = getattr(exc, "error_stage_id", None) + engine = req.app.state.engine_client + + if isinstance(exc, EngineDeadError): + # Log Omni-specific diagnostic information for dead engines. + orchestrator_alive = engine.engine.is_alive() if hasattr(engine, "engine") else "N/A" + logger.error( + "EngineDeadError: orchestrator_alive=%s, errored=%s, request_id=%s, error_stage_id=%s", + orchestrator_alive, + engine.errored, + request_id, + error_stage_id, + ) + + terminate_if_errored( + server=req.app.state.server, + engine=engine, + ) + + payload, status_code = _build_engine_error_payload(exc, request_id=request_id) + return JSONResponse(content=payload, status_code=status_code) + + +def _register_omni_exception_handlers(app) -> None: + """Override upstream vLLM exception handlers with Omni-aware versions. + + The upstream ``engine_error_handler`` is designed for ``AsyncLLM`` (single + EngineCore process). Omni uses a multi-stage orchestrator with different + health semantics, so we register our own handlers that: + + - Log multi-stage diagnostic info (orchestrator liveness, per-stage health) + when an ``EngineDeadError`` is caught. + - Call ``terminate_if_errored`` + - Return an OpenAI-compatible error JSON response. + """ + + async def omni_engine_error_handler( + req: Request, + exc: EngineDeadError | EngineGenerateError, + ): + request_id = _get_request_id_from_request(req) + + if req.app.state.args.log_error_stack: + logger.exception("Engine Exception caught. Request id: %s", request_id) + + return _create_engine_error_json_response(req, exc) + + app.exception_handler(EngineGenerateError)(omni_engine_error_handler) + app.exception_handler(EngineDeadError)(omni_engine_error_handler) diff --git a/vllm_omni/entrypoints/serve/utils/routes.py b/vllm_omni/entrypoints/serve/utils/routes.py new file mode 100644 index 00000000000..55e4c0492ad --- /dev/null +++ b/vllm_omni/entrypoints/serve/utils/routes.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Route-table helpers for Omni server assembly. + +Owns FastAPI app **and** router mutation used during server construction. +``endpoint_policy`` imports app-level removal from here; it does not own +route-table helpers itself. +""" + +from collections.abc import Set + +from fastapi import APIRouter, FastAPI +from starlette.routing import Route + + +def remove_route_from_app( + app: FastAPI, + path: str, + methods: Set[str], +) -> None: + """Remove app routes matching a path and one of the given HTTP methods.""" + routes_to_remove = [ + route + for route in app.routes + if isinstance(route, Route) and route.path == path and route.methods is not None and route.methods & methods + ] + for route in routes_to_remove: + app.routes.remove(route) + + +def _remove_route_from_router( + router: APIRouter, + path: str, + methods: set[str] | None = None, +) -> None: + """Remove routes from an ``APIRouter`` by path and optionally by methods.""" + methods_set = {method.upper() for method in methods} if methods else None + for route in list(router.routes): + if getattr(route, "path", None) != path: + continue + if methods_set is not None: + route_methods = {method.upper() for method in (getattr(route, "methods", None) or set())} + if not (route_methods & methods_set): + continue + router.routes.remove(route) From a7c295d5442820d7390a39968aea5b9f135bbb2e Mon Sep 17 00:00:00 2001 From: LHXuuu Date: Thu, 10 Sep 2026 15:40:49 +0800 Subject: [PATCH 50/93] [CI] Stabilize Qwen3-Omni Server VAD E2E (#7356) Signed-off-by: LHXuuu Co-authored-by: amy-why-3459 --- .../test_qwen3_omni_realtime_websocket.py | 38 +++++++++++++++++-- 1 file changed, 35 insertions(+), 3 deletions(-) diff --git a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py index 8c80b964928..406668a4ce9 100644 --- a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py +++ b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py @@ -22,6 +22,7 @@ import pytest import websockets +from tests.e2e.online_serving.helpers.minicpmo_4_5_duplex import validated_input_wav from tests.helpers.mark import hardware_test from tests.helpers.media import ( convert_audio_bytes_to_text, @@ -306,6 +307,11 @@ def _synthetic_pcm16_input( return _pcm16_mono_16k_from_wav_bytes(wav_bytes) +def _server_vad_pcm16_input() -> bytes: + """Load the fixed single-turn speech fixture used by the Server VAD E2E.""" + return _pcm16_mono_16k_from_wav_bytes(validated_input_wav().read_bytes()) + + def _assert_realtime_smoke(result: dict) -> None: out_pcm = result["output_pcm"] assert result["delta_events"] >= 1 @@ -408,7 +414,6 @@ def test_long_audio_response_completes_async_chunk(self, omni_server) -> None: @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", realtime_server_vad_server_params, indirect=True) - @pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/7279") def test_server_vad_multi_turn_without_client_commit( self, cached_silero_vad_artifact: str, @@ -420,7 +425,7 @@ def test_server_vad_multi_turn_without_client_commit( assert deploy.session_mode == "turn" assert deploy.async_chunk is True assert deploy.duplex_session.server_vad_model_path == SERVER_VAD_MODEL_PATH - pcm16 = _synthetic_pcm16_input() + pcm16 = _server_vad_pcm16_input() turns = asyncio.run( _run_server_vad_audio_roundtrips( @@ -434,6 +439,12 @@ def test_server_vad_multi_turn_without_client_commit( ) assert len(turns) == 2 + updated_session = next(event["session"] for event in turns[0] if event["type"] == "session.updated") + effective_turn_detection = updated_session["audio"]["input"]["turn_detection"] + assert effective_turn_detection["type"] == "server_vad" + assert effective_turn_detection["silence_duration_ms"] == 500 + assert effective_turn_detection["create_response"] is True + assert effective_turn_detection["interrupt_response"] is False required_sequence = [ "input_audio_buffer.speech_started", "input_audio_buffer.speech_stopped", @@ -447,6 +458,9 @@ def test_server_vad_multi_turn_without_client_commit( response_ids: list[str] = [] for events in turns: event_types = [event["type"] for event in events] + for event_type in required_sequence: + if event_type != "response.audio.delta": + assert event_types.count(event_type) == 1, event_types positions = [event_types.index(event_type) for event_type in required_sequence] assert positions == sorted(positions) @@ -457,9 +471,27 @@ def test_server_vad_multi_turn_without_client_commit( done = next(event for event in events if event["type"] == "response.done")["response"] assert started["item_id"] == stopped["item_id"] == committed["item_id"] + input_item_id = committed["item_id"] + history_events = [ + event + for event in events + if event["type"] in {"conversation.item.added", "conversation.item.done"} + and event["item"]["id"] == input_item_id + ] + assert [event["type"] for event in history_events] == [ + "conversation.item.added", + "conversation.item.done", + ] + assert all(event["item"]["role"] == "user" for event in history_events) + output_pcm = b"".join( + base64.b64decode(event["delta"]) + for event in events + if event["type"] == "response.audio.delta" and event.get("delta") + ) + assert output_pcm assert created["id"] == done["id"] assert done["status"] == "completed" - input_item_ids.append(committed["item_id"]) + input_item_ids.append(input_item_id) response_ids.append(created["id"]) assert len(set(input_item_ids)) == 2 From 5b927b7ceb92d9b1998e22d50cdc4ee400db08d0 Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Thu, 10 Sep 2026 16:40:01 +0800 Subject: [PATCH 51/93] [CI/Build] Diff-aware source_file_dependencies for CUDA/NPU pipelines (#6597) Signed-off-by: wangyu <410167048@qq.com> Co-authored-by: Cursor --- .../common/ci_source_file_dependencies.yml | 438 ++++++++++++++++++ .buildkite/common/scripts/upload_pipeline.py | 263 +++++++++-- .buildkite/cuda/bootstrap-upload-steps.yml | 28 +- .buildkite/cuda/test-merge.yml | 129 +----- .buildkite/cuda/test-nightly.yml | 104 +++-- .buildkite/cuda/test-ready.yml | 146 +----- .buildkite/cuda/test-weekly.yml | 288 +++++++++++- .buildkite/npu/bootstrap-upload-steps.yml | 9 +- .buildkite/npu/pipeline-npu-a3.yml | 8 +- .buildkite/npu/test-npu-nightly.yml | 42 +- .claude/skills/vllm-omni-test/SKILL.md | 8 +- .../vllm-omni-test/references/test-routing.md | 4 +- docs/contributing/ci/ci_settings.md | 54 ++- tests/buildkite/test_upload_pipeline.py | 334 ++++++++++++- tests/e2e/accuracy/test_hidream_o1_image.py | 4 +- .../test_lingbot_world_v2.py | 4 +- .../test_minimax_h3_dlo_dp2_fl2va.py | 2 +- .../test_minimax_h3_dlo_dp2_ref2va.py | 2 +- .../minimax_h3/test_minimax_h3_fasth3_lora.py | 2 +- .../minimax_h3/test_minimax_h3_turbo_lora.py | 2 +- .../e2e/online_serving/test_lingbot_video.py | 6 +- .../online_serving/test_lingbot_video_moe.py | 6 +- tests/e2e/online_serving/test_magi2.py | 5 +- .../test_minimax_h3_dlo_dp2_t2va.py | 1 - tools/pre_commit/check_buildkite.py | 1 + 25 files changed, 1479 insertions(+), 411 deletions(-) create mode 100644 .buildkite/common/ci_source_file_dependencies.yml diff --git a/.buildkite/common/ci_source_file_dependencies.yml b/.buildkite/common/ci_source_file_dependencies.yml new file mode 100644 index 00000000000..0455922c8f8 --- /dev/null +++ b/.buildkite/common/ci_source_file_dependencies.yml @@ -0,0 +1,438 @@ +# Uploader-only path-prefix presets for CUDA L2–L5 and NPU L4 test YAML. +# Referenced as ``source_file_dependencies: omni_qwen3_omni_function`` (or a list of +# keys) on a step; ``upload_pipeline.py`` expands the key(s) only. +# +# YAML anchors above the registry hold model business-code paths. Job keys +# ({model}_function / _perf / _accuracy / _reliability / _doc / _cov) alias +# those anchors and add the scripts the job actually runs. Feature keys below +# are unchanged. ``source_filter_fallback`` is not a job key: when no job-key +# prefix matches, a matching diff keeps every job for that upload. +# +# YAML anchors are nested lists; the loader flattens them. + +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/snake_activation.py + - vllm_omni/model_executor/models/common/alias_free_activation.py + - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py + - vllm_omni/deploy/qwen3_tts.yaml + - vllm_omni/entrypoints/openai/serving_speech.py + - vllm_omni/entrypoints/openai/tts_adapters/ + - vllm_omni/engine/stage_pool.py + +x-minicpmo-4-5: &minicpmo_4_5 + - vllm_omni/model_executor/models/minicpmo_4_5/ + - vllm_omni/model_executor/models/step_audio2/ + - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ + - vllm_omni/model_executor/models/cosyvoice3/utils.py + - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py + - vllm_omni/deploy/minicpmo_4_5.yaml + - vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja + +x-minicpmo-4-5-duplex: &minicpmo_4_5_duplex + - vllm_omni/deploy/minicpmo_4_5_duplex.yaml + - vllm_omni/engine/duplex/ + - vllm_omni/entrypoints/duplex/ + - vllm_omni/clients/ + - vllm_omni/entrypoints/duplex_request_client.py + - vllm_omni/outputs/duplex.py + - vllm_omni/model_executor/duplex_sampling.py + - vllm_omni/experimental/fullduplex/ + - vllm_omni/core/sched/ + - vllm_omni/outputs/multimodal_accumulation.py + - vllm_omni/engine/stage_init_utils.py + - vllm_omni/engine/async_omni_engine.py + - vllm_omni/engine/orchestrator.py + - vllm_omni/entrypoints/cli/serve.py + - vllm_omni/entrypoints/openai/api_server.py + - vllm_omni/entrypoints/openai/duplex_capability.py + - vllm_omni/config/stage_config.py + +x-qwen2-5-omni: &qwen2_5_omni + - vllm_omni/model_executor/models/qwen2_5_omni/ + - vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py + - vllm_omni/deploy/qwen2_5_omni.yaml + +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/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_thinking.yaml + +x-wan22: &wan22 + - vllm_omni/diffusion/models/wan2_2/ + - vllm_omni/diffusion/models/dmd2/ + +x-cosmos3: &cosmos3 + - vllm_omni/diffusion/models/cosmos3/ + - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py + - vllm_omni/experimental/world_models/ + - vllm_omni/model_extras/cosmos3.py + +x-qwen-image: &qwen_image + - vllm_omni/diffusion/models/qwen_image/ + - vllm_omni/diffusion/models/dmd2/ + - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py + +x-hunyuan-image3: &hunyuan_image3 + - vllm_omni/diffusion/models/hunyuan_image3/ + - vllm_omni/model_executor/models/hunyuan_image3/ + - vllm_omni/deploy/hunyuan_image3_dit.yaml + +x-minimax-h3: &minimax_h3 + - vllm_omni/diffusion/models/minimax_h3/ + +x-hunyuan-video: &hunyuan_video + - vllm_omni/diffusion/models/hunyuan_video/ + - vllm_omni/model_executor/models/hunyuan_video/ + +x-ltx2: <x2 + - vllm_omni/diffusion/models/ltx2/ + - vllm_omni/model_extras/ltx2.py + +x-sana-video: &sana_video + - vllm_omni/diffusion/models/sana_video/ + - vllm_omni/model_extras/sana_video.py + +x-z-image: &z_image + - vllm_omni/diffusion/models/z_image/ + +x-stable-audio: &stable_audio + - vllm_omni/diffusion/models/stable_audio/ + +x-lingbot-video: &lingbot_video + - vllm_omni/diffusion/models/lingbot_video/ + - vllm_omni/model_extras/lingbot_video.py + +x-lingbot-world: &lingbot_world + - vllm_omni/diffusion/models/lingbot_world/ + + +source_file_dependencies: + # Shared uploader/registry paths. Not a job key — do not attach this to a step. + # Used only when no job-key prefix matches: then a hit keeps every job. + source_filter_fallback: + - .buildkite/common/scripts/upload_pipeline.py + - .buildkite/common/ci_source_file_dependencies.yml + - .buildkite/common/ci_mirror_hardwares.yml + - tests/helpers/ + + # --- features --- + diffusion_distributed_attention: + - tests/diffusion/distributed/ + - tests/diffusion/attention/test_attention_sp.py + - vllm_omni/diffusion/distributed/ + - vllm_omni/diffusion/attention/ + - vllm_omni/diffusion/data.py + - vllm_omni/diffusion/models/ + + custom_pipeline: + - tests/e2e/features/custom_pipeline/ + - tests/e2e/features/helpers/custom_pipeline.py + - vllm_omni/diffusion/ + - vllm_omni/engine/ + - vllm_omni/entrypoints/ + - vllm_omni/inputs/ + - vllm_omni/outputs/ + + rlhf_verl_omni: + - tests/e2e/features/rlhf_test/test_verl_omni_e2e.py + - tests/e2e/features/helpers/verl_omni_server.py + - tests/e2e/features/helpers/custom_pipeline.py + - vllm_omni/diffusion/ + - vllm_omni/engine/ + - vllm_omni/entrypoints/ + - vllm_omni/inputs/ + - vllm_omni/outputs/ + + # Tiny-model diffusion jobs (shared across many model trees under models/). + diffusion_tiny_model: + - tests/model_tests/diffusion/ + - vllm_omni/diffusion/models/ + + diffusion_offloader: + - tests/diffusion/offloader/ + - vllm_omni/diffusion/offloader/ + + diffusion_quantization: + - tests/diffusion/quantization/ + - vllm_omni/diffusion/quantization/ + + # --- models --- + # --- omni / tts jobs --- + omni_qwen3_omni_function: + - *qwen3_omni + - tests/e2e/offline_inference/test_qwen3_omni.py + - tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py + - tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16_expansion.py + - tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py + - tests/e2e/online_serving/test_qwen3_omni.py + - tests/e2e/online_serving/test_qwen3_omni_expansion.py + - tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py + + omni_qwen3_omni_accuracy: + - *qwen3_omni + - tests/e2e/accuracy/qwen3_omni/ + + omni_qwen3_omni_perf: + - *qwen3_omni + - tests/dfx/perf/scripts/run_benchmark.py + - tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json + - tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json + - tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json + - tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json + + omni_qwen3_omni_reliability: + - *qwen3_omni + - tests/dfx/reliability/test_reliability_qwen3_omni.py + + omni_qwen3_omni_doc: + - *qwen3_omni + - examples/online_serving/qwen3_omni/ + - examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py + - tests/examples/online_serving/test_qwen3_omni.py + + omni_qwen2_5_omni_doc: + - *qwen2_5_omni + - examples/online_serving/qwen2_5_omni/ + - examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py + - tests/examples/online_serving/test_qwen2_5_omni.py + + omni_minicpmo_4_5_function: + - *minicpmo_4_5 + - tests/e2e/offline_inference/test_minicpmo_4_5.py + - tests/e2e/online_serving/test_minicpmo_4_5.py + - tests/e2e/online_serving/test_minicpmo_4_5_expansion.py + + omni_minicpmo_4_5_accuracy: + - *minicpmo_4_5 + - tests/e2e/accuracy/minicpmo_4_5/ + + omni_minicpmo_4_5_perf: + - *minicpmo_4_5 + - tests/dfx/perf/scripts/run_benchmark.py + - tests/dfx/perf/tests/test_minicpmo_4_5_ready.json + - tests/dfx/perf/tests/test_minicpmo_4_5.json + + omni_minicpmo_4_5_doc: + - *minicpmo_4_5 + - examples/online_serving/minicpmo/realtime_web/ + - examples/online_serving/minicpmo/realtime_duplex_demo.py + - tests/examples/test_minicpmo_realtime_web_server.py + - tests/examples/test_minicpmo_realtime_web_static.py + - tests/examples/test_minicpmo_realtime_duplex_simple_demo.py + + omni_minicpmo_4_5_duplex_function: + - *minicpmo_4_5 + - *minicpmo_4_5_duplex + - tests/e2e/online_serving/test_minicpmo_4_5_duplex.py + - tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py + - tests/e2e/online_serving/test_duplex_client_live.py + - tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py + - tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py + - tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py + - tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py + - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py + - tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py + - tests/assets/minicpmo_4_5/ + - tests/clients/ + - tests/engine/duplex/ + - tests/entrypoints/duplex/ + - tests/model_executor/models/minicpmo_4_5/duplex/ + + omni_minicpmo_4_5_duplex_perf: + - *minicpmo_4_5 + - *minicpmo_4_5_duplex + - tests/dfx/perf/scripts/run_benchmark.py + - tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json + + tts_qwen3_tts_function: + - *qwen3_tts + - tests/e2e/features/prefill_phase/ + - tests/e2e/offline_inference/test_qwen3_tts_base.py + - tests/e2e/offline_inference/test_qwen3_tts_customvoice.py + - tests/e2e/online_serving/test_qwen3_tts_base.py + - tests/e2e/online_serving/test_qwen3_tts_base_expansion.py + - tests/e2e/online_serving/test_qwen3_tts_customvoice.py + - tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py + + tts_qwen3_tts_perf: + - *qwen3_tts + - tests/dfx/perf/scripts/run_benchmark.py + - tests/dfx/perf/tests/test_tts.json + + tts_qwen3_tts_cov: + - *qwen3_tts + - .buildkite/common/scripts/run_cov_split.sh + - pyproject.toml + - tests/e2e/offline_inference/test_qwen3_tts_base.py + - tests/e2e/online_serving/test_qwen3_tts_base.py + + # --- diffusion jobs --- + diffusion_wan22_function: + - *wan22 + - tests/e2e/offline_inference/test_wan22_t2v.py + - tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py + - tests/e2e/online_serving/test_wan22_t2v.py + - tests/e2e/online_serving/test_wan22_expansion.py + + diffusion_wan22_accuracy: + - *wan22 + - tests/e2e/accuracy/wan22_i2v/ + - tests/e2e/accuracy/test_diffusers_backend_similarity.py + + diffusion_wan22_perf: + - *wan22 + - tests/dfx/perf/scripts/run_diffusion_benchmark.py + - tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json + + diffusion_wan22_reliability: + - *wan22 + - tests/dfx/reliability/test_reliability_wan22.py + + diffusion_image_to_video_doc: + - *wan22 + - *hunyuan_video + - *ltx2 + - *sana_video + - *lingbot_video + - examples/offline_inference/image_to_video/ + - tests/examples/offline_inference/test_image_to_video.py + + diffusion_cosmos3_function: + - *cosmos3 + - tests/e2e/online_serving/test_cosmos3.py + + diffusion_cosmos3_accuracy: + - *cosmos3 + - tests/e2e/accuracy/test_cosmos3_similarity.py + + diffusion_cosmos3_perf: + - *cosmos3 + - tests/dfx/perf/scripts/run_diffusion_benchmark.py + - tests/dfx/perf/tests/test_cosmos3_vllm_omni.json + + diffusion_qwen_image_function: + - *qwen_image + - tests/e2e/online_serving/test_qwen_image.py + - tests/e2e/online_serving/test_qwen_image_expansion.py + + diffusion_qwen_image_accuracy: + - *qwen_image + - tests/e2e/accuracy/test_qwen_image.py + - tests/e2e/accuracy/test_diffusers_backend_similarity.py + + diffusion_qwen_image_perf: + - *qwen_image + - tests/dfx/perf/scripts/run_diffusion_benchmark.py + - tests/dfx/perf/tests/test_qwen_image_vllm_omni.json + + diffusion_text_to_image_doc: + - *qwen_image + - *z_image + - examples/offline_inference/text_to_image/ + - examples/online_serving/text_to_image/ + - tests/examples/offline_inference/test_text_to_image.py + - tests/examples/online_serving/test_text_to_image.py + - vllm_omni/diffusion/models/nextstep_1_1/ + - vllm_omni/diffusion/models/longcat_image/ + - vllm_omni/diffusion/models/ovis_image/ + - vllm_omni/diffusion/models/omnigen2/ + - vllm_omni/diffusion/models/sd3/ + - vllm_omni/diffusion/models/flux/ + - vllm_omni/diffusion/models/flux2/ + - vllm_omni/diffusion/models/flux2_klein/ + - vllm_omni/diffusion/models/hidream_image/ + - vllm_omni/diffusion/models/krea2/ + - vllm_omni/diffusion/models/lingbot_video/ + + diffusion_hunyuan_image3_function: + - *hunyuan_image3 + - tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py + + diffusion_hunyuan_image3_accuracy: + - *hunyuan_image3 + - tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py + + diffusion_hunyuan_image3_perf: + - *hunyuan_image3 + - tests/dfx/perf/scripts/run_diffusion_benchmark.py + - tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json + - tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json + - tests/dfx/perf/tests/test_hunyuan_image_tp4.json + + diffusion_hunyuan_image3_reliability: + - *hunyuan_image3 + - tests/dfx/reliability/test_reliability_hunyuan_image.py + + diffusion_minimax_h3_function: + - *minimax_h3 + - tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py + - tests/e2e/online_serving/minimax_h3/ + - tests/diffusion/models/minimax_h3/ + - vllm_omni/diffusion/lora/ + - vllm_omni/diffusion/offloader/ + - vllm_omni/diffusion/distributed/ + - vllm_omni/entrypoints/openai/ + + diffusion_minimax_h3_accuracy: + - *minimax_h3 + - tests/e2e/accuracy/minimax_h3/ + + diffusion_minimax_h3_perf: + - *minimax_h3 + - tests/dfx/perf/scripts/run_diffusion_benchmark.py + - tests/dfx/perf/tests/test_minimax_h3_vllm_omni.json + + diffusion_hunyuan_video_function: + - *hunyuan_video + - tests/e2e/online_serving/test_hunyuan_video_15_expansion.py + + diffusion_hunyuan_video_accuracy: + - *hunyuan_video + - tests/e2e/accuracy/hunyuanvideo15_t2v/ + - tests/e2e/accuracy/hunyuanvideo15_i2v/ + + diffusion_hunyuan_video_perf: + - *hunyuan_video + - tests/dfx/perf/scripts/run_diffusion_benchmark.py + - tests/dfx/perf/tests/test_hunyuanvideo15_t2v_vllm_omni.json + - tests/dfx/perf/tests/test_hunyuanvideo15_i2v_vllm_omni.json + + diffusion_z_image_function: + - *z_image + - tests/e2e/online_serving/test_z_image.py + - tests/e2e/online_serving/test_zimage_expansion.py + + diffusion_stable_audio_function: + - *stable_audio + - tests/e2e/offline_inference/test_stable_audio_expansion.py + - tests/e2e/online_serving/test_stable_audio_online_expansion.py + + diffusion_lingbot_function: + - *lingbot_video + - *lingbot_world + - tests/e2e/online_serving/test_lingbot_video.py + - tests/e2e/online_serving/test_lingbot_video_moe.py + - tests/e2e/offline_inference/test_lingbot_world_v2.py + + + # Invalid-param jobs: model business code only, plus the reliability scripts. + reliability_invalid_param_h100: + - *qwen3_omni + - *minicpmo_4_5 + - *qwen_image + - *wan22 + - tests/dfx/reliability/invalid_param_test/ + + reliability_invalid_param_l4: + - *qwen3_tts + - *stable_audio + - tests/dfx/reliability/invalid_param_test/ diff --git a/.buildkite/common/scripts/upload_pipeline.py b/.buildkite/common/scripts/upload_pipeline.py index 749f35f5a68..c8587d8c301 100644 --- a/.buildkite/common/scripts/upload_pipeline.py +++ b/.buildkite/common/scripts/upload_pipeline.py @@ -10,8 +10,15 @@ - Detect docs-only, pytest skip-mark-only, or combined skip-ci from git diff. - When only CI level YAML changes, enable **L2/L3** upload steps for affected levels only. -Test pipeline mode (e.g. test-merge.yml): - - Drop steps whose ``source_file_dependencies`` do not match changed files. +Test pipeline mode (e.g. test-merge.yml, test-nightly.yml, test-weekly.yml): + - Drop steps whose ``source_file_dependencies`` do not match changed files + (string/list keys from ci_source_file_dependencies.yml, or inline path + prefixes). Filtering applies on **PR label** uploads only; ``main`` + env + schedule (NIGHTLY/WEEKLY/merge push) keeps every job and still strips the field. + If no job-key prefix matches, a change to the pipeline YAML being uploaded + or to a path under the ``source_filter_fallback`` registry key keeps every + job so command, env, and hardware edits can be validated before merge. + When any job-key prefix already matches, normal filtering wins. - Expand uploader-only ``mirror_hardwares`` into ``agents`` (+ optional ``image`` for NPU) + ``plugins`` (see ci_mirror_hardwares.yml). - Omit ``mirror_hardwares`` to compose ``{chip}_{n}`` from pytest ``-m`` SKU @@ -87,6 +94,10 @@ } E2E_GROUP_MARKER = "E2E Test" CI_MIRROR_HARDWARES_PATH = ROOT / ".buildkite/common/ci_mirror_hardwares.yml" +CI_SOURCE_FILE_DEPENDENCIES_PATH = ROOT / ".buildkite/common/ci_source_file_dependencies.yml" +# Registry key whose prefixes bypass source filtering. Paths live in +# ci_source_file_dependencies.yml; do not attach this key to a job. +SOURCE_FILTER_FALLBACK_KEY = "source_filter_fallback" CUDA_HF_TOKEN_ENV = "VLLM_CI_HF_TOKEN" CUDA_HF_TOKEN_EXPORT = f'if [ -n "$${{{CUDA_HF_TOKEN_ENV}:-}}" ]; then export HF_TOKEN="$${{{CUDA_HF_TOKEN_ENV}}}"; fi' @@ -95,14 +106,7 @@ # ``*_UPLOAD_IF``: full gate for uploading that child pipeline. NIGHTLY_MAIN_IF = 'build.branch == "main" && build.env("NIGHTLY") == "1"' NIGHTLY_LABEL_IF = ( - f"({NIGHTLY_MAIN_IF}) || " - '(build.branch != "main" && (' - 'build.pull_request.labels includes "nightly-test" || ' - 'build.pull_request.labels includes "omni-test" || ' - 'build.pull_request.labels includes "tts-test" || ' - 'build.pull_request.labels includes "diffusion-x2iat-test" || ' - 'build.pull_request.labels includes "diffusion-x2v-test"' - "))" + f'({NIGHTLY_MAIN_IF}) || (build.branch != "main" && build.pull_request.labels includes "nightly-test")' ) WEEKLY_E2E_IF = 'build.branch == "main" && build.env("WEEKLY") == "1"' WEEKLY_MAIN_IF = 'build.branch == "main" && (build.env("WEEKLY") == "1" || build.env("NON_CRITICAL") == "1")' @@ -248,7 +252,7 @@ def _render_bootstrap_pipeline( return yaml.safe_dump(doc, sort_keys=False) -# --- Test pipeline (test-ready.yml, test-merge.yml) --- +# --- Test pipeline (test-ready.yml, test-merge.yml, test-nightly.yml) --- @lru_cache(maxsize=1) @@ -264,6 +268,117 @@ def _load_mirror_hardwares() -> dict[str, dict[str, Any]]: return presets +@lru_cache(maxsize=1) +def _load_source_file_dependencies() -> dict[str, list[str]]: + def flatten(value: Any, *, key: str) -> list[str]: + """Flatten YAML-anchor nested lists into a de-duplicated prefix list.""" + if isinstance(value, str): + if not value.strip(): + raise ValueError(f"empty path in source_file_dependencies[{key!r}]") + return [value] + if isinstance(value, list): + flattened: list[str] = [] + seen: set[str] = set() + for item in value: + for path in flatten(item, key=key): + if path not in seen: + seen.add(path) + flattened.append(path) + return flattened + raise ValueError( + f"source_file_dependencies[{key!r}] must be a list of path prefixes, got {type(value).__name__}", + ) + + if not CI_SOURCE_FILE_DEPENDENCIES_PATH.is_file(): + raise FileNotFoundError( + f"missing CI source_file_dependencies registry: {CI_SOURCE_FILE_DEPENDENCIES_PATH}", + ) + doc = yaml.safe_load(CI_SOURCE_FILE_DEPENDENCIES_PATH.read_text(encoding="utf-8")) + if not isinstance(doc, dict): + raise ValueError(f"invalid CI source_file_dependencies registry: {CI_SOURCE_FILE_DEPENDENCIES_PATH}") + presets = doc.get("source_file_dependencies") + if not isinstance(presets, dict): + raise ValueError(f"source_file_dependencies must be a mapping in {CI_SOURCE_FILE_DEPENDENCIES_PATH}") + loaded: dict[str, list[str]] = {} + for key, value in presets.items(): + if not isinstance(key, str) or not key.strip(): + raise ValueError( + f"source_file_dependencies keys must be non-empty strings in {CI_SOURCE_FILE_DEPENDENCIES_PATH}", + ) + loaded[key] = flatten(value, key=key) + return loaded + + +_REGISTRY_KEY_RE = re.compile(r"^[a-z][a-z0-9_]*$") + + +def _parse_command_text(step: dict[str, Any]) -> str: + """Parse step ``commands`` / ``command`` into a single string.""" + raw = step.get("commands", step.get("command")) + if raw is None: + return "" + if isinstance(raw, str): + return raw + if isinstance(raw, list): + return "\n".join(str(item) for item in raw if item is not None) + return str(raw) + + +def _resolve_source_file_dependencies(step: dict[str, Any]) -> list[str] | None: + """Expand a registry key (or list of keys / inline path prefixes).""" + deps = step.get("source_file_dependencies") + if deps is None: + return None + + def dedupe(*groups: list[str]) -> list[str]: + merged: list[str] = [] + seen: set[str] = set() + for group in groups: + for path in group: + if path not in seen: + seen.add(path) + merged.append(path) + return merged + + def lookup(key: str) -> list[str]: + if key == SOURCE_FILTER_FALLBACK_KEY: + raise ValueError( + f"source_file_dependencies {key!r} bypasses filtering and cannot be attached to a step " + f"({_get_step_label(step)!r})", + ) + registry = _load_source_file_dependencies() + paths = registry.get(key) + if paths is None: + known = ", ".join(sorted(registry)) + raise ValueError( + f"unknown source_file_dependencies {key!r} in step {_get_step_label(step)!r}; known: {known}", + ) + return list(paths) + + prefixes: list[str] + if isinstance(deps, str): + prefixes = lookup(deps) if _REGISTRY_KEY_RE.fullmatch(deps) else [deps] + elif isinstance(deps, list): + if not all(isinstance(item, str) for item in deps): + raise ValueError( + f"source_file_dependencies must be a string key or list of strings in step {_get_step_label(step)!r}", + ) + keyish = [bool(_REGISTRY_KEY_RE.fullmatch(item)) for item in deps] + if deps and all(keyish): + prefixes = dedupe(*(lookup(key) for key in deps)) + elif any(keyish): + raise ValueError( + f"source_file_dependencies in step {_get_step_label(step)!r} mixes registry keys and path prefixes", + ) + else: + prefixes = list(deps) + else: + raise ValueError( + f"source_file_dependencies must be a string key or list in step {_get_step_label(step)!r}", + ) + return prefixes + + @lru_cache(maxsize=1) def _cuda_mirror_chips() -> tuple[str, ...]: """Lowercase CUDA SKUs that have at least one ``{chip}_{n}`` preset. @@ -278,7 +393,7 @@ def _cuda_mirror_chips() -> tuple[str, ...]: return tuple(sorted(chips, key=lambda chip: (-len(chip), chip))) -def _read_cards_marks(expr: str) -> tuple[set[int], bool]: +def _parse_cards_marks(expr: str) -> tuple[set[int], bool]: """Parse registered ``cards_n`` / ``not cards_n`` from a pytest ``-m`` expr. Returns positive card counts and whether any ``not cards_*`` is present. @@ -295,8 +410,8 @@ def _read_cards_marks(expr: str) -> tuple[set[int], bool]: return positives, has_not_cards -def _read_hardware_marks(expr: str) -> set[str]: - """Return lowercase positive CUDA SKUs that have a mirror preset. +def _parse_hardware_marks(expr: str) -> set[str]: + """Parse lowercase positive CUDA SKUs that have a mirror preset. ``not H100`` is ignored. SKUs with no ``{chip}_*`` preset (e.g. H200) are dropped. """ @@ -338,36 +453,22 @@ def _get_mirror_hw_selector() -> str: return selector -def _read_pytest_marks(commands: Any) -> tuple[set[str], int | Literal["max"] | None]: - """Read pytest ``-m`` from step commands: CUDA SKU chips and cards count. +def _parse_pytest_marks(step: dict[str, Any]) -> tuple[set[str], int | Literal["max"] | None]: + """Parse pytest ``-m`` from step commands: CUDA SKU chips and cards count. *chips* are lowercase SKUs from positive markers (``not H100`` is ignored). *cards* is the max positive ``cards_n``, ``"max"`` when only ``not cards_*`` is present, or ``None``. """ - chunks: list[str] = [] - - def _collect(value: Any) -> None: - if value is None: - return - if isinstance(value, str): - chunks.append(value) - elif isinstance(value, list): - for part in value: - _collect(part) - else: - chunks.append(str(value)) - - _collect(commands) - text = "\n".join(chunks) + text = _parse_command_text(step) chips: set[str] = set() positives: set[int] = set() has_not_cards = False for match in _PYTEST_MARKER_ARG.finditer(text): expr = match.group(1) or match.group(2) or match.group(3) or "" - chips |= _read_hardware_marks(expr) - card_counts, not_cards = _read_cards_marks(expr) + chips |= _parse_hardware_marks(expr) + card_counts, not_cards = _parse_cards_marks(expr) positives.update(card_counts) has_not_cards = has_not_cards or not_cards @@ -484,7 +585,7 @@ def _expand_mirror_hardwares(step: dict[str, Any]) -> dict[str, Any] | None: if "mirror_hardwares" not in step: if has_pool: return step - chips, cards = _read_pytest_marks(step.get("commands")) + chips, cards = _parse_pytest_marks(step) if not chips and cards is None: return step preset_name = _compose_mirror_hardware_name(chips, cards, step_label=step_label) @@ -552,11 +653,7 @@ def _process_test_steps( processed.append(step) continue - deps = step.get("source_file_dependencies") - if deps is not None and not isinstance(deps, list): - raise ValueError( - f"source_file_dependencies must be a list in step {_get_step_label(step)!r}", - ) + deps = _resolve_source_file_dependencies(step) if changed_files is not None and deps is not None and not _match_source_file(changed_files, deps): _log(f"skip {_get_step_label(step)!r} (no changes under {deps})") continue @@ -595,10 +692,25 @@ def _select_e2e_group_steps(steps: list[Any]) -> list[Any]: return selected +def _any_source_dependency_match(steps: list[Any], changed_files: list[str]) -> bool: + """True when any step with ``source_file_dependencies`` matches *changed_files*.""" + for step in steps: + if not isinstance(step, dict): + continue + deps = _resolve_source_file_dependencies(step) + if deps is not None and _match_source_file(changed_files, deps): + return True + nested = step.get("steps") + if isinstance(nested, list) and _any_source_dependency_match(nested, changed_files): + return True + return False + + def _render_test_pipeline( doc: dict[str, Any], changed_files: list[str] | None, *, + pipeline_path: Path | None = None, e2e_only: bool = False, ) -> dict[str, Any]: """Filter steps by PR diff and strip uploader-only ``source_file_dependencies`` metadata.""" @@ -611,6 +723,13 @@ def _render_test_pipeline( return doc if e2e_only: steps = _select_e2e_group_steps(steps) + # Bypass is a fallback: only when no job-key prefix matched. + if changed_files is not None and pipeline_path is not None: + if not _any_source_dependency_match(steps, changed_files): + bypass = _source_filter_fallback_reason(changed_files, pipeline_path) + if bypass is not None: + _log(f"keep all jobs (no source_file_dependencies match; bypassed by {bypass})") + changed_files = None steps = _process_test_steps(steps, changed_files) return {**doc, "steps": steps} @@ -618,6 +737,63 @@ def _render_test_pipeline( # --- Entry (read file → bootstrap or test render → YAML string) --- +def _repo_relative_posix(path: Path) -> str: + resolved = path.resolve() + try: + return resolved.relative_to(ROOT).as_posix() + except ValueError: + return path.as_posix().replace("\\", "/") + + +def _source_filter_fallback_prefixes() -> list[str]: + """Shared uploader/registry prefixes from ``source_filter_fallback``. Not a job key.""" + registry = _load_source_file_dependencies() + prefixes = registry.get(SOURCE_FILTER_FALLBACK_KEY) + if not prefixes: + raise ValueError( + f"source_file_dependencies[{SOURCE_FILTER_FALLBACK_KEY!r}] must list shared uploader paths " + f"in {CI_SOURCE_FILE_DEPENDENCIES_PATH}", + ) + return prefixes + + +def _source_filter_fallback_reason(changed_files: list[str] | None, pipeline_path: Path) -> str | None: + """Return the changed path that should disable source filtering, if any.""" + if not changed_files: + return None + changed = set(changed_files) + pipeline_rel = _repo_relative_posix(pipeline_path) + if pipeline_rel in changed: + return pipeline_rel + prefixes = _source_filter_fallback_prefixes() + for path in changed_files: + if _match_source_file([path], prefixes): + return path + return None + + +def _changed_files_for_source_filter( + ctx, + *, + force_all: bool, + e2e_only: bool, +) -> list[str] | None: + """Return the diff to filter against, or None to keep every step. + + ``source_file_dependencies`` is label-only: scheduled ``main`` uploads + (NIGHTLY/WEEKLY/post-merge) run the full pipeline. ``--all`` / ``--e2e`` + also disable filtering. Pipeline YAML / ``source_filter_fallback`` matches + are applied later in ``_render_test_pipeline`` only when no job-key prefix + already matched. + """ + if force_all or e2e_only: + return None + if os.environ.get("BUILDKITE_BRANCH", "") == "main": + _log("main branch: keep all jobs (source_file_dependencies is label-only)") + return None + return ctx.changed_files + + def _render_pipeline( path: Path, *, @@ -635,13 +811,18 @@ def _render_pipeline( text = path.read_text(encoding="utf-8") ctx = resolve_ci_context_from_git() - changed_files = None if force_all or e2e_only else ctx.changed_files + changed_files = _changed_files_for_source_filter(ctx, force_all=force_all, e2e_only=e2e_only) doc = yaml.safe_load(text) if not isinstance(doc, dict): raise ValueError(f"invalid pipeline YAML: {path}") - doc = _render_test_pipeline(doc, changed_files, e2e_only=e2e_only) + doc = _render_test_pipeline( + doc, + changed_files, + pipeline_path=path, + e2e_only=e2e_only, + ) return yaml.safe_dump(doc, sort_keys=False) diff --git a/.buildkite/cuda/bootstrap-upload-steps.yml b/.buildkite/cuda/bootstrap-upload-steps.yml index ab4254feb98..502a506e6d6 100644 --- a/.buildkite/cuda/bootstrap-upload-steps.yml +++ b/.buildkite/cuda/bootstrap-upload-steps.yml @@ -22,34 +22,48 @@ steps: agents: queue: "cpu_queue_premerge" - # L3 Test — PR merge-test label (diff-aware) or main+WEEKLY=1 (--e2e, full E2E group) + # L3 Test — PR merge-test label (diff-aware) or main (full pipeline / WEEKLY --e2e) - label: "Upload Merge Pipeline" depends_on: image-build key: upload-merge-pipeline commands: - | EXTRA="" - if [ "${WEEKLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then - EXTRA="--e2e" + if [ "${BUILDKITE_BRANCH}" = "main" ]; then + if [ "${WEEKLY:-}" = "1" ]; then + EXTRA="--e2e" + else + EXTRA="--all" + fi fi python3 .buildkite/common/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/cuda/test-merge.yml agents: queue: "cpu_queue_premerge" - # L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild) + # L4 Test — main+NIGHTLY=1 (full pipeline), or PR nightly-test label (diff-aware) - label: "Upload Nightly Pipeline" depends_on: image-build key: upload-nightly-pipeline commands: - - python3 .buildkite/common/scripts/upload_pipeline.py --upload .buildkite/cuda/test-nightly.yml + - | + EXTRA="" + if [ "${NIGHTLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then + EXTRA="--all" + fi + python3 .buildkite/common/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/cuda/test-nightly.yml agents: queue: "cpu_queue_premerge" - # L5 Test — main+WEEKLY=1 / NON_CRITICAL=1 (scheduled), or PR with weekly-test label + # L5 Test — main+WEEKLY=1 / NON_CRITICAL=1 (full pipeline), or PR weekly-test label (diff-aware) - label: "Upload Weekly Pipeline" depends_on: image-build key: upload-weekly-pipeline commands: - - python3 .buildkite/common/scripts/upload_pipeline.py --upload .buildkite/cuda/test-weekly.yml + - | + EXTRA="" + if [ "${BUILDKITE_BRANCH}" = "main" ]; then + EXTRA="--all" + fi + python3 .buildkite/common/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/cuda/test-weekly.yml agents: queue: "cpu_queue_premerge" diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index 5aa49ea6cb1..1a458bccf68 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.buildkite/cuda/test-merge.yml @@ -78,21 +78,7 @@ steps: depends_on: upload-merge-pipeline steps: - label: "TTS · Qwen3-TTS CustomVoice Test" - source_file_dependencies: - - tests/e2e/features/prefill_phase/ - - vllm_omni/worker/gpu_model_runner.py - - tests/helpers/media.py - - tests/e2e/online_serving/test_qwen3_tts_customvoice.py - - tests/e2e/offline_inference/test_qwen3_tts_customvoice.py - - vllm_omni/model_executor/models/qwen3_tts/ - - vllm_omni/model_executor/models/common/qwen3_code_predictor.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 - - vllm_omni/deploy/qwen3_tts.yaml - - vllm_omni/entrypoints/openai/serving_speech.py - - vllm_omni/entrypoints/openai/tts_adapters/ - - vllm_omni/engine/stage_pool.py + source_file_dependencies: tts_qwen3_tts_function commands: - | timeout 40m bash -c " @@ -103,21 +89,7 @@ steps: mirror_hardwares: l4_1 - label: "TTS · Qwen3-TTS Base Test" - source_file_dependencies: - - tests/helpers/media.py - - tests/e2e/online_serving/test_qwen3_tts_base.py - - tests/e2e/offline_inference/test_qwen3_tts_base.py - - vllm_omni/model_executor/models/qwen3_tts/ - - vllm_omni/model_executor/models/common/qwen3_code_predictor.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 - - vllm_omni/deploy/qwen3_tts.yaml - - vllm_omni/entrypoints/openai/serving_speech.py - - vllm_omni/entrypoints/openai/tts_adapters/ - - vllm_omni/engine/stage_pool.py - - .buildkite/common/scripts/run_cov_split.sh - - pyproject.toml + source_file_dependencies: tts_qwen3_tts_cov # 40 -> 60: the coverage split runs the model twice. timeout_in_minutes: 60 commands: @@ -132,36 +104,14 @@ steps: mirror_hardwares: l4_1 - label: "Omni · Qwen3-Omni Test" - source_file_dependencies: - - tests/e2e/offline_inference/test_qwen3_omni.py - - tests/e2e/offline_inference/test_qwen3_omni_colocate_async.py - - tests/e2e/online_serving/test_qwen3_omni.py - - tests/helpers/runtime.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 - - vllm_omni/deploy/qwen3_omni_moe_thinking.yaml + source_file_dependencies: omni_qwen3_omni_function timeout_in_minutes: 50 commands: - 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 · MiniCPM-o 4.5 Test" - source_file_dependencies: - - tests/e2e/offline_inference/test_minicpmo_4_5.py - - tests/e2e/online_serving/test_minicpmo_4_5.py - - vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja - - vllm_omni/model_executor/models/minicpmo_4_5/ - - vllm_omni/model_executor/models/step_audio2/ - - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ - - vllm_omni/model_executor/models/cosyvoice3/utils.py - - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py - - vllm_omni/deploy/minicpmo_4_5.yaml + source_file_dependencies: omni_minicpmo_4_5_function timeout_in_minutes: 50 commands: - pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5.py -m 'advanced_model and cuda' --run-level 'advanced_model' @@ -174,59 +124,14 @@ steps: mirror_hardwares: h100_2 - label: "Omni · MiniCPM-o 4.5 Duplex Test" - source_file_dependencies: - - tests/assets/minicpmo_4_5/response_required_16k.wav - - tests/assets/minicpmo_4_5/soft_interrupt_16k.wav - - tests/e2e/online_serving/test_minicpmo_4_5_duplex.py - - tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py - - tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py - - tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py - - tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py - - tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py - - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py - - vllm_omni/deploy/minicpmo_4_5.yaml - - tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py - - vllm_omni/engine/duplex/ - - vllm_omni/entrypoints/duplex/ - # Public duplex client library + its unit suites: client-side - # changes must re-run the duplex e2e group. The live client probe - # (test_duplex_client_live.py, single GPU, core tier) runs in the - # ready pipeline's duplex group. - - vllm_omni/clients/ - - tests/clients/ - - vllm_omni/entrypoints/duplex_request_client.py - - vllm_omni/outputs/duplex.py - - vllm_omni/model_executor/duplex_sampling.py - - tests/engine/duplex/ - - tests/entrypoints/duplex/ - - tests/model_executor/models/minicpmo_4_5/duplex/ - # Legacy Realtime probe client + video stacking used by the e2e drivers. - - vllm_omni/experimental/fullduplex/ - - vllm_omni/model_executor/models/minicpmo_4_5/ - - vllm_omni/model_executor/models/step_audio2/ - - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ - - vllm_omni/model_executor/models/cosyvoice3/utils.py - - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py - - vllm_omni/core/sched/ - - vllm_omni/outputs/multimodal_accumulation.py - - vllm_omni/engine/stage_init_utils.py - - vllm_omni/engine/async_omni_engine.py - - vllm_omni/engine/orchestrator.py - - vllm_omni/entrypoints/cli/serve.py - - vllm_omni/entrypoints/openai/api_server.py - - vllm_omni/entrypoints/openai/duplex_capability.py - - vllm_omni/config/stage_config.py + source_file_dependencies: omni_minicpmo_4_5_duplex_function timeout_in_minutes: 50 commands: - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'advanced_model and cuda' --run-level 'advanced_model' mirror_hardwares: h100_1 - label: "Diffusion · Qwen Image Test" - source_file_dependencies: - - tests/e2e/online_serving/test_qwen_image.py - - vllm_omni/diffusion/models/qwen_image/ - - vllm_omni/diffusion/models/dmd2/ - - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py + source_file_dependencies: diffusion_qwen_image_function commands: - | timeout 40m bash -c ' @@ -236,11 +141,7 @@ steps: mirror_hardwares: h100_1 - label: "Diffusion · Wan22 Test" - source_file_dependencies: - - tests/e2e/offline_inference/test_wan22_t2v.py - - tests/e2e/online_serving/test_wan22_t2v.py - - vllm_omni/diffusion/models/wan2_2/ - - vllm_omni/diffusion/models/dmd2/ + source_file_dependencies: diffusion_wan22_function commands: - | timeout 40m bash -c " @@ -250,9 +151,7 @@ steps: mirror_hardwares: h100_1 - label: "Diffusion · HunyuanVideo-1.5 Test" - source_file_dependencies: - - tests/e2e/online_serving/test_hunyuan_video_15_expansion.py - - vllm_omni/diffusion/models/hunyuan_video/ + source_file_dependencies: diffusion_hunyuan_video_function commands: - | timeout 40m bash -c " @@ -262,17 +161,7 @@ steps: mirror_hardwares: h100_1 - label: "Diffusion · MiniMax-H3 L3 DLO DP2 and FastH3" - source_file_dependencies: - - tests/e2e/online_serving/minimax_h3/ - - tests/helpers/assertions.py - - tests/helpers/mark.py - - tests/helpers/media.py - - tests/helpers/runtime.py - - vllm_omni/diffusion/models/minimax_h3/ - - vllm_omni/diffusion/lora/ - - vllm_omni/diffusion/offloader/ - - vllm_omni/diffusion/distributed/ - - vllm_omni/entrypoints/openai/ + source_file_dependencies: diffusion_minimax_h3_function # Keep four server lifecycles in one Buildkite job, but never in one # pytest process. DLO uses multiprocessing shared-memory queues; the # parent's resource_tracker can retain a stale /psm_* entry during diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index c85e9abc10f..ee4fc6f3ea4 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -12,26 +12,34 @@ steps: - group: ":card_index_dividers: Omni Model Test" key: nightly-omni-test-group depends_on: upload-nightly-pipeline - if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "omni-test" + if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" steps: - label: ":full_moon: Omni · Function Test with H100 · Single-GPU" + source_file_dependencies: + - omni_qwen3_omni_function + - omni_minicpmo_4_5_function timeout_in_minutes: 120 commands: - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and B200 and omni and cards_1" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py - label: ":full_moon: Omni · Function Test with H100 · 2-GPU" + source_file_dependencies: + - omni_qwen3_omni_function + - omni_minicpmo_4_5_function timeout_in_minutes: 90 commands: - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and B200 and omni and cards_2" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py - label: ":full_moon: Omni · MiniCPM-o 4.5 Duplex Test" + source_file_dependencies: omni_minicpmo_4_5_duplex_function timeout_in_minutes: 50 commands: - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py -m "full_model and cuda and H100 and omni and cards_1" --run-level "full_model" mirror_hardwares: h100_1 - label: ":full_moon: Omni · Doc Test with L4 · 4-GPU" + source_file_dependencies: omni_qwen2_5_omni_doc timeout_in_minutes: 90 commands: - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" @@ -39,12 +47,14 @@ steps: mirror_hardwares: l4_4 - label: ":full_moon: Omni · Doc Test with H100 · 2-GPU" + source_file_dependencies: omni_qwen3_omni_doc timeout_in_minutes: 90 commands: - pytest -s -v tests/examples/ -m "full_model and omni and H100 and cards_2" --run-level "full_model" mirror_hardwares: h100_2 - label: ":full_moon: Omni · Accuracy Test" + source_file_dependencies: omni_qwen3_omni_accuracy timeout_in_minutes: 180 commands: - export SEED_TTS_WER_EVAL=1 @@ -58,6 +68,7 @@ steps: mirror_hardwares: h100_2 - label: ":full_moon: Omni · MiniCPM-o 4.5 · Accuracy Test" + source_file_dependencies: omni_minicpmo_4_5_accuracy timeout_in_minutes: 180 commands: - export SEED_TTS_WER_EVAL=1 @@ -71,6 +82,7 @@ steps: mirror_hardwares: h100_2 - label: ":full_moon: Omni · Perf Test · No Async Chunk" + source_file_dependencies: omni_qwen3_omni_perf key: nightly-omni-performance-no-async-chunk timeout_in_minutes: 300 commands: @@ -84,6 +96,7 @@ steps: mirror_hardwares: h100_2 - label: ":full_moon: Omni · Perf Test · Async Chunk" + source_file_dependencies: omni_qwen3_omni_perf key: nightly-omni-performance-async-chunk timeout_in_minutes: 180 commands: @@ -97,19 +110,21 @@ steps: mirror_hardwares: h100_2 - label: ":full_moon: Omni · MiniCPM-o 4.5 · Perf Test" + source_file_dependencies: omni_minicpmo_4_5_perf key: nightly-omni-performance-minicpmo-4-5 timeout_in_minutes: 180 commands: - export BENCHMARK_DIR=tests/dfx/perf/results - | set +e - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5.json + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5.json -m "H100 and cards_1" EXIT=$$? buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT mirror_hardwares: h100_1 - label: ":full_moon: Omni · MiniCPM-o 4.5 · Duplex Seed-TTS Perf Test" + source_file_dependencies: omni_minicpmo_4_5_duplex_perf key: nightly-omni-performance-minicpmo-4-5-duplex-seed-tts timeout_in_minutes: 180 commands: @@ -123,6 +138,7 @@ steps: mirror_hardwares: h100_1 - label: ":full_moon: Omni · Multi-Replica Startup Test with 4x H100" + source_file_dependencies: omni_qwen3_omni_function timeout_in_minutes: 45 commands: - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py -m "full_model" --run-level "core_model" @@ -131,10 +147,11 @@ steps: - group: ":card_index_dividers: TTS Model Test" key: nightly-tts-test-group depends_on: upload-nightly-pipeline - if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test" + if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" steps: - label: ":full_moon: TTS · Function Test with L4" + source_file_dependencies: tts_qwen3_tts_function timeout_in_minutes: 120 commands: - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" @@ -144,6 +161,7 @@ steps: mirror_hardwares: l4_1 - label: ":full_moon: TTS · Perf Test · Single-GPU" + source_file_dependencies: tts_qwen3_tts_perf key: nightly-tts-performance-single timeout_in_minutes: 180 commands: @@ -158,6 +176,7 @@ steps: mirror_hardwares: h100_1 - label: ":full_moon: TTS · Perf Test · 2-GPU" + source_file_dependencies: tts_qwen3_tts_perf key: nightly-tts-performance-2gpu timeout_in_minutes: 180 commands: @@ -180,6 +199,7 @@ steps: steps: - label: ":full_moon: Tiny Model Tests · Diffusion (Multi-GPU accelerations) · 2-GPU" + source_file_dependencies: diffusion_tiny_model timeout_in_minutes: 30 # NOTE: Single-GPU (core_model) tiny model tests run on the ready label, # so nightly only adds the multi-GPU parallelism configurations. @@ -193,6 +213,7 @@ steps: mirror_hardwares: l4_2 - label: ":full_moon: Tiny Model Tests · Diffusion (Multi-GPU accelerations) · 4-GPU" + source_file_dependencies: diffusion_tiny_model timeout_in_minutes: 30 # CFG+TP extra_test_groups request 2×2=4 devices (see get_required_device_count). commands: @@ -204,11 +225,11 @@ steps: depends_on: upload-nightly-pipeline if: >- build.env("NIGHTLY") == "1" || - build.pull_request.labels includes "nightly-test" || - build.pull_request.labels includes "diffusion-x2iat-test" + build.pull_request.labels includes "nightly-test" steps: - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Single-GPU" + source_file_dependencies: diffusion_qwen_image_function timeout_in_minutes: 120 commands: - >- @@ -219,6 +240,7 @@ steps: mirror_hardwares: h100_1 - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · 2-GPU" + source_file_dependencies: diffusion_qwen_image_function timeout_in_minutes: 120 commands: - >- @@ -228,30 +250,26 @@ steps: --run-level "full_model" mirror_hardwares: h100_2 - - label: ":full_moon: Diffusion X2I(&A&T) · Boogu Edit CFG · 3-GPU" + - label: ":full_moon: Diffusion X2I(&A&T) · Z-Image Function Test · L4 · 4-GPU" + source_file_dependencies: diffusion_z_image_function timeout_in_minutes: 120 commands: - >- pytest -sv - tests/e2e/online_serving/test_boogu_image_edit_expansion.py - -k "test_double_guidance_cfg_parallel and cfg_parallel_3_full" - -m "slow and diffusion and cuda and H100 and cards_3" + tests/e2e/online_serving/test_zimage_expansion.py + -m "full_model and diffusion and L4 and cards_4" --run-level "full_model" - mirror_hardwares: h100_3 - - - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4 · 4-GPU" - timeout_in_minutes: 120 - commands: - - pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4 and cards_4" --run-level "full_model" --ignore=tests/e2e/accuracy mirror_hardwares: l4_4 - label: ":full_moon: Diffusion X2I(&A&T) · Doc Test" + source_file_dependencies: diffusion_text_to_image_doc timeout_in_minutes: 60 commands: - pytest -s -v tests/examples/*/test_text_to_image.py -m "full_model and example and H100" --run-level "full_model" mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2I(&A&T) · Accuracy Test" + source_file_dependencies: diffusion_qwen_image_accuracy timeout_in_minutes: 180 commands: - export VLLM_HTTP_TIMEOUT_KEEP_ALIVE=120 @@ -260,6 +278,7 @@ steps: mirror_hardwares: h100_1 - label: ":full_moon: HunyuanImage3-DIT · Accuracy Test" + source_file_dependencies: diffusion_hunyuan_image3_accuracy timeout_in_minutes: 180 commands: - export HUNYUAN_IMAGE3_MODEL="tencent/HunyuanImage-3.0-Instruct" @@ -267,13 +286,8 @@ steps: - pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m full_model --run-level full_model mirror_hardwares: h100_4 - - label: ":full_moon: HiDream-O1-Image · Accuracy Test" - timeout_in_minutes: 180 - commands: - - pytest -s -v tests/e2e/accuracy/test_hidream_o1_image.py -m full_model --run-level full_model - mirror_hardwares: h100_2 - - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image · Single-GPU" + source_file_dependencies: diffusion_qwen_image_perf key: nightly-diffusion-x2iat-performance-qwen-image-single timeout_in_minutes: 180 commands: @@ -289,6 +303,7 @@ steps: mirror_hardwares: h100_1 - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image · 4-GPU" + source_file_dependencies: diffusion_qwen_image_perf key: nightly-diffusion-x2iat-performance-qwen-image-4gpu timeout_in_minutes: 180 commands: @@ -309,59 +324,73 @@ steps: depends_on: upload-nightly-pipeline if: >- build.env("NIGHTLY") == "1" || - build.pull_request.labels includes "nightly-test" || - build.pull_request.labels includes "diffusion-x2v-test" + build.pull_request.labels includes "nightly-test" steps: - label: ":full_moon: Diffusion X2V · Wan2.2 T2V Function Test · Single-GPU" + source_file_dependencies: diffusion_wan22_function timeout_in_minutes: 90 commands: - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "t2v" -m "full_model and cuda and H100 and cards_1" --run-level "full_model" mirror_hardwares: h100_1 - label: ":full_moon: Diffusion X2V · Wan2.2 T2V Function Test · 2-GPU" + source_file_dependencies: diffusion_wan22_function timeout_in_minutes: 90 commands: - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "t2v" -m "full_model and cuda and H100 and cards_2" --run-level "full_model" mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2V · Wan2.2 I2V/TI2V Function Test · Single-GPU" + source_file_dependencies: diffusion_wan22_function timeout_in_minutes: 90 commands: - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda and H100 and cards_1" --run-level "full_model" mirror_hardwares: h100_1 - label: ":full_moon: Diffusion X2V · Wan2.2 I2V/TI2V Function Test · 2-GPU" + source_file_dependencies: diffusion_wan22_function timeout_in_minutes: 90 commands: - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda and H100 and cards_2" --run-level "full_model" mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2V · HunyuanVideo-1.5 Function Test" + source_file_dependencies: diffusion_hunyuan_video_function timeout_in_minutes: 90 commands: - pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m "full_model and cuda" --run-level "full_model" mirror_hardwares: h100_2 - - label: ":full_moon: Diffusion X2V · MAGI-2 Preview Serving Smoke" - timeout_in_minutes: 180 + - label: ":full_moon: Diffusion X2V · LingBot Function Test" + source_file_dependencies: diffusion_lingbot_function + timeout_in_minutes: 120 commands: - - pytest -s -v tests/e2e/online_serving/test_magi2.py -m "advanced_model and diffusion and cuda and H100 and cards_4" --run-level "advanced_model" - mirror_hardwares: h100_4 + - >- + pytest -s -v + tests/e2e/online_serving/test_lingbot_video.py + tests/e2e/online_serving/test_lingbot_video_moe.py + tests/e2e/offline_inference/test_lingbot_world_v2.py + -m "full_model and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 - label: ":full_moon: Diffusion X2V · Doc Test" + source_file_dependencies: diffusion_image_to_video_doc timeout_in_minutes: 60 commands: - pytest -s -v tests/examples/offline_inference/test_image_to_video.py -m "full_model and example and H100" --run-level "full_model" mirror_hardwares: h100_1 - - label: ":full_moon: Diffusion X2V · Other Function Test" + - label: ":full_moon: Diffusion X2V · Wan2.2 Autoround Function Test" + source_file_dependencies: diffusion_wan22_function timeout_in_minutes: 90 commands: - pytest -s -v tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model" mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test" + source_file_dependencies: diffusion_wan22_accuracy timeout_in_minutes: 180 commands: - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m full_model --run-level full_model @@ -369,12 +398,14 @@ steps: mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2V · MiniMax H3 I2VA/Ref2VA Accuracy Test" + source_file_dependencies: diffusion_minimax_h3_accuracy timeout_in_minutes: 180 commands: - pytest -s -v tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py -m full_model --run-level full_model mirror_hardwares: h100_4 - - label: ":full_moon: Diffusion X2V · Other Accuracy Test" + - label: ":full_moon: Diffusion X2V · HunyuanVideo-1.5 Accuracy Test" + source_file_dependencies: diffusion_hunyuan_video_accuracy timeout_in_minutes: 180 commands: - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model @@ -382,6 +413,7 @@ steps: mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2V · Perf Test · Single-GPU" + source_file_dependencies: diffusion_wan22_perf key: nightly-diffusion-x2v-performance-single timeout_in_minutes: 180 commands: @@ -396,6 +428,9 @@ steps: mirror_hardwares: h100_1 - label: ":full_moon: Diffusion X2V · Perf Test · 2-GPU" + source_file_dependencies: + - diffusion_wan22_perf + - diffusion_cosmos3_perf key: nightly-diffusion-x2v-performance-2gpu timeout_in_minutes: 180 commands: @@ -414,6 +449,7 @@ steps: mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2V · Perf Test · HunyuanVideo-1.5" + source_file_dependencies: diffusion_hunyuan_video_perf key: nightly-diffusion-x2v-performance-hunyuanvideo15 timeout_in_minutes: 180 commands: @@ -431,6 +467,7 @@ steps: mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2V · Perf Test · MiniMax-H3" + source_file_dependencies: diffusion_minimax_h3_perf key: nightly-diffusion-x2v-performance-minimax-h3 timeout_in_minutes: 180 commands: @@ -450,54 +487,63 @@ steps: if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" steps: - label: ":full_moon: Diffusion · Distributed Test with L4 · 2-GPU" + source_file_dependencies: diffusion_distributed_attention timeout_in_minutes: 60 commands: - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4 and cards_2' --run-level "full_model" mirror_hardwares: l4_2 - label: ":full_moon: Diffusion · Distributed Test with L4 · 3-GPU" + source_file_dependencies: diffusion_distributed_attention timeout_in_minutes: 60 commands: - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4 and cards_3' --run-level "full_model" mirror_hardwares: l4_3 - label: ":full_moon: Diffusion · Distributed Test with L4 · 4-GPU" + source_file_dependencies: diffusion_distributed_attention timeout_in_minutes: 60 commands: - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4 and cards_4' --run-level "full_model" mirror_hardwares: l4_4 - label: ":full_moon: Diffusion · Distributed Test with H100 · 2-GPU" + source_file_dependencies: diffusion_distributed_attention timeout_in_minutes: 60 commands: - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100 and cards_2' --run-level "full_model" mirror_hardwares: h100_2 - label: ":full_moon: Diffusion · Distributed Test with H100 · 4-GPU" + source_file_dependencies: diffusion_distributed_attention timeout_in_minutes: 60 commands: - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100 and cards_4' --run-level "full_model" mirror_hardwares: h100_4 - label: ":full_moon: Diffusion · Offloader Test" + source_file_dependencies: diffusion_offloader timeout_in_minutes: 60 commands: - pytest -sv tests/diffusion/offloader -m "full_model and cuda" mirror_hardwares: l4_1 - label: ":full_moon: Diffusion · Quantization Test with H100 · Single-GPU" + source_file_dependencies: diffusion_quantization timeout_in_minutes: 60 commands: - pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model" mirror_hardwares: h100_1 - label: ":full_moon: Diffusion · MiniMax H3 FP8 Quality Test with H100 · 2-GPU" + source_file_dependencies: diffusion_minimax_h3_function timeout_in_minutes: 60 commands: - pytest -sv tests/diffusion/models/minimax_h3/test_minimax_h3_quantization_quality.py -m 'full_model and cuda and H100' --run-level "full_model" mirror_hardwares: h100_2 - label: ":full_moon: Diffusion · Quantization Test with L4" + source_file_dependencies: diffusion_quantization timeout_in_minutes: 60 commands: - pip install "vllm-gguf-plugin==0.0.4" diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index 6d8d1e8b98a..67f03dbb680 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.buildkite/cuda/test-ready.yml @@ -54,16 +54,19 @@ steps: mirror_hardwares: l4_1 - label: "Diffusion · Distributed Test · Single-GPU" + source_file_dependencies: diffusion_distributed_attention commands: - timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda and L4 and cards_1' --run-level "core_model" mirror_hardwares: l4_1 - label: "Diffusion · Distributed Test · 2-GPU" + source_file_dependencies: diffusion_distributed_attention commands: - timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda and L4 and cards_2' --run-level "core_model" mirror_hardwares: l4_2 - label: "Diffusion · Distributed & Attention Test · 4-GPU" + source_file_dependencies: diffusion_distributed_attention commands: - timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda and L4 and cards_4' --run-level "core_model" - timeout 20m pytest -sv tests/diffusion/attention/test_attention_sp.py -m 'core_model and cuda and L4 and cards_4' --run-level "core_model" @@ -120,29 +123,14 @@ steps: depends_on: upload-ready-pipeline steps: - label: "Custom Pipeline Test" - source_file_dependencies: - - tests/e2e/features/custom_pipeline/ - - tests/e2e/features/helpers/custom_pipeline.py - - vllm_omni/diffusion/ - - vllm_omni/engine/ - - vllm_omni/entrypoints/ - - vllm_omni/inputs/ - - vllm_omni/outputs/ + source_file_dependencies: custom_pipeline commands: - timeout 20m pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model" mirror_hardwares: l4_1 - label: "RLHF · VeRL-Omni E2E Test" - source_file_dependencies: - - tests/e2e/features/rlhf_test/test_verl_omni_e2e.py - - tests/e2e/features/helpers/verl_omni_server.py - - tests/e2e/features/helpers/custom_pipeline.py - - vllm_omni/diffusion/ - - vllm_omni/engine/ - - vllm_omni/entrypoints/ - - vllm_omni/inputs/ - - vllm_omni/outputs/ + source_file_dependencies: rlhf_verl_omni commands: - | timeout 20m bash -c " @@ -151,17 +139,7 @@ steps: mirror_hardwares: l4_1 - label: "Omni · Qwen3-Omni Test" - source_file_dependencies: - - tests/e2e/online_serving/test_qwen3_omni.py - - tests/helpers/runtime.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 + source_file_dependencies: omni_qwen3_omni_function commands: - | timeout 20m bash -c " @@ -170,14 +148,7 @@ steps: mirror_hardwares: h100_2 - label: "Omni · MiniCPM-o 4.5 Offline Test" - source_file_dependencies: - - tests/e2e/offline_inference/test_minicpmo_4_5.py - - vllm_omni/model_executor/models/minicpmo_4_5/ - - vllm_omni/model_executor/models/step_audio2/ - - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ - - vllm_omni/model_executor/models/cosyvoice3/utils.py - - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py - - vllm_omni/deploy/minicpmo_4_5.yaml + source_file_dependencies: omni_minicpmo_4_5_function commands: - | timeout 20m bash -c " @@ -186,15 +157,7 @@ steps: mirror_hardwares: h100_1 - label: "Omni · MiniCPM-o 4.5 Online Test" - source_file_dependencies: - - tests/e2e/online_serving/test_minicpmo_4_5.py - - vllm_omni/transformers_utils/chat_templates/minicpmo45_native.jinja - - vllm_omni/model_executor/models/minicpmo_4_5/ - - vllm_omni/model_executor/models/step_audio2/ - - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ - - vllm_omni/model_executor/models/cosyvoice3/utils.py - - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py - - vllm_omni/deploy/minicpmo_4_5.yaml + source_file_dependencies: omni_minicpmo_4_5_function commands: - | timeout 20m bash -c " @@ -203,48 +166,7 @@ steps: mirror_hardwares: h100_1 - label: "Omni · MiniCPM-o 4.5 Duplex Test" - source_file_dependencies: - - tests/assets/minicpmo_4_5/response_required_16k.wav - - tests/assets/minicpmo_4_5/soft_interrupt_16k.wav - - tests/e2e/online_serving/test_minicpmo_4_5_duplex.py - - tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py - - tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py - - tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py - - tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py - - tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py - - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py - - vllm_omni/deploy/minicpmo_4_5.yaml - - tests/e2e/online_serving/helpers/minicpmo_4_5_duplex.py - - vllm_omni/engine/duplex/ - - vllm_omni/entrypoints/duplex/ - # Public duplex client library + its unit suites: client-side - # changes must re-run the duplex e2e group, which includes the - # live client probe (test_duplex_client_live.py, single GPU). - - vllm_omni/clients/ - - tests/clients/ - - tests/e2e/online_serving/test_duplex_client_live.py - - vllm_omni/entrypoints/duplex_request_client.py - - vllm_omni/outputs/duplex.py - - vllm_omni/model_executor/duplex_sampling.py - - tests/engine/duplex/ - - tests/entrypoints/duplex/ - - tests/model_executor/models/minicpmo_4_5/duplex/ - # Legacy Realtime probe client + video stacking used by the e2e drivers. - - vllm_omni/experimental/fullduplex/ - - vllm_omni/model_executor/models/minicpmo_4_5/ - - vllm_omni/model_executor/models/step_audio2/ - - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ - - vllm_omni/model_executor/models/cosyvoice3/utils.py - - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py - - vllm_omni/core/sched/ - - vllm_omni/outputs/multimodal_accumulation.py - - vllm_omni/engine/stage_init_utils.py - - vllm_omni/engine/async_omni_engine.py - - vllm_omni/engine/orchestrator.py - - vllm_omni/entrypoints/cli/serve.py - - vllm_omni/entrypoints/openai/api_server.py - - vllm_omni/entrypoints/openai/duplex_capability.py - - vllm_omni/config/stage_config.py + source_file_dependencies: omni_minicpmo_4_5_duplex_function commands: - | timeout 35m bash -c " @@ -254,29 +176,14 @@ steps: - label: "Omni · MiniCPM-o 4.5 Perf Test" timeout_in_minutes: 60 - source_file_dependencies: - - tests/dfx/perf/tests/test_minicpmo_4_5_ready.json - - vllm_omni/deploy/minicpmo_4_5.yaml - - vllm_omni/model_executor/models/minicpmo_4_5/ - - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py + source_file_dependencies: omni_minicpmo_4_5_perf commands: - export BENCHMARK_DIR=tests/dfx/perf/results - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_ready.json -m 'core_model and cuda' mirror_hardwares: h100_1 - label: "TTS · Qwen3-TTS CustomVoice Test" - source_file_dependencies: - - tests/helpers/media.py - - tests/e2e/online_serving/test_qwen3_tts_customvoice.py - - vllm_omni/model_executor/models/qwen3_tts/ - - vllm_omni/model_executor/models/common/qwen3_code_predictor.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 - - vllm_omni/deploy/qwen3_tts.yaml - - vllm_omni/entrypoints/openai/serving_speech.py - - vllm_omni/entrypoints/openai/tts_adapters/ - - vllm_omni/engine/stage_pool.py + source_file_dependencies: tts_qwen3_tts_function commands: - | timeout 20m bash -c " @@ -286,10 +193,7 @@ steps: mirror_hardwares: l4_1 - label: "Diffusion · Wan22 Test" - source_file_dependencies: - - tests/e2e/online_serving/test_wan22_t2v.py - - vllm_omni/diffusion/models/wan2_2/ - - vllm_omni/diffusion/models/dmd2/ + source_file_dependencies: diffusion_wan22_function commands: - | timeout 40m bash -c " @@ -299,12 +203,7 @@ steps: mirror_hardwares: h100_1 - label: "Diffusion · Cosmos3 Test" - source_file_dependencies: - - tests/e2e/online_serving/test_cosmos3.py - - vllm_omni/diffusion/models/cosmos3/ - - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py - - vllm_omni/experimental/world_models/ - - vllm_omni/model_extras/cosmos3.py + source_file_dependencies: diffusion_cosmos3_function commands: - | timeout 40m bash -c " @@ -315,14 +214,7 @@ steps: - label: "Diffusion · MiniMax-H3 DLO DP2 T2VA Test" timeout_in_minutes: 90 - source_file_dependencies: - - tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py - - tests/helpers/assertions.py - - tests/helpers/runtime.py - - vllm_omni/diffusion/models/minimax_h3/ - - vllm_omni/diffusion/offloader/ - - vllm_omni/diffusion/distributed/ - - vllm_omni/entrypoints/openai/ + source_file_dependencies: diffusion_minimax_h3_function commands: - | timeout 80m bash -c " @@ -332,11 +224,7 @@ steps: mirror_hardwares: h100_2 - label: "Diffusion · Qwen Image Test" - source_file_dependencies: - - tests/e2e/online_serving/test_qwen_image.py - - vllm_omni/diffusion/models/qwen_image/ - - vllm_omni/diffusion/models/dmd2/ - - vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py + source_file_dependencies: diffusion_qwen_image_function commands: - | timeout 40m bash -c " @@ -346,9 +234,7 @@ steps: mirror_hardwares: h100_1 - label: "Diffusion · HunyuanVideo-1.5 Test" - source_file_dependencies: - - tests/e2e/online_serving/test_hunyuan_video_15_expansion.py - - vllm_omni/diffusion/models/hunyuan_video/ + source_file_dependencies: diffusion_hunyuan_video_function commands: - | timeout 40m bash -c " diff --git a/.buildkite/cuda/test-weekly.yml b/.buildkite/cuda/test-weekly.yml index 614af957262..d53bc56e485 100644 --- a/.buildkite/cuda/test-weekly.yml +++ b/.buildkite/cuda/test-weekly.yml @@ -29,36 +29,42 @@ steps: steps: - label: "Reliability Test · qwen3-omni" + source_file_dependencies: omni_qwen3_omni_reliability timeout_in_minutes: 180 commands: - pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow" mirror_hardwares: h100_2 - label: "Reliability Test · wan22" + source_file_dependencies: diffusion_wan22_reliability timeout_in_minutes: 180 commands: - pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow" mirror_hardwares: h100_2 - label: "Reliability Test · hunyuan-image" + source_file_dependencies: diffusion_hunyuan_image3_reliability timeout_in_minutes: 180 commands: - pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow" mirror_hardwares: h100_4 - label: "Reliability Test · Invalid parameters · H100 · Single-GPU" + source_file_dependencies: reliability_invalid_param_h100 timeout_in_minutes: 90 commands: - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100 and cards_1" mirror_hardwares: h100_1 - label: "Reliability Test · Invalid parameters · H100 · 2-GPU" + source_file_dependencies: reliability_invalid_param_h100 timeout_in_minutes: 90 commands: - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100 and cards_2" mirror_hardwares: h100_2 - label: "Reliability Test · Invalid parameters · L4" + source_file_dependencies: reliability_invalid_param_l4 timeout_in_minutes: 120 commands: - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and L4" @@ -70,6 +76,7 @@ steps: steps: - label: ":full_moon: Omni · Perf Test · vLLM Text" + source_file_dependencies: omni_qwen3_omni_perf key: weekly-omni-performance-vllm-text timeout_in_minutes: 180 commands: @@ -96,6 +103,7 @@ steps: mirror_hardwares: h100_2 - label: ":full_moon: Omni · Perf Test · Multi-Replica" + source_file_dependencies: omni_qwen3_omni_perf key: weekly-omni-performance-multi-replicas timeout_in_minutes: 180 commands: @@ -221,21 +229,262 @@ steps: --run-level "full_model" mirror_hardwares: l4_4 - - label: "Diffusion · H100 · Single-GPU" - timeout_in_minutes: 150 + # Split the old shared H100 Single-GPU sweep: one pytest process collected + # every slow+cards_1 diffusion file and exceeded 150 minutes. + - label: "Diffusion · Qwen-Image · H100 · Single-GPU" + timeout_in_minutes: 120 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_qwen_image_edit.py + tests/e2e/online_serving/test_qwen_image_edit_expansion.py + tests/e2e/online_serving/test_qwen_image_layered.py + tests/e2e/online_serving/test_qwen_image_layered_expansion.py + tests/e2e/accuracy/test_qwen_image_edit.py + tests/e2e/accuracy/test_qwen_image_layered.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · Wan2.1 VACE · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_wan_2_1_vace_expansion.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · LongCat · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_longcat_image_expansion.py + tests/e2e/online_serving/test_longcat_image_edit_expansion.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · Bagel · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_bagel.py + tests/e2e/online_serving/test_bagel_expansion.py + tests/e2e/offline_inference/test_bagel.py + tests/e2e/offline_inference/test_bagel_expansion.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · Boogu · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_boogu_image.py + tests/e2e/online_serving/test_boogu_image_edit.py + tests/e2e/online_serving/test_boogu_image_expansion.py + tests/e2e/online_serving/test_boogu_image_edit_expansion.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · Flux · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_flux_2_dev_expansion.py + tests/e2e/offline_inference/test_flux1_schnell_expansion.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · LTX · H100 · Single-GPU" + timeout_in_minutes: 120 commands: - >- pytest -sv - tests/e2e/ -m "slow and diffusion and H100 and cards_1" + tests/e2e/online_serving/test_ltx.py + tests/e2e/online_serving/test_ltx25.py + tests/e2e/accuracy/ltx/test_ltx_official_similarity.py + tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py + -m "slow and diffusion and H100 and cards_1" --run-level "full_model" mirror_hardwares: h100_1 - - label: "Diffusion · H100 · 2-GPU" + - label: "Diffusion · Krea2 · H100 · Single-GPU" timeout_in_minutes: 90 commands: - >- pytest -sv - tests/e2e/ -m "slow and diffusion and H100 and cards_2" + tests/e2e/online_serving/test_krea2_expansion.py + tests/e2e/offline_inference/test_krea2_expansion.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · HiDream-I1 · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_hidream_i1_full_expansion.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · SANA-Video · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_sana_video_expansion.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · SenseNova-U1 · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/offline_inference/test_sensenova_u1_expansion.py + tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py + tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · MammothModa2 · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/offline_inference/test_mammoth_moda2_expansion.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + - label: "Diffusion · Helios · H100 · Single-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/accuracy/test_video_streaming_output_similarity.py + -m "slow and diffusion and H100 and cards_1" + --run-level "full_model" + mirror_hardwares: h100_1 + + # Split the old shared H100 2-GPU sweep: one pytest process collected every + # slow+cards_2 diffusion file and exceeded 90 minutes. + - label: "Diffusion · Qwen-Image · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_qwen_image_edit_expansion.py + tests/e2e/online_serving/test_qwen_image_layered_expansion.py + -m "slow and diffusion and H100 and cards_2" + --run-level "full_model" + mirror_hardwares: h100_2 + + - label: "Diffusion · Wan2.1 VACE · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_wan_2_1_vace_expansion.py + -m "slow and diffusion and H100 and cards_2" + --run-level "full_model" + mirror_hardwares: h100_2 + + - label: "Diffusion · LongCat · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_longcat_image_expansion.py + tests/e2e/online_serving/test_longcat_image_edit_expansion.py + -m "slow and diffusion and H100 and cards_2" + --run-level "full_model" + mirror_hardwares: h100_2 + + - label: "Diffusion · Bagel · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_bagel_expansion.py + -m "slow and diffusion and H100 and cards_2" + --run-level "full_model" + mirror_hardwares: h100_2 + + - label: "Diffusion · Boogu · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_boogu_image_expansion.py + tests/e2e/online_serving/test_boogu_image_edit_expansion.py + -m "slow and diffusion and H100 and cards_2" + --run-level "full_model" + mirror_hardwares: h100_2 + + - label: "Diffusion · Flux · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_flux_2_dev_expansion.py + tests/e2e/offline_inference/test_flux1_kontext_expansion.py + -m "slow and diffusion and H100 and cards_2" + --run-level "full_model" + mirror_hardwares: h100_2 + + - label: "Diffusion · GLM-Image · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_glm_image_expansion.py + tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py + -m "slow and diffusion and H100 and cards_2" + --run-level "full_model" + mirror_hardwares: h100_2 + + - label: "Diffusion · LTX · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_ltx.py + -m "slow and diffusion and H100 and cards_2" + --run-level "full_model" + mirror_hardwares: h100_2 + + - label: "Diffusion · Krea2 · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/offline_inference/test_krea2_expansion.py + -m "slow and diffusion and H100 and cards_2" + --run-level "full_model" + mirror_hardwares: h100_2 + + - label: "Diffusion · DreamZero · H100 · 2-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_dreamzero_expansion.py + tests/e2e/accuracy/test_dreamzero.py + -m "slow and diffusion and H100 and cards_2" --run-level "full_model" mirror_hardwares: h100_2 @@ -246,4 +495,33 @@ steps: pytest -sv tests/e2e/ -m "slow and diffusion and H100 and cards_4" --run-level "full_model" + --ignore=tests/e2e/online_serving/test_magi2.py + mirror_hardwares: h100_4 + + # Dedicated: HiDream-O1 accuracy needs a long TP2 run; the shared H100 + # 2-GPU sweep is only 90 minutes. + - label: "Diffusion · HiDream-O1-Image Accuracy · H100 · 2-GPU" + timeout_in_minutes: 180 + commands: + - pytest -s -v tests/e2e/accuracy/test_hidream_o1_image.py -m "slow and diffusion and cuda and H100 and cards_2" --run-level "full_model" + mirror_hardwares: h100_2 + + # Dedicated: Boogu CFG-parallel size 3 needs 3 GPUs; weekly has no H100 3-GPU sweep. + - label: "Diffusion · Boogu Edit CFG · H100 · 3-GPU" + timeout_in_minutes: 120 + commands: + - >- + pytest -sv + tests/e2e/online_serving/test_boogu_image_edit_expansion.py + -k "test_double_guidance_cfg_parallel and cfg_parallel_3_full" + -m "slow and diffusion and cuda and H100 and cards_3" + --run-level "full_model" + mirror_hardwares: h100_3 + + # Dedicated: MAGI-2 Preview serving smoke needs ~30 min server init; the + # shared H100 4-GPU sweep is only 90 minutes. + - label: "Diffusion · MAGI-2 Preview Serving Smoke · H100 · 4-GPU" + timeout_in_minutes: 180 + commands: + - pytest -s -v tests/e2e/online_serving/test_magi2.py -m "slow and diffusion and cuda and H100 and cards_4" --run-level "full_model" mirror_hardwares: h100_4 diff --git a/.buildkite/npu/bootstrap-upload-steps.yml b/.buildkite/npu/bootstrap-upload-steps.yml index 398fbc6dced..d88fcc70bd2 100644 --- a/.buildkite/npu/bootstrap-upload-steps.yml +++ b/.buildkite/npu/bootstrap-upload-steps.yml @@ -93,7 +93,7 @@ steps: command: | python3 .buildkite/common/scripts/upload_pipeline.py --upload .buildkite/npu/test-npu-ready.yml - # L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild) + # L4 Test — main+NIGHTLY=1 (scheduled), or PR nightly-test label - label: "Upload Nightly Pipeline" depends_on: - image-build-a2 @@ -102,6 +102,11 @@ steps: - image-build-310p key: upload-nightly-pipeline commands: - - python3 .buildkite/common/scripts/upload_pipeline.py --upload .buildkite/npu/test-npu-nightly.yml + - | + EXTRA="" + if [ "${NIGHTLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then + EXTRA="--all" + fi + python3 .buildkite/common/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/npu/test-npu-nightly.yml agents: queue: "cpu_queue_premerge" diff --git a/.buildkite/npu/pipeline-npu-a3.yml b/.buildkite/npu/pipeline-npu-a3.yml index a57f5f8e569..3544b384d4c 100644 --- a/.buildkite/npu/pipeline-npu-a3.yml +++ b/.buildkite/npu/pipeline-npu-a3.yml @@ -96,13 +96,7 @@ # key: upload-nightly-pipeline # if: >- # (build.branch == "main" && build.env("NIGHTLY") == "1") || - # (build.branch != "main" && ( - # build.pull_request.labels includes "nightly-test" || - # build.pull_request.labels includes "omni-test" || - # build.pull_request.labels includes "tts-test" || - # build.pull_request.labels includes "diffusion-x2iat-test" || - # build.pull_request.labels includes "diffusion-x2v-test" - # )) + # (build.branch != "main" && build.pull_request.labels includes "nightly-test") # commands: # - buildkite-agent pipeline upload .buildkite/npu/test-npu-nightly.yml # agents: diff --git a/.buildkite/npu/test-npu-nightly.yml b/.buildkite/npu/test-npu-nightly.yml index e049783ce8c..f64fa711f64 100644 --- a/.buildkite/npu/test-npu-nightly.yml +++ b/.buildkite/npu/test-npu-nightly.yml @@ -4,10 +4,10 @@ steps: depends_on: upload-nightly-pipeline if: >- build.env("NIGHTLY") == "1" || - build.pull_request.labels includes "nightly-test" || - build.pull_request.labels includes "omni-test" + build.pull_request.labels includes "nightly-test" steps: - label: ":full_moon: Omni · Perf Test · No Async Chunk" + source_file_dependencies: omni_qwen3_omni_perf key: nightly-omni-performance-no-async-chunk timeout_in_minutes: 180 mirror_hardwares: a3_npu_4 @@ -24,6 +24,7 @@ steps: buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT - label: ":full_moon: Omni · Perf Test · Async Chunk" + source_file_dependencies: omni_qwen3_omni_perf key: nightly-omni-performance-async-chunk timeout_in_minutes: 180 mirror_hardwares: a3_npu_4 @@ -40,6 +41,7 @@ steps: buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT - label: ":full_moon: Omni · MiniCPM-o 4.5 Test" + source_file_dependencies: omni_minicpmo_4_5_function mirror_hardwares: a3_npu_2 env: VLLM_WORKER_MULTIPROC_METHOD: spawn @@ -48,6 +50,7 @@ steps: commands: - pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5_expansion.py -m 'full_model' --run-level 'full_model' - label: ":full_moon: Omni · MiniCPM-o 4.5 · Duplex Test" + source_file_dependencies: omni_minicpmo_4_5_duplex_function timeout_in_minutes: 180 mirror_hardwares: a3_npu_2 env: @@ -62,6 +65,7 @@ steps: -m "npu and A3 and omni" --run-level full_model - label: ":full_moon: Omni · MiniCPM-o 4.5 · Accuracy Test" + source_file_dependencies: omni_minicpmo_4_5_accuracy timeout_in_minutes: 180 mirror_hardwares: a3_npu_2 env: @@ -82,6 +86,7 @@ steps: buildkite-agent artifact upload "tests/e2e/accuracy/minicpmo_4_5/results/*.json" exit $$EXIT - label: ":full_moon: Omni · MiniCPM-o 4.5 · Perf Test" + source_file_dependencies: omni_minicpmo_4_5_perf key: nightly-omni-performance-minicpmo-4-5-npu timeout_in_minutes: 180 mirror_hardwares: a3_npu_2 @@ -98,6 +103,7 @@ steps: buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT - label: ":full_moon: Omni · MiniCPM-o 4.5 · Perf Test (A2)" + source_file_dependencies: omni_minicpmo_4_5_perf key: nightly-omni-performance-minicpmo-4-5-npu-a2 timeout_in_minutes: 180 mirror_hardwares: a2b3_npu_1 @@ -115,6 +121,7 @@ steps: buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT - label: ":full_moon: Omni · MiniCPM-o 4.5 · Duplex Seed-TTS Perf Test" + source_file_dependencies: omni_minicpmo_4_5_duplex_perf key: nightly-omni-performance-minicpmo-4-5-duplex-seed-tts-npu timeout_in_minutes: 180 mirror_hardwares: a3_npu_2 @@ -133,9 +140,10 @@ steps: - group: ":card_index_dividers: TTS Model Test" key: nightly-tts-test-group depends_on: upload-nightly-pipeline - if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test" + if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" steps: - label: ":full_moon: TTS · Perf Test" + source_file_dependencies: tts_qwen3_tts_perf timeout_in_minutes: 180 mirror_hardwares: a3_npu_2 env: @@ -151,15 +159,15 @@ steps: EXIT=$$? buildkite-agent artifact upload "tests/dfx/perf/results/*.json" exit $$EXIT - - group: ":card_index_dividers: Diffusion I2I Model Test" - key: nightly-HunyuanImage3-x2i-test-group + - group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test" + key: nightly-diffusion-x2iat-group depends_on: upload-nightly-pipeline if: >- build.env("NIGHTLY") == "1" || - build.pull_request.labels includes "nightly-test" || - build.pull_request.labels includes "diffusion-x2iat-test" + build.pull_request.labels includes "nightly-test" steps: - label: ":full_moon: HunyuanImage3-DIT · Accuracy Test" + source_file_dependencies: diffusion_hunyuan_image3_accuracy timeout_in_minutes: 180 mirror_hardwares: a3_npu_4 env: @@ -169,15 +177,8 @@ steps: DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA commands: - pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m "full_model and A3" --run-level full_model - - group: ":card_index_dividers: Diffusion Benchmark Perf Test" - key: nightly-diffusion-benchmark-group - depends_on: upload-nightly-pipeline - if: >- - build.env("NIGHTLY") == "1" || - build.pull_request.labels includes "nightly-test" || - build.pull_request.labels includes "diffusion-x2iat-test" - steps: - label: ":full_moon: Diffusion Perf · HunyuanImage3 TP=2 CFGP=2" + source_file_dependencies: diffusion_hunyuan_image3_perf timeout_in_minutes: 180 mirror_hardwares: a3_npu_4 env: @@ -194,6 +195,7 @@ steps: buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" exit $$EXIT - label: ":full_moon: Diffusion Perf · HunyuanImage3 TP=2 SP=2" + source_file_dependencies: diffusion_hunyuan_image3_perf timeout_in_minutes: 180 mirror_hardwares: a3_npu_4 env: @@ -210,6 +212,7 @@ steps: buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" exit $$EXIT - label: ":full_moon: Diffusion Perf · HunyuanImage3 TP=4" + source_file_dependencies: diffusion_hunyuan_image3_perf timeout_in_minutes: 180 mirror_hardwares: a3_npu_4 env: @@ -230,10 +233,10 @@ steps: depends_on: upload-nightly-pipeline if: >- build.env("NIGHTLY") == "1" || - build.pull_request.labels includes "nightly-test" || - build.pull_request.labels includes "diffusion-x2v-test" + build.pull_request.labels includes "nightly-test" steps: - - label: ":full_moon: Diffusion X2V · Function Test" + - label: ":full_moon: Diffusion X2V · Wan 2.2 Function Test" + source_file_dependencies: diffusion_wan22_function mirror_hardwares: a2b3_npu_4 env: VLLM_WORKER_MULTIPROC_METHOD: spawn @@ -241,7 +244,8 @@ steps: HF_TOKEN: "${HF_TOKEN}" commands: - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k npu_i2v --run-level "full_model" - - label: ":full_moon: Diffusion X2V · Perf Test" + - label: ":full_moon: Diffusion X2V · Wan 2.2 Perf Test" + source_file_dependencies: diffusion_wan22_perf key: nightly-diffusion-x2v-performance timeout_in_minutes: 180 mirror_hardwares: a3_npu_4 diff --git a/.claude/skills/vllm-omni-test/SKILL.md b/.claude/skills/vllm-omni-test/SKILL.md index 2d42267882d..f0fb009d56c 100644 --- a/.claude/skills/vllm-omni-test/SKILL.md +++ b/.claude/skills/vllm-omni-test/SKILL.md @@ -468,10 +468,10 @@ When the user asks for **invalid parameter validation**, **invalid request bodie | **Invalid parameters Test · H100** | `pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100"` | Diffusion / omni / video invalid-param tests with `H100` hardware mark | | **Invalid parameters Test · L4** | `pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and L4"` | TTS / lighter models marked `L4` | -- **Trigger:** `build.env("WEEKLY") == "1"` or PR label `weekly-test`. +- **Trigger:** `build.env("WEEKLY") == "1"` (all weekly jobs) or PR label `weekly-test` (Reliability / Perf, then filtered by `source_file_dependencies`). - **Default:** extending an existing `invalid_param_test` script needs **no YAML edit** — the weekly steps already sweep the whole directory. - **Edit YAML only** when adding a **new hardware queue**, a **new top-level script** that must run in isolation, or a **model-specific weekly shard** (mirror neighboring reliability steps). -- Weekly steps **do not** use `source_file_dependencies`. +- Weekly Reliability / Perf jobs declare `source_file_dependencies`. Filtering applies only on the `weekly-test` label; `WEEKLY=1` on `main` keeps every job in those groups. CPU coverage (`WEEKLY=1` only) and the E2E Tests group (`NON_CRITICAL=1` only) do **not** use this field. **Example — append to `test_invalid_image_generation.py`:** @@ -644,7 +644,7 @@ If the test is not already collected by an existing pipeline command (for exampl | **L2** | `test-ready.yml` with `core_model` + `source_file_dependencies` | merge / nightly pipelines | | **Invalid param** | `test-weekly.yml` — **Invalid parameters Test** group (`-m "slow and H100"` / `-m "slow and L4"`). Usually **no YAML edit** when appending cases to existing `invalid_param_test` scripts | `test-ready.yml`, `test-merge.yml`, `test-nightly.yml` | -`test-nightly.yml` shards use **explicit pytest file paths** in `commands` (and PR labels like `diffusion-x2iat-test`); they generally **do not** use `source_file_dependencies`. Only suggest merge/ready YAML when the requested level is L3/L2 (or the user explicitly asks for multi-level CI). +`test-nightly.yml` shards use **explicit pytest file paths** in `commands` plus `source_file_dependencies`; PR trigger is `nightly-test` (path filter), or `NIGHTLY=1` on `main` (all jobs). Only suggest merge/ready YAML when the requested level is L3/L2 (or the user explicitly asks for multi-level CI). **`test-nightly.yml` top-level groups** (each uses `full_model` + type marker in pytest `-m`): @@ -653,7 +653,7 @@ If the test is not already collected by an existing pipeline command (for exampl | **Omni Model Test** | `omni` | Function / doc / accuracy / **perf** for Qwen-Omni family | | **TTS Model Test** | `tts` | `-m "full_model and L4 and tts"` (and H100 variants when enabled); **Perf Test** uses `run_benchmark.py` + `test_tts.json` | | **Diffusion X2I(&A&T) Model Test** | `diffusion` | Image/audio/text diffusion; explicit file shards for **Function**; **Perf Test · <Model>** uses `run_diffusion_benchmark.py` + JSON config | -| **Diffusion X2V Model Test** | `diffusion` | Video-only (Wan, HunyuanVideo, …); PR label `diffusion-x2v-test`; separate **Perf** steps when in-tree | +| **Diffusion X2V Model Test** | `diffusion` | Video-only (Wan, HunyuanVideo, …); separate **Perf** steps when in-tree | Guidelines when editing YAML: diff --git a/.claude/skills/vllm-omni-test/references/test-routing.md b/.claude/skills/vllm-omni-test/references/test-routing.md index 5b90957d68e..a61324e6df4 100644 --- a/.claude/skills/vllm-omni-test/references/test-routing.md +++ b/.claude/skills/vllm-omni-test/references/test-routing.md @@ -37,7 +37,7 @@ Every model-centric e2e test must declare **exactly one** type marker: | **Diffusion X2I(&A&T) Model Test** | x2**i** / x2**a** / x2**t** — image, audio, text (non-video) | Qwen-Image*, BAGEL, FLUX, SD3, Z-Image, LongCat, DreamZero | | **Diffusion X2V Model Test** | x2**v** — video only | Wan2.2, HunyuanVideo 1.5, Wan VACE | -PR labels for selective nightly runs: `diffusion-x2iat-test`, `diffusion-x2v-test` (plus `nightly-test` / `NIGHTLY=1`). +PR label for L4 on a branch: `nightly-test` (jobs then filtered by `source_file_dependencies`). Scheduled `main` uses `NIGHTLY=1` and runs every nightly job. L5 Reliability / Perf use `weekly-test` with the same path filter; CPU coverage and the weekly E2E group are env-only (`WEEKLY=1` / `NON_CRITICAL=1`) and are not source-filtered. ## Level and Marker Mapping @@ -238,7 +238,7 @@ pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and L4" pytest -s -v dfx/reliability/invalid_param_test/test_invalid_image_generation.py::test_images_generations_invalid_requests -m "slow and H100" ``` -**Trigger:** `WEEKLY=1` or PR label `weekly-test`. Appending cases to existing scripts usually needs **no YAML edit** (directory sweep). No `source_file_dependencies` on weekly steps. +**Trigger:** `WEEKLY=1` (all jobs) or PR label `weekly-test` (Reliability / Perf, filtered by `source_file_dependencies`). Appending cases to existing scripts usually needs **no YAML edit** (directory sweep). ### Platform-targeted examples diff --git a/docs/contributing/ci/ci_settings.md b/docs/contributing/ci/ci_settings.md index 885ebd5280c..fb3af0ed3cb 100644 --- a/docs/contributing/ci/ci_settings.md +++ b/docs/contributing/ci/ci_settings.md @@ -15,7 +15,8 @@ Canonical layout (prefer these paths for new changes): │ │ ├── skip_ci.py # skip-ci decision (docs / skip-mark / CI YAML paths) │ │ ├── upload_pipeline.py # Bootstrap + test-pipeline uploader (CUDA/NPU) │ │ └── resolve_skip_ci.sh # Shell helpers for AMD/Intel bootstrap -│ └── ci_mirror_hardwares.yml # CUDA uploader presets (referenced by name only) +│ ├── ci_mirror_hardwares.yml # CUDA/NPU hardware presets (referenced by name only) +│ └── ci_source_file_dependencies.yml # L2–L5 path-prefix presets (referenced by name only) ├── cuda/ # Primary NVIDIA CUDA CI │ ├── pipeline.yml # Bootstrap entry (hook upload) │ ├── bootstrap-upload-steps.yml # Bootstrap child steps (upload_pipeline --upload) @@ -108,19 +109,21 @@ Canonical layout (prefer these paths for new changes): Inferred form: unset/`empty` `MIRROR_HW` matches `H100` or `L4` in `-m` (both present → H100); no match skips the step. `MIRROR_HW=b200` must appear in `-m` or the step is skipped. The uploader does **not** rewrite pytest `-m`; B200 collection comes from the YAML expression and from tests that declare `B200` in `hardware_test` / `hardware_marks` (for example `res={"cuda": ["H100", "B200"]}`). `MIRROR_HW` must be empty or `b200` (case-insensitive); unknown values (for example `b20o`) **fail the upload**. A CUDA preset string such as `mirror_hardwares: h100_4` is skipped when `MIRROR_HW=b200`. Unset `MIRROR_HW` keeps string presets. NPU presets ignore `MIRROR_HW`. A `mirror_hardwares` name that is not a key in `ci_mirror_hardwares.yml` **fails the upload**. + **Path filter in YAML:** `source_file_dependencies: ` (string, or a list of preset names)—preset names in [`common/ci_source_file_dependencies.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/common/ci_source_file_dependencies.yml). Required on **E2E Test** leaf jobs in L2/L3, and on L4/L5 leaf jobs that should be diff-gated on PR labels; see [Step filtering](#step-filtering). + **Conventions** - **`depends_on`:** leaf jobs depend on `upload-ready-pipeline`, `upload-merge-pipeline`, etc. - **`group` / `label`:** `:card_index_dividers:` groups; labels like `Diffusion · Qwen Image Test`. - **`commands`:** `timeout … pytest …` with markers and `--run-level` for the pipeline level. - - **`source_file_dependencies`:** required on **E2E Test** leaf jobs in L2/L3; see [Step filtering](#step-filtering). + - **`source_file_dependencies`:** required on **E2E Test** leaf jobs in L2/L3, and on L4/L5 leaf jobs that should be diff-gated on PR labels; name a job key from `ci_source_file_dependencies.yml` (see [Step filtering](#step-filtering)). The key must list every script the step runs; `commands` is not scanned. **Adding a job** 1. Pick the level file (`test-ready.yml` / `test-merge.yml` / `test-nightly.yml` / `test-weekly.yml`). 2. Add a step under the right **group** (usually **E2E Test** for model pytest). - 3. Set `label`, `commands`, `mirror_hardwares`, `depends_on: upload-ready-pipeline` (or `upload-merge-pipeline` / `upload-nightly-pipeline` / `upload-weekly-pipeline`). - 4. For L2/L3 E2E, add `source_file_dependencies` (pytest + model + deploy YAML prefixes). + 3. Set `label`, `commands`, `mirror_hardwares`, `depends_on: upload--pipeline`. + 4. For L2–L5 jobs that should be diff-gated, add a job key in `common/ci_source_file_dependencies.yml` and set `source_file_dependencies: ` on the step. YAML anchors hold model business-code paths only. The job key (`{model}_function` / `_perf` / `_accuracy` / `_reliability` / `_doc` / `_cov`) aliases that anchor and lists the pytest / example / benchmark scripts the step actually runs. The uploader does **not** parse `commands` for path filters—omitting a script from the key means editing that script will not select the job. 5. Dry-run: ```bash @@ -144,6 +147,8 @@ Canonical layout (prefer these paths for new changes): - A5: `a5_npu_2`, `a5_npu_4`, `a5_npu_8`, `a5_npu_16` - 310P: `310p_npu_1`, `310p_npu_2`, `310p_npu_4` + **Path filter in YAML:** same `source_file_dependencies` presets as CUDA, on **L4** leaf jobs in `test-npu-nightly.yml`. Filtering applies on the PR `nightly-test` label; `main` + `NIGHTLY=1` keeps every job. If no job-key prefix matches, a change to the pipeline YAML being uploaded, or to any path under `source_filter_fallback`, keeps every job in that upload. See [Step filtering](#step-filtering). + **Conventions** - **`depends_on: upload-ready-pipeline`** (or `upload-nightly-pipeline`) ties jobs to bootstrap upload keys. @@ -153,7 +158,7 @@ Canonical layout (prefer these paths for new changes): 1. Edit `test-npu-ready.yml` (L2) or `test-npu-nightly.yml` (L4). 2. Add a step with `mirror_hardwares` (add a new preset in `ci_mirror_hardwares.yml` first if needed). - 3. Set `commands` to your pytest file and markers. + 3. Set `commands` to your pytest file and markers. For L4, set `source_file_dependencies` to a job key from `ci_source_file_dependencies.yml` that aliases the model anchor and lists that pytest file. 4. Dry-run: ```bash @@ -213,7 +218,7 @@ PR diffs drive **two independent skip layers**. Both read changed files from git | Layer | Script | When | What is skipped | Where you configure | | ----- | ------ | ---- | --------------- | ------------------- | | **Bootstrap** | [`skip_ci.py`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/common/scripts/skip_ci.py) + [`upload_pipeline.py`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/common/scripts/upload_pipeline.py) | Before child test pipelines upload (`cuda/pipeline.yml`, `npu/pipeline-npu.yml`, AMD/Intel bootstraps) | Entire default CI, or whole L2/L3 upload for a platform | Whitelists in `skip_ci.py`; bootstrap `if` injected by step `key` in `upload_pipeline.py` | -| **Step filter** | [`upload_pipeline.py`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/common/scripts/upload_pipeline.py) | While uploading CUDA L2/L3 YAML | Individual Buildkite steps inside `test-ready.yml` / `test-merge.yml` | `source_file_dependencies` on each step or group | +| **Step filter** | [`upload_pipeline.py`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/common/scripts/upload_pipeline.py) | While uploading CUDA L2–L5 YAML | Individual Buildkite steps inside `test-ready.yml` / `test-merge.yml` / `test-nightly.yml` / `test-weekly.yml` | `source_file_dependencies` key on each step (presets in `ci_source_file_dependencies.yml`) | **Changed files** (both layers): @@ -306,40 +311,38 @@ Unit coverage: `tests/buildkite/test_skip_ci.py`. #### Step filtering {#step-filtering} -CUDA **L2** (`.buildkite/cuda/test-ready.yml`) and **L3** (`.buildkite/cuda/test-merge.yml`) only. Bootstrap upload entry: `upload_pipeline.py --upload .buildkite/cuda/bootstrap-upload-steps.yml`. +CUDA **L2** (`.buildkite/cuda/test-ready.yml`), **L3** (`.buildkite/cuda/test-merge.yml`), **L4** (`.buildkite/cuda/test-nightly.yml`), **L5** (`.buildkite/cuda/test-weekly.yml`), and NPU **L4** (`.buildkite/npu/test-npu-nightly.yml`). Bootstrap upload entry: `upload_pipeline.py --upload .buildkite/cuda/bootstrap-upload-steps.yml` (or `npu/bootstrap-upload-steps.yml`). + +**When filtering applies:** only **PR label** uploads (`ready`, `merge-test`, `nightly-test`, `weekly-test`). On **`main` + env** (`NIGHTLY=1`, post-merge L3, `WEEKLY=1` / `NON_CRITICAL=1`, `WEEKLY=1 --e2e`) the uploader keeps every job and still strips the key. If **no** job-key prefix matches the diff, a change to the pipeline YAML being uploaded, or any path under the `source_filter_fallback` key in `ci_source_file_dependencies.yml`, keeps every job in that upload. `source_filter_fallback` is not a job key — do not attach it to a step. When any job-key prefix already matches, normal filtering wins (the fallback does not expand the selection). Without this fallback, command / env / hardware-only edits would match nothing (NPU drops every step; CUDA nightly keeps only the email aggregator). A different pipeline YAML does not trigger the fallback for this upload. **Uploader-only keys** — removed before Buildkite sees the YAML; never used at runtime on agents: | Key | Purpose | | --- | ------- | -| `source_file_dependencies` | List of path **prefixes**. If any changed file equals a prefix or starts with `prefix/`, keep the step (or group); otherwise omit it. | +| `source_file_dependencies` | Named preset (or list of presets) from [`ci_source_file_dependencies.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/common/ci_source_file_dependencies.yml). The uploader expands only that key; it does **not** scan `commands`. If any changed file equals a listed prefix or starts with `prefix/`, keep the step (or group); otherwise omit it. Inline path-prefix lists are still accepted. | | `mirror_hardwares` | Expand to `agents` + `plugins` (+ optional `image`) from `ci_mirror_hardwares.yml`. | ### Policy -- **Always uploaded** (no key): groups outside **E2E Test**—Simple Test, Diffusion unit tests, Engine/Model Executor, Distributed, Custom Pipeline, Entrypoints (L2), LoRA / Entrypoints (L3). -- **Diff-gated**: every **E2E Test** leaf job. List the smallest prefix set per step—pytest file(s), model code under `vllm_omni/model_executor/models/` or `vllm_omni/diffusion/models/`, plus `stage_input_processors/` and `vllm_omni/deploy/*.yaml` when applicable. A **group** may define the key instead; the whole group drops if no prefix matches. +- **Always uploaded** (no key): L2/L3 groups outside **E2E Test**—Simple Test, Diffusion unit tests, Engine/Model Executor, Distributed, Custom Pipeline, Entrypoints (L2), LoRA / Entrypoints (L3). +- **Diff-gated**: every **E2E Test** leaf job in L2/L3, L4 nightly leaf jobs (CUDA and NPU), and L5 weekly jobs whose `if` includes the `weekly-test` label (Reliability / Perf). Jobs gated only by `WEEKLY=1` or `NON_CRITICAL=1` have no `source_file_dependencies`. Registry layout: YAML anchors are model business code only (`vllm_omni/model_executor/models/`, `vllm_omni/diffusion/models/`, plus `stage_input_processors/` and `vllm_omni/deploy/*.yaml` when applicable). Job keys alias the anchor and **also** list the scripts that job runs (pytest files, `run_cov_split.sh --offline` / `--online` targets, example README runners). Feature keys stay cross-cutting. A **group** may define the key instead; the whole group drops if no prefix matches. ### YAML examples ```yaml - label: "Diffusion · Qwen Image Test" - source_file_dependencies: - - tests/e2e/online_serving/test_qwen_image.py - - vllm_omni/diffusion/models/qwen_image/ + source_file_dependencies: diffusion_qwen_image_function + # diffusion_qwen_image_function: *qwen_image + tests/e2e/online_serving/test_qwen_image.py commands: - pytest -s -v tests/e2e/online_serving/test_qwen_image.py -m 'core_model' ... mirror_hardwares: h100_1 - label: "TTS · Qwen3-TTS CustomVoice Test" - source_file_dependencies: - - tests/e2e/online_serving/test_qwen3_tts_customvoice.py - - vllm_omni/model_executor/models/qwen3_tts/ - - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py - - vllm_omni/deploy/qwen3_tts.yaml + source_file_dependencies: tts_qwen3_tts_function + # tts_qwen3_tts_function: *qwen3_tts + tests/e2e/online_serving/test_qwen3_tts_customvoice.py commands: - pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py ... - mirror_hardwares: l4_4 + mirror_hardwares: l4_1 ``` ### Local dry-run @@ -443,15 +446,16 @@ plugin, which provides the agent. `l4_*` jobs run on the EKS `l4-k8s` queue still needs `mount-buildkite-agent: true`. List both `run_cov_split.sh` and `pyproject.toml` in every opted-in job's -`source_file_dependencies` — both change what the job measures, so without them a +`source_file_dependencies` preset — both change what the job measures, so without them a change there is filtered out of normal PR builds and only surfaces in a later nightly. `tests/buildkite/test_upload_pipeline.py` covers the filter behavior with a synthetic job (it does not pin real merge labels). -Editing only the surrounding CI YAML still does not schedule them, so a PR that -touches just the wiring needs a full E2E run (or the commands run on a GPU host) -to produce artifacts. When checking a new model's -artifacts, compare `lines-covered` between the online and offline XML rather than -just confirming both files exist. +If no job-key prefix matches, editing the pipeline YAML being uploaded, or any +path under the `source_filter_fallback` registry key, keeps every job in that +upload (including coverage jobs) so the wiring can be validated before merge. +When a job-key prefix already matches, bypass does not expand the selection. A +change to a different pipeline YAML still does not schedule them. When checking a new model's artifacts, compare `lines-covered` +between the online and offline XML rather than just confirming both files exist. ### Validation checklist diff --git a/tests/buildkite/test_upload_pipeline.py b/tests/buildkite/test_upload_pipeline.py index 7492a4b8281..06e2b975c3c 100644 --- a/tests/buildkite/test_upload_pipeline.py +++ b/tests/buildkite/test_upload_pipeline.py @@ -13,11 +13,15 @@ from skip_ci import resolve_ci_decision # noqa: E402 from upload_pipeline import ( # noqa: E402 CUDA_HF_TOKEN_EXPORT, + NIGHTLY_LABEL_IF, + _changed_files_for_source_filter, _expand_mirror_hardwares, _get_mirror_hw_selector, _load_bootstrap_steps, + _load_source_file_dependencies, _render_bootstrap_pipeline, _render_test_pipeline, + _resolve_source_file_dependencies, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -101,6 +105,26 @@ def test_npu_docs_only_does_not_upload_ready_on_nightly() -> None: assert 'build.env("NIGHTLY") == "1"' in by_key["upload-nightly-pipeline"]["if"] +def test_nightly_label_if_is_only_nightly_test() -> None: + assert 'labels includes "nightly-test"' in NIGHTLY_LABEL_IF + forbidden = ( + 'includes "omni-test"', + 'includes "tts-test"', + 'includes "diffusion-x2iat-test"', + 'includes "diffusion-x2v-test"', + ) + for needle in forbidden: + assert needle not in NIGHTLY_LABEL_IF + for path in ( + Path(".buildkite/cuda/test-nightly.yml"), + Path(".buildkite/npu/test-npu-nightly.yml"), + Path(".buildkite/common/scripts/upload_pipeline.py"), + ): + text = path.read_text(encoding="utf-8") + for needle in forbidden: + assert needle not in text, f"{path} still references {needle}" + + def test_yaml_gated_l45_only_does_not_unconditionally_build_image() -> None: rendered = _render([".buildkite/cuda/test-nightly.yml"]) assert "if: true" not in rendered @@ -366,8 +390,17 @@ def test_mirror_hw_typo_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None: ) -def _surviving_labels(doc: dict, changed_files: list[str]) -> set[str]: - rendered = _render_test_pipeline(doc, changed_files=changed_files) +def _surviving_labels( + doc: dict, + changed_files: list[str], + *, + pipeline_path: Path | None = None, +) -> set[str]: + rendered = _render_test_pipeline( + doc, + changed_files=changed_files, + pipeline_path=pipeline_path, + ) labels: set[str] = set() def walk(steps: list | None) -> None: @@ -382,6 +415,17 @@ def walk(steps: list | None) -> None: return labels +def _iter_steps(doc: dict): + def walk(steps: list | None): + for step in steps or []: + if not isinstance(step, dict): + continue + yield step + yield from walk(step.get("steps")) + + yield from walk(doc.get("steps")) + + # Synthetic coverage-style job: shared inputs that change what the split measures. _COVERAGE_SHARED_INPUTS_DOC = { "steps": [ @@ -426,3 +470,289 @@ def test_coverage_shared_inputs_ignored_for_unrelated_change() -> None: ) assert "Coverage Pilot" not in labels assert "Unrelated Model Test" not in labels + + +def _pipeline_dep_keys(path: Path) -> set[str]: + doc = yaml.safe_load(path.read_text(encoding="utf-8")) + keys: set[str] = set() + + def walk(steps: list | None) -> None: + for step in steps or []: + if not isinstance(step, dict): + continue + deps = step.get("source_file_dependencies") + if isinstance(deps, str) and "/" not in deps: + keys.add(deps) + elif isinstance(deps, list) and deps and all("/" not in item for item in deps): + keys.update(deps) + walk(step.get("steps")) + + walk((doc or {}).get("steps")) + return keys + + +def test_pipeline_source_file_dependency_keys_are_registered() -> None: + _load_source_file_dependencies.cache_clear() + registry = _load_source_file_dependencies() + used = ( + _pipeline_dep_keys(Path(".buildkite/cuda/test-ready.yml")) + | _pipeline_dep_keys(Path(".buildkite/cuda/test-merge.yml")) + | _pipeline_dep_keys(Path(".buildkite/cuda/test-nightly.yml")) + | _pipeline_dep_keys(Path(".buildkite/cuda/test-weekly.yml")) + | _pipeline_dep_keys(Path(".buildkite/npu/test-npu-nightly.yml")) + ) + missing = used - set(registry) + assert not missing, f"unregistered source_file_dependencies keys: {sorted(missing)}" + + +def test_source_file_dependencies_key_expands_from_registry() -> None: + doc = { + "steps": [ + { + "label": "Diffusion · Qwen Image Test", + "source_file_dependencies": "diffusion_qwen_image_function", + "commands": ["pytest"], + }, + ], + } + assert "Diffusion · Qwen Image Test" in _surviving_labels( + doc, + ["vllm_omni/diffusion/models/qwen_image/transformer.py"], + ) + assert "Diffusion · Qwen Image Test" not in _surviving_labels(doc, ["vllm_omni/unrelated.py"]) + + +def test_registry_lists_pytest_targets() -> None: + _load_source_file_dependencies.cache_clear() + resolved = _resolve_source_file_dependencies( + { + "label": "Diffusion · Wan22 Test", + "source_file_dependencies": "diffusion_wan22_function", + "commands": [ + "pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py " + "tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model'", + ], + }, + ) + assert resolved is not None + assert "tests/e2e/offline_inference/test_wan22_t2v.py" in resolved + assert "tests/e2e/online_serving/test_wan22_t2v.py" in resolved + assert "vllm_omni/diffusion/models/wan2_2/" in resolved + + +def test_coverage_key_lists_offline_online_scripts() -> None: + resolved = _resolve_source_file_dependencies( + { + "label": "TTS · Qwen3-TTS Base Test", + "source_file_dependencies": "tts_qwen3_tts_cov", + "commands": [ + ".buildkite/common/scripts/run_cov_split.sh \\\n" + " --offline tests/e2e/offline_inference/test_qwen3_tts_base.py \\\n" + " --online tests/e2e/online_serving/test_qwen3_tts_base.py", + ], + }, + ) + assert resolved is not None + assert "tests/e2e/offline_inference/test_qwen3_tts_base.py" in resolved + assert "tests/e2e/online_serving/test_qwen3_tts_base.py" in resolved + assert ".buildkite/common/scripts/run_cov_split.sh" in resolved + assert "pyproject.toml" in resolved + + +def test_source_file_dependencies_list_of_keys_concatenates() -> None: + resolved = _resolve_source_file_dependencies( + { + "label": "composed", + "source_file_dependencies": ["omni_qwen3_omni_function", "tts_qwen3_tts_function"], + }, + ) + assert resolved is not None + assert "vllm_omni/model_executor/models/qwen3_omni/" in resolved + assert "vllm_omni/model_executor/models/qwen3_tts/" in resolved + assert resolved.count("vllm_omni/model_executor/models/common/snake_activation.py") == 1 + + +def test_unknown_source_file_dependencies_key() -> None: + with pytest.raises(ValueError, match="unknown source_file_dependencies"): + _render_test_pipeline( + {"steps": [{"label": "bad", "source_file_dependencies": "not_a_real_key"}]}, + changed_files=None, + ) + + +def test_source_file_dependencies_rejects_mixed_keys_and_paths() -> None: + with pytest.raises(ValueError, match="mixes registry keys and path prefixes"): + _resolve_source_file_dependencies( + { + "label": "bad", + "source_file_dependencies": ["omni_qwen3_omni_function", "tests/e2e/online_serving/test_qwen3_omni.py"], + }, + ) + + +# Synthetic pipeline: selection depends only on listed deps, not on live job names. +_SOURCE_FILTER_DOC = { + "steps": [ + { + "group": "E2E Tests", + "if": 'build.env("NON_CRITICAL") == "1"', + "steps": [ + { + "label": "Dedicated E2E", + "commands": ["pytest -sv tests/e2e/online_serving/test_magi2.py"], + }, + ], + }, + { + "label": "Omni Sweep", + "source_file_dependencies": ["tests/e2e/online_serving/test_qwen3_omni.py"], + "commands": ["pytest -sv tests/e2e/ -m omni"], + }, + { + "label": "Z-Image Function", + "source_file_dependencies": [ + "tests/e2e/online_serving/test_zimage_expansion.py", + "vllm_omni/diffusion/models/z_image/", + ], + "mirror_hardwares": "l4_4", + "commands": ["pytest -sv tests/e2e/online_serving/test_zimage_expansion.py"], + }, + { + "label": "Tiny Model", + "source_file_dependencies": ["tests/e2e/online_serving/test_tiny.py"], + "commands": ["pytest -sv tests/e2e/ -m tiny"], + }, + { + "label": "Wan Function", + "source_file_dependencies": [ + "tests/e2e/offline_inference/test_wan22_t2v.py", + "tests/e2e/online_serving/test_wan22_t2v.py", + "vllm_omni/diffusion/models/wan2_2/", + ], + "commands": ["pytest -sv tests/e2e/offline_inference/test_wan22_t2v.py"], + }, + { + "label": "Wan Perf", + "source_file_dependencies": ["vllm_omni/diffusion/models/wan2_2/"], + "commands": ["pytest -sv tests/dfx/perf/scripts/run_benchmark.py"], + }, + { + "label": "Doc Test", + "source_file_dependencies": [ + "tests/examples/offline_inference/test_text_to_image.py", + "tests/examples/online_serving/test_text_to_image.py", + "vllm_omni/diffusion/models/qwen_image/", + "vllm_omni/diffusion/models/z_image/", + ], + "commands": ["pytest -sv tests/examples/*/test_text_to_image.py"], + }, + ], +} + + +def test_source_filter_ignores_unrelated_e2e_file() -> None: + """A sweep command is not a dependency; only listed paths select a keyed job.""" + labels = _surviving_labels(_SOURCE_FILTER_DOC, ["tests/e2e/online_serving/test_magi2.py"]) + assert labels == {"Dedicated E2E"} + + +def test_source_filter_selects_job_by_listed_script_not_sibling() -> None: + labels = _surviving_labels(_SOURCE_FILTER_DOC, ["tests/e2e/online_serving/test_zimage_expansion.py"]) + assert labels == {"Dedicated E2E", "Z-Image Function"} + + +def test_source_filter_selects_every_job_sharing_a_path() -> None: + labels = _surviving_labels(_SOURCE_FILTER_DOC, ["vllm_omni/diffusion/models/wan2_2/transformer.py"]) + assert labels == {"Dedicated E2E", "Wan Function", "Wan Perf"} + + +def test_source_filter_selects_composed_doc_job_by_example_or_model() -> None: + for changed in ( + "tests/examples/online_serving/test_text_to_image.py", + "vllm_omni/diffusion/models/z_image/transformer.py", + "vllm_omni/diffusion/models/qwen_image/foo.py", + ): + labels = _surviving_labels(_SOURCE_FILTER_DOC, [changed]) + assert "Doc Test" in labels, changed + assert "Tiny Model" not in labels + assert "Omni Sweep" not in labels + assert "Doc Test" not in _surviving_labels(_SOURCE_FILTER_DOC, ["tests/e2e/online_serving/test_magi2.py"]) + + +def test_source_filter_strips_deps_and_expands_hardware(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("upload_pipeline._get_mirror_hw_selector", lambda: "") + rendered = _render_test_pipeline( + _SOURCE_FILTER_DOC, + changed_files=["tests/e2e/online_serving/test_zimage_expansion.py"], + ) + dumped = yaml.safe_dump(rendered) + assert "source_file_dependencies" not in dumped + assert "mirror_hardwares" not in dumped + z_image = next(step for step in _iter_steps(rendered) if step.get("label") == "Z-Image Function") + assert z_image["agents"]["queue"] == "l4-k8s" + + +def test_source_filter_disabled_on_main_branch(monkeypatch: pytest.MonkeyPatch) -> None: + class _Ctx: + changed_files = ["vllm_omni/unrelated.py"] + + monkeypatch.setenv("BUILDKITE_BRANCH", "main") + assert _changed_files_for_source_filter(_Ctx(), force_all=False, e2e_only=False) is None + + monkeypatch.setenv("BUILDKITE_BRANCH", "feat/source-filter") + assert _changed_files_for_source_filter(_Ctx(), force_all=False, e2e_only=False) == [ + "vllm_omni/unrelated.py", + ] + monkeypatch.setenv("BUILDKITE_BRANCH", "main") + assert _changed_files_for_source_filter(_Ctx(), force_all=True, e2e_only=False) is None + + +def test_source_filter_fallback_is_fallback_when_no_job_key_matches( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Fallback only when no listed prefix matched; a matching prefix still wins.""" + monkeypatch.setattr("upload_pipeline._get_mirror_hw_selector", lambda: "") + monkeypatch.setenv("BUILDKITE_BRANCH", "feat/nightly-yaml") + pipeline_yaml = Path(".buildkite/npu/test-npu-nightly.yml") + shared_paths = _load_source_file_dependencies()["source_filter_fallback"] + assert ".buildkite/common/scripts/upload_pipeline.py" in shared_paths + + # Synthetic steps only — do not pin live Buildkite job labels. + doc = { + "steps": [ + {"key": "ungated", "commands": ["true"]}, + { + "key": "gated_a", + "source_file_dependencies": ["pkg/model_a/"], + "commands": ["true"], + }, + { + "key": "gated_b", + "source_file_dependencies": ["pkg/model_b/"], + "commands": ["true"], + }, + ], + } + + def surviving_keys(changed_files: list[str]) -> set[str]: + rendered = _render_test_pipeline( + doc, + changed_files=changed_files, + pipeline_path=pipeline_yaml, + ) + return {step["key"] for step in _iter_steps(rendered) if isinstance(step.get("key"), str)} + + # Only pipeline YAML / shared uploader → no listed prefix match → keep every step. + for changed in [pipeline_yaml.as_posix(), *shared_paths]: + assert surviving_keys([changed]) == {"ungated", "gated_a", "gated_b"}, changed + + # A different pipeline YAML is not a fallback for this upload. + assert surviving_keys([".buildkite/cuda/test-nightly.yml"]) == {"ungated"} + + # A matching source prefix wins over fallback files in the same diff. + assert surviving_keys( + [ + "pkg/model_a/transformer.py", + ".buildkite/common/scripts/upload_pipeline.py", + ], + ) == {"ungated", "gated_a"} diff --git a/tests/e2e/accuracy/test_hidream_o1_image.py b/tests/e2e/accuracy/test_hidream_o1_image.py index 923e53a4789..440050be1bf 100644 --- a/tests/e2e/accuracy/test_hidream_o1_image.py +++ b/tests/e2e/accuracy/test_hidream_o1_image.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 __future__ import annotations @@ -16,7 +16,7 @@ from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs from vllm_omni.inputs.data import OmniDiffusionSamplingParams -pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] +pytestmark = [pytest.mark.slow, pytest.mark.diffusion] MODEL_ID = "HiDream-ai/HiDream-O1-Image" PROMPT = "A cat is sitting next to a sign that says 'HiDream-O1 vLLM-Omni'" diff --git a/tests/e2e/offline_inference/test_lingbot_world_v2.py b/tests/e2e/offline_inference/test_lingbot_world_v2.py index 64578af01f9..88e5979dcf2 100644 --- a/tests/e2e/offline_inference/test_lingbot_world_v2.py +++ b/tests/e2e/offline_inference/test_lingbot_world_v2.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 """End-to-end smoke test for LingBot-World v2 generation.""" from __future__ import annotations @@ -37,7 +37,7 @@ ) pytestmark = [ - pytest.mark.slow, + pytest.mark.full_model, pytest.mark.diffusion, pytest.mark.skipif( not _HAS_INPUT_ASSETS, diff --git a/tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_fl2va.py b/tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_fl2va.py index 24cc2f7e9bf..7f14e403dfc 100644 --- a/tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_fl2va.py +++ b/tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_fl2va.py @@ -12,7 +12,7 @@ from ._common import assert_h3_video, dlo_params, run_dlo_wave, run_fl2va -pytestmark = [pytest.mark.core_model, pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.slow] +pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion] H100_TWO_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) diff --git a/tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_ref2va.py b/tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_ref2va.py index f988a44318f..2c61d2a3cec 100644 --- a/tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_ref2va.py +++ b/tests/e2e/online_serving/minimax_h3/test_minimax_h3_dlo_dp2_ref2va.py @@ -12,7 +12,7 @@ from ._common import assert_h3_video, dlo_params, run_dlo_wave, run_ref2va -pytestmark = [pytest.mark.core_model, pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.slow] +pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion] H100_TWO_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) diff --git a/tests/e2e/online_serving/minimax_h3/test_minimax_h3_fasth3_lora.py b/tests/e2e/online_serving/minimax_h3/test_minimax_h3_fasth3_lora.py index 9fc3e8e31c1..1fd80454b60 100644 --- a/tests/e2e/online_serving/minimax_h3/test_minimax_h3_fasth3_lora.py +++ b/tests/e2e/online_serving/minimax_h3/test_minimax_h3_fasth3_lora.py @@ -21,7 +21,7 @@ post_sync, ) -pytestmark = [pytest.mark.core_model, pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.slow] +pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion] H100_TWO_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) diff --git a/tests/e2e/online_serving/minimax_h3/test_minimax_h3_turbo_lora.py b/tests/e2e/online_serving/minimax_h3/test_minimax_h3_turbo_lora.py index 0190598d5fa..8a1a6542a88 100644 --- a/tests/e2e/online_serving/minimax_h3/test_minimax_h3_turbo_lora.py +++ b/tests/e2e/online_serving/minimax_h3/test_minimax_h3_turbo_lora.py @@ -20,7 +20,7 @@ turbo_params, ) -pytestmark = [pytest.mark.core_model, pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.slow] +pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion] H100_TWO_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) diff --git a/tests/e2e/online_serving/test_lingbot_video.py b/tests/e2e/online_serving/test_lingbot_video.py index 8f0f45639cd..6ada319f5dd 100644 --- a/tests/e2e/online_serving/test_lingbot_video.py +++ b/tests/e2e/online_serving/test_lingbot_video.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 """Online serving smoke coverage for the dense LingBot-Video checkpoint.""" @@ -41,7 +41,7 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.slow +@pytest.mark.full_model @pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: @@ -63,7 +63,7 @@ def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientH ) -@pytest.mark.slow +@pytest.mark.full_model @pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) def test_video_generation_modes_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: diff --git a/tests/e2e/online_serving/test_lingbot_video_moe.py b/tests/e2e/online_serving/test_lingbot_video_moe.py index 7e4d1656820..a75fa78e0df 100644 --- a/tests/e2e/online_serving/test_lingbot_video_moe.py +++ b/tests/e2e/online_serving/test_lingbot_video_moe.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 """Single-GPU serving smoke for ``robbyant/lingbot-video-moe-30b-a3b``.""" @@ -41,7 +41,7 @@ def _get_server_cases(model: str): ] -@pytest.mark.slow +@pytest.mark.full_model @pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_server_cases(MODEL), indirect=True) def test_text_to_image_moe(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: @@ -63,7 +63,7 @@ def test_text_to_image_moe(omni_server: OmniServer, openai_client: OpenAIClientH ) -@pytest.mark.slow +@pytest.mark.full_model @pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_server_cases(MODEL), indirect=True) def test_video_generation_modes_moe(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: diff --git a/tests/e2e/online_serving/test_magi2.py b/tests/e2e/online_serving/test_magi2.py index 22a30b495e0..d713abab9f6 100644 --- a/tests/e2e/online_serving/test_magi2.py +++ b/tests/e2e/online_serving/test_magi2.py @@ -8,7 +8,7 @@ From ``tests/``:: - pytest -s -v e2e/online_serving/test_magi2.py -m "advanced_model and diffusion" --run-level=advanced_model + pytest -s -v e2e/online_serving/test_magi2.py -m "slow and diffusion" --run-level=full_model """ import json @@ -27,8 +27,7 @@ FOUR_CARD_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=4) -@pytest.mark.core_model -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", diff --git a/tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py b/tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py index 0bf5466831e..4fa71aead15 100644 --- a/tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py +++ b/tests/e2e/online_serving/test_minimax_h3_dlo_dp2_t2va.py @@ -111,7 +111,6 @@ def _run_t2va_request(client: OpenAIClientHandler, seed: int) -> bytes: @pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.diffusion -@pytest.mark.slow @pytest.mark.parametrize( "omni_server", [ diff --git a/tools/pre_commit/check_buildkite.py b/tools/pre_commit/check_buildkite.py index 3622534941c..f470f09f921 100644 --- a/tools/pre_commit/check_buildkite.py +++ b/tools/pre_commit/check_buildkite.py @@ -27,6 +27,7 @@ SKIP_FILES = { ".buildkite/common/ci_mirror_hardwares.yml", + ".buildkite/common/ci_source_file_dependencies.yml", ".buildkite/npu/pipeline-npu-a3.yml", } From a2486dd00bcc4f6206ebfe4c99f40e8a784e8d06 Mon Sep 17 00:00:00 2001 From: NancyFyong <88076188+NancyFyong@users.noreply.github.com> Date: Thu, 10 Sep 2026 17:50:54 +0800 Subject: [PATCH 52/93] [Core][Diffusion] Add a typed pre-D2H video media contract (#6615) Signed-off-by: NancyFyong Signed-off-by: Samit <285365963@qq.com> Co-authored-by: NancyFyong Co-authored-by: Samit <285365963@qq.com> --- .buildkite/cuda/test-nightly.yml | 8 + docs/.nav.yml | 1 + docs/mkdocs/hooks/generate_argparse.py | 2 + .../device_side_video_postprocess.md | 72 +++ tests/config/test_omni_config.py | 55 ++ .../wan2_2/test_wan22_pipeline_diffuse.py | 111 ++++ tests/diffusion/test_async_output_worker.py | 136 ++++- .../test_device_output_preparation.py | 460 +++++++++++++++ tests/diffusion/test_device_reduction.py | 125 ++++ .../diffusion/test_diffusion_model_runner.py | 114 +++- tests/diffusion/test_ipc_async.py | 30 +- tests/diffusion/test_media_output_contract.py | 546 ++++++++++++++++++ ...test_video_transport_defaults_are_inert.py | 61 ++ tests/diffusion/test_wan_device_reduction.py | 348 +++++++++++ tests/docs/test_generate_argparse.py | 24 + .../test_device_postprocess_equivalence.py | 105 ++++ tests/entrypoints/test_serve.py | 24 + vllm_omni/config/omni_config.py | 9 + vllm_omni/config/stage_config.py | 1 + vllm_omni/diffusion/data.py | 37 +- vllm_omni/diffusion/diffusion_engine.py | 46 +- vllm_omni/diffusion/ipc.py | 257 ++++++++- vllm_omni/diffusion/media.py | 187 ++++++ .../models/wan2_2/pipeline_wan2_2.py | 32 +- .../diffusion/postprocess/device_reduction.py | 152 +++++ vllm_omni/diffusion/postprocess/media.py | 92 +++ .../worker/diffusion_model_runner.py | 22 + .../diffusion/worker/diffusion_worker.py | 14 +- vllm_omni/diffusion/worker/request_batch.py | 24 +- vllm_omni/engine/arg_utils.py | 1 + vllm_omni/entrypoints/cli/serve.py | 18 + 31 files changed, 3054 insertions(+), 60 deletions(-) create mode 100644 docs/user_guide/diffusion/device_side_video_postprocess.md create mode 100644 tests/diffusion/test_device_output_preparation.py create mode 100644 tests/diffusion/test_device_reduction.py create mode 100644 tests/diffusion/test_media_output_contract.py create mode 100644 tests/diffusion/test_video_transport_defaults_are_inert.py create mode 100644 tests/diffusion/test_wan_device_reduction.py create mode 100644 tests/docs/test_generate_argparse.py create mode 100644 tests/e2e/offline_inference/test_device_postprocess_equivalence.py create mode 100644 vllm_omni/diffusion/media.py create mode 100644 vllm_omni/diffusion/postprocess/device_reduction.py create mode 100644 vllm_omni/diffusion/postprocess/media.py diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index ee4fc6f3ea4..583f935dc25 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -389,6 +389,14 @@ steps: - pytest -s -v tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model" mirror_hardwares: h100_2 + - label: ":full_moon: Diffusion X2V · Wan2.2 Device-Postprocess Equivalence Test" + timeout_in_minutes: 90 + commands: + - | + VLLM_OMNI_DEVICE_POSTPROCESS_MODEL=Wan-AI/Wan2.2-TI2V-5B-Diffusers \ + pytest -s -v tests/e2e/offline_inference/test_device_postprocess_equivalence.py -m "full_model and cuda" --run-level "full_model" + mirror_hardwares: h100_1 + - label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test" source_file_dependencies: diffusion_wan22_accuracy timeout_in_minutes: 180 diff --git a/docs/.nav.yml b/docs/.nav.yml index aa84f28199e..8dac28d2233 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -71,6 +71,7 @@ nav: - Overview and Compatibility: - Overview: user_guide/diffusion_features.md - Feature Compatibility: user_guide/feature_compatibility.md + - Device-side Video Postprocessing: user_guide/diffusion/device_side_video_postprocess.md - CPU Offloading: - Overview: user_guide/diffusion/cpu_offload.md - Model-Level: user_guide/diffusion/offloader/module_offload.md diff --git a/docs/mkdocs/hooks/generate_argparse.py b/docs/mkdocs/hooks/generate_argparse.py index 4c3ad3373b4..ea0fc9e0cc8 100644 --- a/docs/mkdocs/hooks/generate_argparse.py +++ b/docs/mkdocs/hooks/generate_argparse.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import ast +import json import logging import re import sys @@ -124,6 +125,7 @@ def add_parser(self, name, **kwargs): "make_arg_parser": lambda parser: parser, # no-op for doc "_ensure_vllm_platform": lambda: None, # no-op for doc "_nonneg_finite_float": float, + "_json_object": json.loads, "_parse_stage_overrides": __import__("json").loads, "VLLM_SUBCMD_PARSER_EPILOG": "", "logger": logger, diff --git a/docs/user_guide/diffusion/device_side_video_postprocess.md b/docs/user_guide/diffusion/device_side_video_postprocess.md new file mode 100644 index 00000000000..910e644570f --- /dev/null +++ b/docs/user_guide/diffusion/device_side_video_postprocess.md @@ -0,0 +1,72 @@ +# Device-side video postprocessing + +vLLM-Omni can convert decoded WAN2.2 video from normalized floating-point +`[B, C, T, H, W]` tensors to contiguous `uint8` `[B, T, H, W, C]` frames before +the worker-to-engine device-to-host transfer. This reduces the transported video +payload by four times when the original IPC path widens bfloat16 to float32. + +The optimization is disabled by default. Enable it with a structured diffusion +configuration: + +```python +from vllm_omni import Omni + +engine = Omni( + model="Wan-AI/Wan2.2-TI2V-5B-Diffusers", + video_output_transport={"enable_device_postprocess": True}, +) +``` + +For the server CLI, pass the same object as JSON: + +```bash +vllm serve Wan-AI/Wan2.2-TI2V-5B-Diffusers --omni \ + --video-output-transport '{"enable_device_postprocess": true}' +``` + +For a deploy configuration, place it on the diffusion stage: + +```yaml +stages: + - stage_id: 0 + video_output_transport: + enable_device_postprocess: true +``` + +## Runtime contract + +WAN2.2 emits decoded video through the typed `DiffusionOutput.media` field. The +model runner validates and splits the batch into request-local tensors before +preparing each video for transport. The worker then performs the D2H copy and +shared-memory IPC using the prepared representation. The engine uses the generic +media finalizer instead of a model-specific postprocessor. + +The runtime leaves the video in normalized floating-point form when: + +- `enable_device_postprocess` is false; +- the requested output type is not `np`; or +- frame interpolation still requires the floating-point tensor. + +These are policy fallbacks, not malformed-contract errors. If the device cannot +allocate the temporary conversion buffer, the runner logs a warning and prepares +the request-local float representation instead. Invalid tensor layout, encoding, +value range, lifecycle state, and non-memory runtime failures still fail before +IPC. + +## Precision + +Device preparation converts to float32 before denormalization and quantization. +For bfloat16 WAN output this is more precise than the historical host path, +which denormalizes in bfloat16. Relative to that path, values may differ by at +most one level out of 255. Float32 inputs remain byte-identical after +quantization. + +## Scope + +WAN2.2 is the reference migration for the first contract version. Other video +pipelines continue using their existing output paths until they can declare the +same typed media contract without bypassing model-specific float consumers such +as safety checks or audio/video packaging. + +See [RFC #6541](https://github.com/vllm-project/vllm-omni/issues/6541) for the +contract, lifecycle, batching rules, and migration plan. diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 40229f52154..da8ffb4525e 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -1542,6 +1542,61 @@ def test_diffusion_quantization_mapping_reaches_terminal_config(monkeypatch): assert cfg.quantization_config.get_name() == "int8" +def test_video_output_transport_mapping_is_normalized() -> None: + from vllm_omni.diffusion.data import VideoOutputTransportConfig + + cfg = omni_config_module._DiffusionConfigProjection( + video_output_transport={"enable_device_postprocess": True}, + ) + + assert isinstance(cfg.video_output_transport, VideoOutputTransportConfig) + assert cfg.video_output_transport.enable_device_postprocess is True + + +def test_omni_diffusion_config_normalizes_video_output_transport_mapping() -> None: + from vllm_omni.diffusion.data import OmniDiffusionConfig, VideoOutputTransportConfig + + cfg = OmniDiffusionConfig( + model=None, + video_output_transport={"enable_device_postprocess": True}, + ) + + assert isinstance(cfg.video_output_transport, VideoOutputTransportConfig) + assert cfg.video_output_transport.enable_device_postprocess is True + + +def test_video_output_transport_rejects_non_boolean_flag() -> None: + from vllm_omni.diffusion.data import VideoOutputTransportConfig + + with pytest.raises(TypeError, match="enable_device_postprocess must be a bool"): + VideoOutputTransportConfig(enable_device_postprocess="true") # type: ignore[arg-type] + + +def test_video_output_transport_survives_stage_override_filtering() -> None: + from vllm_omni.config.stage_config import build_stage_runtime_overrides, deploy_runtime_override_keys + + transport = {"enable_device_postprocess": True} + overrides = build_stage_runtime_overrides(0, {"video_output_transport": transport}) + + assert "video_output_transport" in deploy_runtime_override_keys() + assert overrides["video_output_transport"] == transport + + +def test_video_output_transport_reaches_default_diffusion_stage() -> None: + from vllm_omni.engine.async_omni_engine import AsyncOmniEngine + + transport = {"enable_device_postprocess": True} + stages = AsyncOmniEngine._create_default_diffusion_stage_cfg( + { + "model": "unused", + "model_class_name": "UnknownPipeline", + "video_output_transport": transport, + } + ) + + assert stages[0]["engine_args"]["video_output_transport"] == transport + + def test_compact_offload_config_reaches_terminal_config(monkeypatch): from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.offloader.config import OffloadStrategy, resolve_offload_strategy diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py index de5abeab0ed..6f7833ff820 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py +++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py @@ -9,6 +9,7 @@ import torch from torch import nn +from vllm_omni.diffusion.media import VideoTensorEncoding, VideoTensorLayout, VideoValueRange from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention from vllm_omni.diffusion.request import OmniDiffusionRequest @@ -30,6 +31,22 @@ def dtype(self) -> torch.dtype: return torch.float32 +class _StubVaeConfig: + latents_mean = [0.0, 0.0, 0.0, 0.0] + latents_std = [1.0, 1.0, 1.0, 1.0] + z_dim = 4 + + +class _StubVae(nn.Module): + dtype = torch.float32 + config = _StubVaeConfig() + + def decode(self, latents, return_dict=False): + del return_dict + batch, _, frames, height, width = latents.shape + return (torch.zeros(batch, 3, frames, height, width),) + + class _StubScheduler: def __init__(self, timesteps: list[int]) -> None: self.timesteps = torch.tensor(timesteps, dtype=torch.int64) @@ -75,6 +92,7 @@ def _make_pipeline() -> Wan22Pipeline: pipeline.transformer = _StubTransformer() pipeline.transformer_2 = None pipeline.text_encoder = _StubTextEncoder() + pipeline.vae = _StubVae() pipeline.transformer_config = SimpleNamespace(patch_size=(1, 2, 2), in_channels=4, out_channels=4) pipeline.scheduler = _StubScheduler([9, 5]) pipeline.od_config = SimpleNamespace(flow_shift=5.0) @@ -227,6 +245,99 @@ def _fake_prepare_latents(**kwargs): torch.testing.assert_close(outputs[1].output, latents_b) +def test_forward_emits_request_local_typed_media_after_vae_decode() -> None: + pipeline = _make_pipeline() + + def _fake_diffuse( + *, + latents, + timesteps, + prompt_embeds, + negative_prompt_embeds, + guidance_low, + guidance_high, + boundary_timestep, + dtype, + attention_kwargs, + latent_condition, + first_frame_mask, + generator, + ): + del ( + timesteps, + prompt_embeds, + negative_prompt_embeds, + guidance_low, + guidance_high, + boundary_timestep, + dtype, + attention_kwargs, + latent_condition, + first_frame_mask, + generator, + ) + return torch.zeros_like(latents) + + pipeline.diffuse = _fake_diffuse # type: ignore[method-assign] + batch = DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest( + prompt="prompt", + request_id="request-0", + sampling_params=OmniDiffusionSamplingParams( + num_frames=1, + num_inference_steps=2, + max_sequence_length=32, + output_type="np", + ), + ) + ] + ) + + outputs = pipeline.forward(batch) + + assert len(outputs) == 1 + assert outputs[0].output is None + assert outputs[0].media is not None + assert outputs[0].media.prepared_for_transport is False + assert outputs[0].media.video.tensor.shape == (1, 3, 1, 8, 8) + assert outputs[0].media.video.spec.layout is VideoTensorLayout.BCTHW + assert outputs[0].media.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT + assert outputs[0].media.video.spec.value_range is VideoValueRange.NEGATIVE_ONE_TO_ONE + + +def test_forward_keeps_legacy_output_on_non_owner_vae_rank() -> None: + # Distributed VAE decode uses broadcast_result=False, so non-owner ranks get + # an empty placeholder instead of the full video. Wrapping that as typed media + # would fail split_diffusion_output_by_request's batch check on every non-owner + # rank, so the pipeline must keep the placeholder on the legacy output field. + pipeline = _make_pipeline() + pipeline.vae.decode = lambda latents, return_dict=False: (torch.empty(0),) # type: ignore[assignment] + pipeline.diffuse = lambda **kwargs: torch.zeros_like(kwargs["latents"]) # type: ignore[method-assign] + + batch = DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest( + prompt="prompt", + request_id="request-0", + sampling_params=OmniDiffusionSamplingParams( + num_frames=1, + num_inference_steps=2, + max_sequence_length=32, + output_type="np", + ), + ) + ] + ) + + outputs = pipeline.forward(batch) + + assert len(outputs) == 1 + assert outputs[0].media is None + assert outputs[0].output is not None + assert outputs[0].output.numel() == 0 + + def test_forward_batches_precomputed_prompt_embeddings() -> None: pipeline = _make_pipeline() diffuse_call = {} diff --git a/tests/diffusion/test_async_output_worker.py b/tests/diffusion/test_async_output_worker.py index a4bff4625f5..8a81f30d4b3 100644 --- a/tests/diffusion/test_async_output_worker.py +++ b/tests/diffusion/test_async_output_worker.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 """Unit tests for WorkerProc async output logic.""" import queue @@ -8,13 +8,23 @@ from unittest.mock import MagicMock import pytest +import torch +import vllm_omni.diffusion.ipc as diffusion_ipc from vllm_omni.diffusion.data import ( AsyncDiffusionOutput, AsyncOutputKind, DiffusionOutput, OmniACK, ) +from vllm_omni.diffusion.media import ( + DiffusionMediaOutput, + VideoMediaOutput, + VideoTensorEncoding, + VideoTensorLayout, + VideoTensorSpec, + VideoValueRange, +) from vllm_omni.diffusion.worker.diffusion_worker import WorkerProc from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput @@ -182,6 +192,66 @@ def test_diffusion_output_uses_sync_path_when_step_execution(self, mocker): mock_pack.assert_called_once_with(output) proc.result_mq.enqueue.assert_called_once_with(output) + def test_media_contract_error_is_not_sent_as_an_unpacked_result(self, mocker): + proc = _make_worker_proc(step_execution=True) + proc._async_output_queue = None + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm", + side_effect=ValueError("media was not prepared"), + ) + + with pytest.raises(ValueError, match="media was not prepared"): + proc._return_result(DiffusionOutput()) + + proc.result_mq.enqueue.assert_not_called() + + def test_media_packing_memory_failure_is_reraised_for_typed_media(self, mocker): + """A non-ValueError packing failure on typed media must re-raise, not + swallow it and enqueue a half-packed payload.""" + proc = _make_worker_proc(step_execution=True) + proc._async_output_queue = None + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm", + side_effect=RuntimeError("SHM pack out of memory"), + ) + + output = DiffusionOutput(media=MagicMock(spec=DiffusionMediaOutput)) + with pytest.raises(RuntimeError, match="SHM pack out of memory"): + proc._return_result(output) + + proc.result_mq.enqueue.assert_not_called() + + def test_media_packing_failure_is_reraised_for_dp_container(self, mocker): + """Container dispatch and typed-media failure detection share a walker.""" + proc = _make_worker_proc(step_execution=True) + proc._async_output_queue = None + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm", + side_effect=RuntimeError("SHM pack out of memory"), + ) + + output = DiffusionOutput(media=MagicMock(spec=DiffusionMediaOutput)) + tagged_output = {"dp_rank": 0, "output": output} + with pytest.raises(RuntimeError, match="SHM pack out of memory"): + proc._return_result(tagged_output) + + proc.result_mq.enqueue.assert_not_called() + + def test_legacy_output_packing_failure_is_still_swallowed(self, mocker): + """Non-typed (legacy output) failures keep the historic swallow-and- + enqueue behavior; only typed media re-raises.""" + proc = _make_worker_proc(step_execution=True) + proc._async_output_queue = None + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm", + side_effect=RuntimeError("SHM pack out of memory"), + ) + + proc._return_result(DiffusionOutput(output="data")) + + # Legacy path: warning logged, output still enqueued once. + proc.result_mq.enqueue.assert_called_once_with(DiffusionOutput(output="data")) + def test_non_diffusion_output_skips_async_in_request_mode(self, mocker): """Non-DiffusionOutput/BatchRunnerOutput in request mode uses sync path.""" proc = _make_worker_proc(step_execution=False) @@ -243,6 +313,70 @@ def test_bg_thread_enqueues_output_ready_after_pack(self, mocker): ) assert output_ready_enqueued, "Expected OUTPUT_READY with async_output_id='abc123' in enqueue calls" + def test_late_field_failure_sends_error_without_partial_payload_or_shm_leak( + self, + mocker, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + from multiprocessing import shared_memory + + proc = _make_worker_proc(step_execution=False) + media = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=torch.zeros(1, 3, 2, 4, 5), + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ), + ), + prepared_for_transport=True, + ) + output = DiffusionOutput(media=media, trajectory_latents=torch.ones(30)) + original_latents = output.trajectory_latents + + real_array_to_shm = diffusion_ipc._array_to_shm + created_handles: list[dict[str, object]] = [] + + def _fail_on_second(array): + if created_handles: + raise RuntimeError("late-field packing failure") + handle = real_array_to_shm(array) + created_handles.append(handle) + return handle + + monkeypatch.setattr(diffusion_ipc, "_SHM_TENSOR_THRESHOLD", 1) + monkeypatch.setattr(diffusion_ipc, "_array_to_shm", _fail_on_second) + + def _pack_without_stream(value, d2h_stream=None): + # The worker loop behavior is under test here. CPU tensors avoid + # introducing a device dependency into this unit test. + del d2h_stream + return diffusion_ipc.pack_diffusion_output_shm(value) + + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm", + side_effect=_pack_without_stream, + ) + mock_torch = mocker.MagicMock() + mock_torch.accelerator.current_accelerator.return_value.type = "cpu" + mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.torch", mock_torch) + + proc._async_output_queue.put((output, "late-failure", None)) + proc._async_output_queue.put(None) + proc._async_output_loop() + + messages = [call.args[0] for call in proc.result_mq.enqueue.call_args_list] + assert len(messages) == 1 + assert messages[0].kind is AsyncOutputKind.OUTPUT_READY + assert messages[0].async_output_id == "late-failure" + assert messages[0].error == "Background D2H/SHM packing failed" + assert messages[0].output is None + assert output.media is media + assert output.trajectory_latents is original_latents + with pytest.raises(FileNotFoundError): + shared_memory.SharedMemory(name=str(created_handles[0]["name"])) + class TestDrainAsyncOutputs: """Async outputs must finish packing before device memory is released.""" diff --git a/tests/diffusion/test_device_output_preparation.py b/tests/diffusion/test_device_output_preparation.py new file mode 100644 index 00000000000..0a6f1ab90cb --- /dev/null +++ b/tests/diffusion/test_device_output_preparation.py @@ -0,0 +1,460 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from dataclasses import FrozenInstanceError + +import numpy as np +import pytest +import torch + +import vllm_omni.diffusion.ipc as diffusion_ipc +import vllm_omni.diffusion.postprocess.device_reduction as device_reduction +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, VideoOutputTransportConfig +from vllm_omni.diffusion.ipc import pack_diffusion_output_shm, unpack_diffusion_output_shm +from vllm_omni.diffusion.media import ( + DiffusionMediaOutput, + FloatVideoConsumer, + VideoMediaOutput, + VideoTensorEncoding, + VideoTensorLayout, + VideoTensorSpec, + VideoTransportConstraints, + VideoValueRange, + ensure_request_owned_tensor, + slice_diffusion_media_output, +) +from vllm_omni.diffusion.postprocess.device_reduction import prepare_diffusion_media_for_transport +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _config(*, enabled: bool) -> OmniDiffusionConfig: + return OmniDiffusionConfig( + model=None, + video_output_transport=VideoOutputTransportConfig(enable_device_postprocess=enabled), + ) + + +def _sampling(*, output_type: str | None = "np", interpolate: bool = False) -> OmniDiffusionSamplingParams: + return OmniDiffusionSamplingParams( + output_type=output_type, + enable_frame_interpolation=interpolate, + ) + + +def _media( + tensor: torch.Tensor, + *, + value_range: VideoValueRange = VideoValueRange.NEGATIVE_ONE_TO_ONE, + consumers: frozenset[FloatVideoConsumer] = frozenset(), +) -> DiffusionMediaOutput: + return DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=tensor, + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=value_range, + ), + constraints=VideoTransportConstraints(pending_float_consumers=consumers), + ) + ) + + +def test_runtime_prepares_request_local_video_as_uint8() -> None: + video = torch.linspace(-1.0, 1.0, 2 * 3 * 2 * 4 * 5).reshape(2, 3, 2, 4, 5) + + prepared = prepare_diffusion_media_for_transport( + _media(video), + od_config=_config(enabled=True), + sampling_params=_sampling(), + ) + + assert prepared.prepared_for_transport is True + assert prepared.video.tensor.dtype == torch.uint8 + assert prepared.video.tensor.shape == (2, 2, 4, 5, 3) + assert prepared.video.spec == VideoTensorSpec( + layout=VideoTensorLayout.BTHWC, + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + ) + assert prepared.video.constraints.pending_float_consumers == frozenset() + + +def test_disabled_runtime_preserves_float_but_marks_media_prepared() -> None: + video = torch.randn(1, 3, 2, 4, 5) + + prepared = prepare_diffusion_media_for_transport( + _media(video), + od_config=_config(enabled=False), + sampling_params=_sampling(), + ) + + assert prepared.prepared_for_transport is True + assert prepared.video.tensor is video + assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT + + +def test_device_oom_falls_back_to_request_owned_float_media(monkeypatch: pytest.MonkeyPatch) -> None: + source = torch.randn(2, 3, 2, 4, 5) + sliced = slice_diffusion_media_output(_media(source), 0, 1) + + def raise_oom(video: torch.Tensor, *, do_denormalize: bool) -> torch.Tensor: + raise torch.OutOfMemoryError("injected") + + monkeypatch.setattr(device_reduction, "reduce_video_to_uint8_frames", raise_oom) + + prepared = prepare_diffusion_media_for_transport( + sliced, + od_config=_config(enabled=True), + sampling_params=_sampling(), + ) + + assert prepared.prepared_for_transport is True + assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT + assert prepared.video.tensor.dtype == source.dtype + assert prepared.video.tensor.shape == (1, 3, 2, 4, 5) + assert prepared.video.tensor.is_contiguous() + assert prepared.video.tensor._base is None + + +def test_device_non_oom_error_is_not_swallowed(monkeypatch: pytest.MonkeyPatch) -> None: + def raise_runtime_error(video: torch.Tensor, *, do_denormalize: bool) -> torch.Tensor: + raise RuntimeError("injected") + + monkeypatch.setattr(device_reduction, "reduce_video_to_uint8_frames", raise_runtime_error) + + with pytest.raises(RuntimeError, match="injected"): + prepare_diffusion_media_for_transport( + _media(torch.randn(1, 3, 2, 4, 5)), + od_config=_config(enabled=True), + sampling_params=_sampling(), + ) + + +def test_declared_and_request_float_consumers_are_unioned() -> None: + video = torch.randn(1, 3, 2, 4, 5) + media = _media( + video, + consumers=frozenset({FloatVideoConsumer.VIDEO_GUARDRAILS}), + ) + + prepared = prepare_diffusion_media_for_transport( + media, + od_config=_config(enabled=True), + sampling_params=_sampling(interpolate=True), + ) + + assert prepared.video.tensor is video + assert prepared.video.constraints.pending_float_consumers == frozenset( + { + FloatVideoConsumer.FRAME_INTERPOLATION, + FloatVideoConsumer.VIDEO_GUARDRAILS, + } + ) + + +def test_non_numpy_presentation_preserves_float() -> None: + prepared = prepare_diffusion_media_for_transport( + _media(torch.randn(1, 3, 2, 4, 5)), + od_config=_config(enabled=True), + sampling_params=_sampling(output_type="pil"), + ) + + assert prepared.prepared_for_transport is True + assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT + + +def test_zero_to_one_video_is_not_denormalized() -> None: + video = torch.tensor([[[[[0.0, 0.5, 1.0]]]]]).expand(1, 3, 1, 1, 3) + + prepared = prepare_diffusion_media_for_transport( + _media(video, value_range=VideoValueRange.ZERO_TO_ONE), + od_config=_config(enabled=True), + sampling_params=_sampling(), + ) + + expected = np.rint(video.permute(0, 2, 3, 4, 1).numpy() * 255.0).astype(np.uint8) + np.testing.assert_array_equal(prepared.video.tensor.numpy(), expected) + + +def test_invalid_contract_raises_even_when_feature_is_off() -> None: + media = _media(torch.randn(3, 2, 4, 5)) + + with pytest.raises(ValueError, match="rank 5"): + prepare_diffusion_media_for_transport( + media, + od_config=_config(enabled=False), + sampling_params=_sampling(), + ) + + +@pytest.mark.parametrize( + ("tensor", "spec", "message"), + [ + ( + torch.randn(1, 2, 4, 5, 3), + VideoTensorSpec( + layout=VideoTensorLayout.BTHWC, + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + ), + "torch.uint8", + ), + ( + torch.zeros(1, 2, 4, 5, 3, dtype=torch.uint8), + VideoTensorSpec( + layout=VideoTensorLayout.BTHWC, + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_ONE, + ), + "incompatible value range", + ), + ( + torch.randn(1, 4, 2, 4, 5), + VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ), + "3 channels", + ), + ], +) +def test_invalid_spec_tensor_combinations_are_contract_errors( + tensor: torch.Tensor, + spec: VideoTensorSpec, + message: str, +) -> None: + media = DiffusionMediaOutput(video=VideoMediaOutput(tensor=tensor, spec=spec)) + + with pytest.raises(ValueError, match=message): + media.validate() + + +def test_pipeline_cannot_emit_unprepared_uint8_media() -> None: + media = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=torch.zeros(1, 2, 4, 5, 3, dtype=torch.uint8), + spec=VideoTensorSpec( + layout=VideoTensorLayout.BTHWC, + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + ), + ) + ) + + with pytest.raises(ValueError, match="must emit unprepared media as NORMALIZED_FLOAT"): + prepare_diffusion_media_for_transport( + media, + od_config=_config(enabled=True), + sampling_params=_sampling(), + ) + + +def test_descriptors_are_immutable() -> None: + media = _media(torch.randn(1, 3, 2, 4, 5)) + + with pytest.raises(FrozenInstanceError): + media.prepared_for_transport = True # type: ignore[misc] + with pytest.raises(FrozenInstanceError): + media.video.spec.layout = VideoTensorLayout.BTHWC # type: ignore[misc] + + +def test_float_fallback_detaches_a_request_view_before_transport() -> None: + source = torch.randn(4, 3, 2, 4, 5) + sliced = slice_diffusion_media_output(_media(source), 1, 3) + + prepared = prepare_diffusion_media_for_transport( + sliced, + od_config=_config(enabled=False), + sampling_params=_sampling(), + ) + + tensor = prepared.video.tensor + assert tensor.shape == (2, 3, 2, 4, 5) + assert tensor.is_contiguous() + assert tensor._base is None + assert tensor.untyped_storage().nbytes() == tensor.numel() * tensor.element_size() + assert tensor.untyped_storage().data_ptr() != source.untyped_storage().data_ptr() + + +def test_full_batch_slice_does_not_duplicate_owned_storage() -> None: + source = torch.randn(2, 3, 2, 4, 5) + + sliced = slice_diffusion_media_output(_media(source), 0, 2) + + assert sliced.video.tensor.untyped_storage().data_ptr() == source.untyped_storage().data_ptr() + + +def test_diffusion_output_to_cpu_moves_prepared_typed_media() -> None: + prepared = prepare_diffusion_media_for_transport( + _media(torch.randn(1, 3, 2, 4, 5)), + od_config=_config(enabled=False), + sampling_params=_sampling(), + ) + output = DiffusionOutput(media=prepared, to_cpu=True) + + assert output.media is not None + assert output.media.video.tensor.device.type == "cpu" + + +def test_diffusion_output_to_cpu_rejects_unprepared_media() -> None: + with pytest.raises(ValueError, match="prepared before to_cpu"): + DiffusionOutput(media=_media(torch.randn(1, 3, 2, 4, 5)), to_cpu=True) + + +def test_prepared_media_round_trips_through_typed_ipc(monkeypatch: pytest.MonkeyPatch) -> None: + expected = prepare_diffusion_media_for_transport( + _media(torch.randn(1, 3, 2, 256, 256)), + od_config=_config(enabled=True), + sampling_params=_sampling(), + ) + output = DiffusionOutput(media=expected) + + monkeypatch.setattr(diffusion_ipc, "_SHM_TENSOR_THRESHOLD", 1) + pack_diffusion_output_shm(output) + + assert isinstance(output.media, dict) + assert output.media["__type__"] == "diffusion_media_v1" + assert output.media["video"]["tensor"]["__tensor_shm__"] is True + + unpack_diffusion_output_shm(output) + + assert isinstance(output.media, DiffusionMediaOutput) + assert output.media.prepared_for_transport is True + assert output.media.video.spec == expected.video.spec + torch.testing.assert_close(output.media.video.tensor, expected.video.tensor) + + +def test_float_media_constraints_round_trip_through_typed_ipc(monkeypatch: pytest.MonkeyPatch) -> None: + expected = prepare_diffusion_media_for_transport( + _media(torch.randn(1, 3, 2, 256, 256, dtype=torch.bfloat16)), + od_config=_config(enabled=True), + sampling_params=_sampling(interpolate=True), + ) + output = DiffusionOutput(media=expected) + + monkeypatch.setattr(diffusion_ipc, "_SHM_TENSOR_THRESHOLD", 1) + pack_diffusion_output_shm(output) + unpack_diffusion_output_shm(output) + + assert isinstance(output.media, DiffusionMediaOutput) + assert output.media.video.tensor.dtype is torch.bfloat16 + assert output.media.video.constraints.pending_float_consumers == frozenset({FloatVideoConsumer.FRAME_INTERPOLATION}) + torch.testing.assert_close(output.media.video.tensor, expected.video.tensor) + + +def test_ipc_rejects_media_that_skipped_runner_preparation() -> None: + output = DiffusionOutput(media=_media(torch.randn(1, 3, 2, 4, 5))) + + with pytest.raises(ValueError, match="must be prepared"): + pack_diffusion_output_shm(output) + + +def test_ipc_rejects_unknown_media_schema_version() -> None: + prepared = prepare_diffusion_media_for_transport( + _media(torch.randn(1, 3, 2, 4, 5)), + od_config=_config(enabled=False), + sampling_params=_sampling(), + ) + output = DiffusionOutput(media=prepared) + pack_diffusion_output_shm(output) + assert isinstance(output.media, dict) + output.media["schema_version"] = 2 + + with pytest.raises(ValueError, match="schema version"): + unpack_diffusion_output_shm(output) + + +def test_prepare_rejects_media_a_pipeline_already_prepared(monkeypatch: pytest.MonkeyPatch) -> None: + # A pipeline must emit unprepared media so the runner applies request policy + # (e.g. frame interpolation). Accepting a forged prepared uint8 payload would + # skip the interpolation constraint; interpolator never runs and gets dropped. + forged = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=torch.zeros(1, 2, 4, 5, 3, dtype=torch.uint8), + spec=VideoTensorSpec( + layout=VideoTensorLayout.BTHWC, + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + ), + ), + prepared_for_transport=True, + ) + + with pytest.raises(ValueError, match="already prepared"): + prepare_diffusion_media_for_transport( + forged, + od_config=_config(enabled=True), + sampling_params=_sampling(interpolate=True), + ) + + +def test_prepare_rejects_prepared_media_despite_pending_interpolation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # The forged payload is the exact case the reviewer described: device + # postprocessing enabled plus frame_interpolation, but the media already being + # prepared. It must be rejected, not silently returned with interpolation dropped. + forged = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=torch.zeros(1, 3, 2, 2, 3, dtype=torch.float32), + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ), + ), + prepared_for_transport=True, + ) + + with pytest.raises(ValueError, match="already prepared"): + prepare_diffusion_media_for_transport( + forged, + od_config=_config(enabled=True), + sampling_params=_sampling(interpolate=True), + ) + + +def test_request_owned_tensor_rejects_a_detached_storage_view() -> None: + # ``tensor._base is None`` is not proof of ownership: a detached slice keeps a + # nonzero storage offset and the whole source storage. Serializing it would + # leak every other request's pixels, so it must be cloned before transport. + source = torch.randn(4, 3, 2, 4, 5) + detached_slice = source[1:2].detach() + assert detached_slice.is_contiguous() + assert detached_slice._base is None + assert detached_slice.storage_offset() != 0 + + owned = ensure_request_owned_tensor(detached_slice) + + assert owned.shape == detached_slice.shape + assert owned.storage_offset() == 0 + assert owned.untyped_storage().nbytes() == owned.numel() * owned.element_size() + assert owned.untyped_storage().data_ptr() != source.untyped_storage().data_ptr() + torch.testing.assert_close(owned, detached_slice) + + +def test_prepared_media_rejects_non_compact_request_storage() -> None: + # A zero-offset view of a bigger buffer is contiguous but not compact, so + # prepared transport must refuse it just as it refuses a storage view. + source = torch.randn(2, 3, 2, 4, 5) + zero_offset_view = source[0:1] + media = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=zero_offset_view, + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ), + ), + prepared_for_transport=True, + ) + + with pytest.raises(ValueError, match="request-local storage"): + media.validate() diff --git a/tests/diffusion/test_device_reduction.py b/tests/diffusion/test_device_reduction.py new file mode 100644 index 00000000000..f69bb5eb5da --- /dev/null +++ b/tests/diffusion/test_device_reduction.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""GPU correctness/size tests for device-side video reduction. + +The reduction must be byte-identical to the current production output (SHM +widens bf16->f32, then diffusers postprocess, then the API server's ``*255`` +rounding), otherwise enabling it would silently change pixels. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from tests.helpers.mark import hardware_marks +from vllm_omni.diffusion.postprocess.device_reduction import reduce_video_to_uint8_frames + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.diffusion, + *hardware_marks(res={"cuda": "L4"}, num_cards=1), +] + +_GPU = torch.accelerator.is_available() if hasattr(torch, "accelerator") else torch.cuda.is_available() +requires_gpu = pytest.mark.skipif(not _GPU, reason="device reduction is a GPU path") + + +def _device() -> str: + return "cuda:0" + + +def _reference_uint8(video: torch.Tensor, *, compute_dtype: torch.dtype | None = torch.float32) -> np.ndarray: + """The diffusers postprocess plus the API server's rounding, for the same input. + + ``compute_dtype`` selects the precision the denormalize runs at. + ``torch.float32`` is the precision the device reduction uses; ``None`` keeps + the input's own dtype, which is what the engine's postprocess actually does. + """ + from diffusers.video_processor import VideoProcessor + + source = video if compute_dtype is None else video.to(compute_dtype) + processed = VideoProcessor(vae_scale_factor=8).postprocess_video(source, output_type="np") + scaled = np.clip(processed.astype(np.float32), 0.0, 1.0) * 255.0 + np.rint(scaled, out=scaled) + return scaled.astype(np.uint8) + + +@requires_gpu +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +def test_device_reduction_matches_a_float32_reference(dtype: torch.dtype) -> None: + """The reduction must be the float32 computation, exactly, for every input dtype.""" + torch.manual_seed(0) + # Values slightly outside [-1, 1] so the clamp is actually exercised. + video = (torch.rand(1, 3, 8, 64, 96, device=_device()) * 2.4 - 1.2).to(dtype) + + reference = _reference_uint8(video) + produced = reduce_video_to_uint8_frames(video).cpu().numpy() + + assert produced.shape == reference.shape == (1, 8, 64, 96, 3) + assert produced.dtype == np.uint8 + np.testing.assert_array_equal(produced, reference) + + +@requires_gpu +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_device_reduction_deviates_from_a_narrow_postprocess_by_at_most_one(dtype: torch.dtype) -> None: + """Pin the cost of computing in float32: at most one 255th. + + A postprocess that denormalizes in the VAE's own dtype (WAN, HunyuanVideo, + LTX-2, Cosmos3) does not match this reduction bit for bit; ~20% of pixels + differ by one on a real 480x832x81 WAN generation. Keep that a bound rather + than a surprise, so a regression widening it fails here. + """ + torch.manual_seed(0) + video = (torch.rand(1, 3, 8, 64, 96, device=_device()) * 2.4 - 1.2).to(dtype) + + narrow = _reference_uint8(video, compute_dtype=None) + produced = reduce_video_to_uint8_frames(video).cpu().numpy() + + deviation = np.abs(produced.astype(np.int16) - narrow.astype(np.int16)) + assert deviation.max() <= 1 + # The narrow path really is different, so the bound above is not vacuous. + assert deviation.any() + + +@requires_gpu +def test_device_reduction_shrinks_hop1_payload_4x() -> None: + """The whole point: the D2H payload drops from float32 to uint8 (4x).""" + video = torch.rand(1, 3, 8, 704, 1280, device=_device(), dtype=torch.bfloat16) * 2 - 1 + + # Production transports bf16 widened to float32 over SHM. + widened_f32_bytes = video.numel() * 4 + reduced = reduce_video_to_uint8_frames(video) + reduced_bytes = reduced.numel() * reduced.element_size() + + assert reduced.dtype == torch.uint8 + assert widened_f32_bytes / reduced_bytes == pytest.approx(4.0, rel=1e-6) + + +@requires_gpu +def test_device_reduction_stays_on_device() -> None: + """Reduction must happen before D2H, so the result is still on the GPU.""" + video = torch.rand(1, 3, 4, 32, 48, device=_device(), dtype=torch.bfloat16) * 2 - 1 + reduced = reduce_video_to_uint8_frames(video) + assert reduced.device.type == "cuda" + assert reduced.is_contiguous() + + +@requires_gpu +def test_device_reduction_without_denormalize_matches_already_unit_range() -> None: + """Pipelines with do_normalize=False emit [0,1]; skip the denorm step.""" + video = torch.rand(1, 3, 4, 16, 16, device=_device(), dtype=torch.float32) # already [0,1] + + produced = reduce_video_to_uint8_frames(video, do_denormalize=False).cpu().numpy() + + expected = video.permute(0, 2, 3, 4, 1).cpu().numpy() + expected = np.rint(np.clip(expected, 0, 1) * 255.0).astype(np.uint8) + np.testing.assert_array_equal(produced, expected) + + +def test_device_reduction_rejects_non_5d() -> None: + """A wrong-rank tensor is a programming error, not a silent reshape.""" + with pytest.raises(ValueError, match=r"\[B, C, F, H, W\]"): + reduce_video_to_uint8_frames(torch.rand(3, 8, 64, 64)) diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index 9f8d4dec682..69a9e4395d1 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -10,9 +10,19 @@ import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module from tests.helpers.mark import hardware_test -from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.data import DiffusionOutput, VideoOutputTransportConfig from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend +from vllm_omni.diffusion.media import ( + DiffusionMediaOutput, + VideoMediaOutput, + VideoTensorEncoding, + VideoTensorLayout, + VideoTensorSpec, + VideoValueRange, +) +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput, NewRequestData from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -144,6 +154,19 @@ def _make_request(): ) +def _raw_video_media(batch_size: int = 1, tensor: torch.Tensor | None = None) -> DiffusionMediaOutput: + return DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=tensor if tensor is not None else torch.randn(batch_size, 3, 2, 4, 5), + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ), + ) + ) + + def _make_request_with_params(req_id: str, sampling_params): return SimpleNamespace( request_id=req_id, @@ -214,6 +237,7 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar enable_cache_dit_summary=enable_cache_dit_summary, parallel_config=SimpleNamespace(use_hsdp=False), streaming_output=False, + video_output_transport=VideoOutputTransportConfig(), ) runner.diffusion_kv_backend = DiffusionKVModelRunnerBackend( vllm_config=runner.vllm_config, @@ -546,6 +570,48 @@ def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch): assert fourth.get_request_output("req").finished is True +@pytest.mark.core_model +@pytest.mark.cpu +def test_execute_stepwise_prepares_typed_media(monkeypatch): + runner = _make_runner(cache_backend=None, cache_backend_name=None) + runner.pipeline = _ChunkStepPipeline([DiffusionOutput(media=_raw_video_media())]) + runner.od_config.streaming_output = True + runner.od_config.step_execution = True + runner.od_config.video_output_transport = VideoOutputTransportConfig(enable_device_postprocess=True) + req = OmniDiffusionRequest( + prompt="a prompt", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=4), + request_id="req", + ) + + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0) + scheduler_output = DiffusionSchedulerOutput( + step_id=0, + scheduled_new_reqs=[NewRequestData(request_id="req", req=req)], + scheduled_cached_reqs=CachedRequestData.make_empty(), + finished_req_ids=set(), + num_running_reqs=1, + num_waiting_reqs=0, + ) + + DiffusionModelRunner.execute_stepwise(runner, scheduler_output) + scheduler_output = DiffusionSchedulerOutput( + step_id=1, + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData(request_ids=["req"]), + finished_req_ids=set(), + num_running_reqs=1, + num_waiting_reqs=0, + ) + result = DiffusionModelRunner.execute_stepwise(runner, scheduler_output).get_request_output("req").result + + assert result.media.prepared_for_transport is True + assert result.media.video.spec.encoding is VideoTensorEncoding.UINT8_FRAMES + + @pytest.mark.core_model @pytest.mark.cpu def test_execute_stepwise_streaming_decodes_final_only_pipeline(monkeypatch): @@ -819,6 +885,7 @@ def _make_batch_runner(pipeline): cache_backend="none", enable_cache_dit_summary=False, parallel_config=SimpleNamespace(use_hsdp=False), + video_output_transport=VideoOutputTransportConfig(), ) runner.kv_transfer_manager = SimpleNamespace( receive_multi_kv_cache_distributed=lambda req, cfg_kv_collect_func=None, target_device=None: None, @@ -938,6 +1005,21 @@ def _recording_from_outputs(self, reqs, outputs): assert helper_calls == [(["req-0", "req-1"], ["a", "b"])] +@pytest.mark.core_model +@pytest.mark.cpu +def test_execute_model_batch_prepares_typed_media_before_return(monkeypatch): + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory()) + runner = _make_batch_runner(_BatchPipeline(outputs=[DiffusionOutput(media=_raw_video_media())])) + runner.od_config.video_output_transport = VideoOutputTransportConfig(enable_device_postprocess=True) + + result = DiffusionModelRunner.execute_model_batch(runner, _make_scheduler_output(num_reqs=1), runner.od_config) + + prepared = result.runner_outputs[0].result.media + assert prepared.prepared_for_transport is True + assert prepared.video.spec.encoding is VideoTensorEncoding.UINT8_FRAMES + + @pytest.mark.core_model @pytest.mark.cpu def test_split_diffusion_output_by_request_slices_single_and_multi_request_outputs(): @@ -959,6 +1041,36 @@ def test_split_diffusion_output_by_request_slices_single_and_multi_request_outpu assert single[0].output == ["img-0a", "img-0b"] +@pytest.mark.core_model +@pytest.mark.cpu +def test_split_diffusion_output_by_request_preserves_typed_media(): + reqs = [_make_request(), _make_request()] + batch = DiffusionRequestBatch(requests=reqs) + source = torch.randn(4, 3, 2, 4, 5) + result = DiffusionOutput(media=_raw_video_media(tensor=source)) + + outputs = split_diffusion_output_by_request(result, batch, num_outputs_per_prompt=2) + + assert [output.media.video.tensor.shape for output in outputs] == [ + torch.Size([2, 3, 2, 4, 5]), + torch.Size([2, 3, 2, 4, 5]), + ] + assert all(output.media.prepared_for_transport is False for output in outputs) + torch.testing.assert_close(outputs[0].media.video.tensor, source[:2]) + torch.testing.assert_close(outputs[1].media.video.tensor, source[2:]) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_split_diffusion_output_by_request_validates_media_batch_dimension(): + reqs = [_make_request(), _make_request()] + batch = DiffusionRequestBatch(requests=reqs) + result = DiffusionOutput(media=_raw_video_media(batch_size=3)) + + with pytest.raises(ValueError, match=r"request_count \* num_outputs_per_prompt \(4\), got 3"): + split_diffusion_output_by_request(result, batch, num_outputs_per_prompt=2) + + @pytest.mark.core_model @pytest.mark.cpu def test_split_diffusion_output_by_request_slices_tuple_outputs(): diff --git a/tests/diffusion/test_ipc_async.py b/tests/diffusion/test_ipc_async.py index 03919095b9a..43650c58119 100644 --- a/tests/diffusion/test_ipc_async.py +++ b/tests/diffusion/test_ipc_async.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for IPC async D2H stream packing.""" +from unittest.mock import ANY + import pytest import torch @@ -43,7 +45,7 @@ def test_pack_value_if_large_passes_stream_to_tensor_pack(self, mocker): tensor = torch.zeros(1) d2h = mocker.MagicMock() _pack_value_if_large(tensor, d2h_stream=d2h) - mock_pack.assert_called_once_with(tensor, d2h_stream=d2h) + mock_pack.assert_called_once_with(tensor, d2h_stream=d2h, created=None) def test_pack_value_if_large_passes_stream_through_dict(self, mocker): """d2h_stream propagates through dict container packing.""" @@ -93,7 +95,7 @@ def test_output_field_packed_with_stream(self, mocker): d2h = mocker.MagicMock() out = DiffusionOutput(output=torch.zeros(1)) _pack_diffusion_fields(out, d2h_stream=d2h) - mock_pack_value.assert_called_once_with(torch.zeros(1), d2h_stream=d2h) + mock_pack_value.assert_called_once_with(torch.zeros(1), d2h_stream=d2h, created=None) def test_trajectory_fields_packed_with_stream(self, mocker): mock_pack_tensor = mocker.patch( @@ -120,7 +122,12 @@ def test_diffusion_output_path_with_stream(self, mocker): d2h = mocker.MagicMock() out = DiffusionOutput() pack_diffusion_output_shm(out, d2h_stream=d2h) - mock_pack_fields.assert_called_once_with(out, d2h_stream=d2h) + mock_pack_fields.assert_called_once_with( + out, + d2h_stream=d2h, + created=ANY, + pending_updates=ANY, + ) def test_rpc_envelope_path_with_stream(self, mocker): mock_pack_fields = mocker.patch( @@ -130,11 +137,11 @@ def test_rpc_envelope_path_with_stream(self, mocker): result = DiffusionOutput() envelope = {"type": "diffusion_rpc_result", "result": result} pack_diffusion_output_shm(envelope, d2h_stream=d2h) - # envelope["result"] is replaced in-place by pack_diffusion_output_shm; - # assert with the original reference, not the replaced one. mock_pack_fields.assert_called_once_with( result, d2h_stream=d2h, + created=ANY, + pending_updates=ANY, ) @pytest.mark.parametrize("wrapped_in_rpc_envelope", [False, True]) @@ -147,7 +154,12 @@ def test_dp_tagged_output_path_with_stream(self, mocker, wrapped_in_rpc_envelope tagged = {"dp_rank": 1, "output": result} value = {"type": "diffusion_rpc_result", "result": tagged} if wrapped_in_rpc_envelope else tagged pack_diffusion_output_shm(value, d2h_stream=d2h) - mock_pack_fields.assert_called_once_with(result, d2h_stream=d2h) + mock_pack_fields.assert_called_once_with( + result, + d2h_stream=d2h, + created=ANY, + pending_updates=ANY, + ) def test_runner_output_path_with_stream(self, mocker): mock_pack_fields = mocker.patch( @@ -163,9 +175,9 @@ def test_runner_output_path_with_stream(self, mocker): result=result, ) pack_diffusion_output_shm(runner_out, d2h_stream=d2h) - # runner_out.result is replaced in-place by pack_diffusion_output_shm; - # assert with the original reference, not the replaced one. mock_pack_fields.assert_called_once_with( result, d2h_stream=d2h, + created=ANY, + pending_updates=ANY, ) diff --git a/tests/diffusion/test_media_output_contract.py b/tests/diffusion/test_media_output_contract.py new file mode 100644 index 00000000000..ae5c2fdbfc7 --- /dev/null +++ b/tests/diffusion/test_media_output_contract.py @@ -0,0 +1,546 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +import vllm_omni.diffusion.diffusion_engine as diffusion_engine_module +import vllm_omni.diffusion.ipc as diffusion_ipc +import vllm_omni.diffusion.postprocess.media as media_postprocess +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, VideoOutputTransportConfig +from vllm_omni.diffusion.diffusion_engine import DiffusionEngine +from vllm_omni.diffusion.ipc import pack_diffusion_output_shm, unpack_diffusion_output_shm +from vllm_omni.diffusion.media import ( + DiffusionMediaOutput, + FloatVideoConsumer, + VideoMediaOutput, + VideoTensorEncoding, + VideoTensorLayout, + VideoTensorSpec, + VideoTransportConstraints, + VideoValueRange, +) +from vllm_omni.diffusion.postprocess.device_reduction import prepare_diffusion_media_for_transport +from vllm_omni.diffusion.postprocess.media import finalize_diffusion_media +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _video( + tensor: torch.Tensor, + *, + encoding: VideoTensorEncoding, + value_range: VideoValueRange, + consumers: frozenset[FloatVideoConsumer] = frozenset(), +) -> VideoMediaOutput: + layout = VideoTensorLayout.BTHWC if encoding is VideoTensorEncoding.UINT8_FRAMES else VideoTensorLayout.BCTHW + return VideoMediaOutput( + tensor=tensor, + spec=VideoTensorSpec( + layout=layout, + encoding=encoding, + value_range=value_range, + ), + constraints=VideoTransportConstraints(pending_float_consumers=consumers), + ) + + +def _prepared_media(video: VideoMediaOutput) -> DiffusionMediaOutput: + return DiffusionMediaOutput(video=video, prepared_for_transport=True) + + +def _od_config(*, enabled: bool = False) -> OmniDiffusionConfig: + config = OmniDiffusionConfig( + model=None, + video_output_transport=VideoOutputTransportConfig(enable_device_postprocess=enabled), + ) + config.model_class_name = "WanPipeline" + return config + + +def test_finalizer_rejects_media_that_skipped_runner_preparation() -> None: + media = DiffusionMediaOutput( + video=_video( + torch.randn(1, 3, 2, 4, 5), + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ) + ) + + with pytest.raises(ValueError, match="before transport preparation"): + finalize_diffusion_media(media, sampling_params=OmniDiffusionSamplingParams(output_type="np")) + + +def test_uint8_media_finalization_preserves_frames() -> None: + frames = torch.randint(0, 256, (1, 4, 8, 8, 3), dtype=torch.uint8) + + output = finalize_diffusion_media( + _prepared_media( + _video( + frames, + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + ) + ), + sampling_params=OmniDiffusionSamplingParams(output_type="np"), + ) + + np.testing.assert_array_equal(output["payload"]["video"], frames.numpy()) + assert output["metadata"] == {} + + +def test_float_media_finalization_uses_the_declared_range() -> None: + video = torch.tensor([[[[[-1.0, 0.0, 1.0]]]]]).expand(1, 3, 1, 1, 3).contiguous() + + output = finalize_diffusion_media( + _prepared_media( + _video( + video, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ) + ), + sampling_params=OmniDiffusionSamplingParams(output_type="np"), + ) + + expected = (video.permute(0, 2, 3, 4, 1).numpy() * 0.5 + 0.5).clip(0.0, 1.0) + np.testing.assert_array_equal(output["payload"]["video"], expected) + + +def test_zero_to_one_float_media_is_not_denormalized() -> None: + video = torch.tensor([[[[[0.0, 0.5, 1.0]]]]]).expand(1, 3, 1, 1, 3).contiguous() + media = _prepared_media( + _video( + video, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.ZERO_TO_ONE, + ) + ) + + output = finalize_diffusion_media(media, sampling_params=OmniDiffusionSamplingParams(output_type="np")) + + expected = video.permute(0, 2, 3, 4, 1).numpy() + np.testing.assert_array_equal(output["payload"]["video"], expected) + + +def test_presentation_is_resolved_outside_the_tensor_descriptor() -> None: + video = torch.zeros(1, 3, 2, 4, 5) + media = _prepared_media( + _video( + video, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ) + ) + + output = finalize_diffusion_media(media, sampling_params=OmniDiffusionSamplingParams(output_type="pil")) + + assert "presentation" not in media.video.__dataclass_fields__ + assert len(output["payload"]["video"]) == 1 + assert len(output["payload"]["video"][0]) == 2 + + +def test_frame_interpolation_consumes_its_constraint(monkeypatch: pytest.MonkeyPatch) -> None: + video = torch.randn(1, 3, 2, 4, 5) + interpolated = torch.randn(1, 3, 3, 4, 5) + monkeypatch.setattr( + media_postprocess, + "interpolate_video_tensor", + lambda *args, **kwargs: (interpolated, 2.0), + ) + + output = finalize_diffusion_media( + _prepared_media( + _video( + video, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + consumers=frozenset({FloatVideoConsumer.FRAME_INTERPOLATION}), + ) + ), + sampling_params=OmniDiffusionSamplingParams(output_type="np", enable_frame_interpolation=True), + ) + + assert output["payload"]["video"].shape == (1, 3, 4, 5, 3) + assert output["metadata"] == {"video": {"video_fps_multiplier": 2.0}} + + +def test_uint8_media_rejects_a_pending_float_consumer() -> None: + media = _prepared_media( + _video( + torch.zeros(1, 2, 4, 5, 3, dtype=torch.uint8), + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + consumers=frozenset({FloatVideoConsumer.FRAME_INTERPOLATION}), + ) + ) + + with pytest.raises(ValueError, match="cannot have pending float consumers"): + media.validate() + + +def test_unconsumed_float_constraint_fails_before_public_formatting() -> None: + media = _prepared_media( + _video( + torch.randn(1, 3, 2, 4, 5), + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + consumers=frozenset({FloatVideoConsumer.VIDEO_GUARDRAILS}), + ) + ) + + with pytest.raises(ValueError, match="pending float consumers"): + finalize_diffusion_media(media, sampling_params=OmniDiffusionSamplingParams(output_type="np")) + + +def test_engine_routes_typed_media_exclusively(monkeypatch: pytest.MonkeyPatch) -> None: + media = _prepared_media( + _video( + torch.zeros(1, 1, 2, 2, 3, dtype=torch.uint8), + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + ) + ) + request = OmniDiffusionRequest( + prompt="test", + sampling_params=OmniDiffusionSamplingParams(output_type="np"), + request_id="request-0", + ) + engine = object.__new__(DiffusionEngine) + engine.od_config = _od_config() + engine.post_process_func = lambda _: pytest.fail("legacy postprocessor was called") + engine._post_process_accepts_sampling_params = False + + finalized = {"payload": {"video": np.zeros((1, 1, 2, 2, 3), dtype=np.uint8)}, "metadata": {}} + + def _format_outputs( + *, + request: object, + od_config: object, + diffusion_output: object, + output_data: object, + postprocess_output: object, + ) -> list[object]: + del request, od_config, diffusion_output, output_data + return [postprocess_output] + + monkeypatch.setattr(diffusion_engine_module, "finalize_diffusion_media", lambda *args, **kwargs: finalized) + monkeypatch.setattr(diffusion_engine_module, "normalize_diffusion_postprocess_output", lambda value: value) + monkeypatch.setattr(diffusion_engine_module, "format_diffusion_outputs", _format_outputs) + + result = engine.postprocess_output(request, DiffusionOutput(media=media)) + + assert result == [finalized] + + +def test_typed_media_crosses_runner_ipc_and_engine_without_legacy_postprocess( + monkeypatch: pytest.MonkeyPatch, +) -> None: + sampling_params = OmniDiffusionSamplingParams(output_type="np", num_outputs_per_prompt=1) + raw_media = DiffusionMediaOutput( + video=_video( + torch.linspace(-1.0, 1.0, 1 * 3 * 2 * 4 * 5).reshape(1, 3, 2, 4, 5), + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ) + ) + prepared = prepare_diffusion_media_for_transport( + raw_media, + od_config=_od_config(enabled=True), + sampling_params=sampling_params, + ) + diffusion_output = DiffusionOutput(media=prepared) + monkeypatch.setattr(diffusion_ipc, "_SHM_TENSOR_THRESHOLD", 1) + pack_diffusion_output_shm(diffusion_output) + unpack_diffusion_output_shm(diffusion_output) + + engine = object.__new__(DiffusionEngine) + engine.od_config = _od_config() + engine.post_process_func = lambda _: pytest.fail("legacy postprocessor was called") + engine._post_process_accepts_sampling_params = False + request = OmniDiffusionRequest(prompt="test", sampling_params=sampling_params, request_id="request-0") + + outputs = engine.postprocess_output(request, diffusion_output) + + assert len(outputs) == 1 + assert outputs[0].images[0].dtype == np.uint8 + assert outputs[0].images[0].shape == (1, 2, 4, 5, 3) + + +def test_diffusion_output_rejects_typed_and_legacy_output_together() -> None: + media = _prepared_media( + _video( + torch.zeros(1, 1, 2, 2, 3, dtype=torch.uint8), + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + ) + ) + with pytest.raises(ValueError, match="both media and legacy output"): + DiffusionOutput(output=torch.zeros(1), media=media) + + +def test_diffusion_output_rejects_postprocessor_on_typed_media() -> None: + media = _prepared_media( + _video( + torch.zeros(1, 1, 2, 2, 3, dtype=torch.uint8), + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + ) + ) + + with pytest.raises(ValueError, match="model-specific post_process_func"): + DiffusionOutput(media=media, post_process_func=lambda value: value) + + +def test_diffusion_output_preserves_the_legacy_positional_constructor() -> None: + # ``media`` was originally declared right after ``output``, which broke + # ``DiffusionOutput(output, trajectory_timesteps)`` by binding the second + # argument to media. It must stay last so out-of-tree positional callers keep + # the pre-existing field order. + output_tensor = torch.randn(1) + timesteps = torch.tensor([1, 2, 3]) + + output = DiffusionOutput(output_tensor, timesteps) + + assert output.output is output_tensor + assert output.trajectory_timesteps is timesteps + assert output.media is None + + +def test_frame_interpolation_uses_the_declared_range_negative_one_to_one( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # RIFE infers [0,1] vs [-1,1] from the tensor min/max, so an all-nonnegative + # [-1,1] clip is misread as [0,1] and interpolated in the wrong space. The + # interpolator must receive the unit-range form per the declared spec. + video = torch.full((1, 3, 2, 3, 3), 0.25) + media = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=video, + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ), + constraints=VideoTransportConstraints( + pending_float_consumers=frozenset({FloatVideoConsumer.FRAME_INTERPOLATION}) + ), + ), + prepared_for_transport=True, + ) + + seen: list[torch.Tensor] = [] + + def _square_interpolator(tensor: torch.Tensor, **_kwargs: object) -> tuple[torch.Tensor, float]: + seen.append(tensor) + return tensor * tensor, 2.0 + + monkeypatch.setattr(media_postprocess, "interpolate_video_tensor", _square_interpolator) + + output = finalize_diffusion_media( + media, + sampling_params=OmniDiffusionSamplingParams(output_type="np", enable_frame_interpolation=True), + ) + + # The interpolator is fed the unit-range form 0.25→0.625, not raw 0.25. + assert len(seen) == 1 + assert seen[0].shape == video.shape + assert float(seen[0][0, 0, 0, 0, 0]) == pytest.approx(0.625, abs=1e-6) + # Squaring gives 0.390625; restored to [-1,1] (-0.21875) by the fix, then the + # existing denormalize maps it back to [0,1] as 0.390625 in the public frames. + frames = output["payload"]["video"] + assert np.allclose(np.asarray(frames), 0.390625, atol=1e-6) + + +def test_frame_interpolation_uses_the_declared_range_zero_to_one( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # With a declared [0,1] range, the all-0.25 clip is already in unit space, so + # squaring gives 0.390625 and there is no restore step. + video = torch.full((1, 3, 2, 3, 3), 0.25) + media = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=video, + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.ZERO_TO_ONE, + ), + constraints=VideoTransportConstraints( + pending_float_consumers=frozenset({FloatVideoConsumer.FRAME_INTERPOLATION}) + ), + ), + prepared_for_transport=True, + ) + + def _square_interpolator(tensor: torch.Tensor, **_kwargs: object) -> tuple[torch.Tensor, float]: + return tensor * tensor, 2.0 + + monkeypatch.setattr(media_postprocess, "interpolate_video_tensor", _square_interpolator) + + output = finalize_diffusion_media( + media, + sampling_params=OmniDiffusionSamplingParams(output_type="np", enable_frame_interpolation=True), + ) + + # 0.25 is already [0,1]; squaring gives 0.0625, and no denormalize is applied. + frames = output["payload"]["video"] + assert np.allclose(np.asarray(frames), 0.0625, atol=1e-6) + + +def test_frame_interpolation_clamps_zero_to_one_overshoot( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # A declared [0,1] clip whose sample is almost entirely inside [0,1] — but + # with one bf16 value overshooting 1 — is still in unit space, so it must be + # clamped to [0,1] before RIFE and interpolated as [0,1]. Without the clamp + # the interpolator classifies it as [-1,1] (amax > 1) and returns the wrong + # space, which is then written back with no *2-1 restore. + video = torch.full((1, 3, 2, 3, 3), 0.25, dtype=torch.bfloat16) + video[0, 0, 0, 0, 0] = 1.25 # single overshoot > 1 -> amax > 1 + media = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=video, + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.ZERO_TO_ONE, + ), + constraints=VideoTransportConstraints( + pending_float_consumers=frozenset({FloatVideoConsumer.FRAME_INTERPOLATION}) + ), + ), + prepared_for_transport=True, + ) + + seen: list[torch.Tensor] = [] + + def _square_interpolator(tensor: torch.Tensor, **_kwargs: object) -> tuple[torch.Tensor, float]: + seen.append(tensor) + return tensor.float().square(), 2.0 + + monkeypatch.setattr(media_postprocess, "interpolate_video_tensor", _square_interpolator) + + output = finalize_diffusion_media( + media, + sampling_params=OmniDiffusionSamplingParams(output_type="np", enable_frame_interpolation=True), + ) + + # The interpolator is fed the clamped bf16 unit-range tensor (max now 1.0), + # not the raw 1.25 overshoot that would be misread as [-1,1]. + assert len(seen) == 1 + assert seen[0].dtype is torch.bfloat16 + assert float(seen[0].min()) == pytest.approx(0.25, abs=1e-6) + assert float(seen[0].max()) == pytest.approx(1.0, abs=1e-6) + # Squaring the clamped values gives 0.0625 at the 0.25 samples and 1.0 at the + # overshoot; the overall shape is unchanged (no *2-1 for ZERO_TO_ONE). + frames = output["payload"]["video"] + assert np.allclose(np.asarray(frames).reshape(-1)[1:], 0.0625, atol=1e-6) + assert np.isclose(float(np.asarray(frames).reshape(-1)[0]), 1.0, atol=1e-6) + + +def test_packing_failure_unlinks_created_segments_and_leaves_output_untouched( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Packing a late field can fail after an SHM segment is already created; + # the segment must be unlinked and ``output`` left unmutated so a partial + # object is never enqueued. + prepared = prepare_diffusion_media_for_transport( + DiffusionMediaOutput( + video=_video( + torch.linspace(-1.0, 1.0, 2 * 3 * 2 * 4 * 5).reshape(2, 3, 2, 4, 5), + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ) + ), + od_config=_od_config(enabled=False), + sampling_params=OmniDiffusionSamplingParams(output_type="np"), + ) + output = DiffusionOutput(media=prepared, trajectory_latents=torch.randn(30)) + original_media = prepared + original_latents = output.trajectory_latents + + real_array_to_shm = diffusion_ipc._array_to_shm + created_handles: list[dict[str, object]] = [] + + def _fail_on_second(array: np.ndarray) -> dict[str, object]: + if created_handles: + raise RuntimeError("injected packing failure") + handle = real_array_to_shm(array) + created_handles.append(handle) + return handle + + monkeypatch.setattr(diffusion_ipc, "_SHM_TENSOR_THRESHOLD", 1) + monkeypatch.setattr(diffusion_ipc, "_array_to_shm", _fail_on_second) + + with pytest.raises(RuntimeError, match="injected packing failure"): + pack_diffusion_output_shm(output) + + assert created_handles, "no SHM segment was created before the injected failure" + # Output was not mutated despite the media segment being created first. + assert output.media is original_media + assert isinstance(output.media, DiffusionMediaOutput) + assert output.trajectory_latents is original_latents + + from multiprocessing import shared_memory + + with pytest.raises(FileNotFoundError): + shared_memory.SharedMemory(name=str(created_handles[0]["name"])) + + +def test_second_batch_entry_failure_leaves_the_whole_batch_untouched( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from multiprocessing import shared_memory + + from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput + + first_media = _prepared_media( + _video( + torch.zeros(1, 3, 2, 4, 5), + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ) + ) + second_media = _prepared_media( + _video( + torch.ones(1, 3, 2, 4, 5), + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ) + ) + first = DiffusionOutput(media=first_media) + second = DiffusionOutput(media=second_media) + batch = BatchRunnerOutput( + runner_outputs=[ + RunnerOutput(request_id="first", result=first), + RunnerOutput(request_id="second", result=second), + ] + ) + + real_array_to_shm = diffusion_ipc._array_to_shm + created_handles: list[dict[str, object]] = [] + + def _fail_on_second(array: np.ndarray) -> dict[str, object]: + if created_handles: + raise RuntimeError("second batch entry failed") + handle = real_array_to_shm(array) + created_handles.append(handle) + return handle + + monkeypatch.setattr(diffusion_ipc, "_SHM_TENSOR_THRESHOLD", 1) + monkeypatch.setattr(diffusion_ipc, "_array_to_shm", _fail_on_second) + + with pytest.raises(RuntimeError, match="second batch entry failed"): + pack_diffusion_output_shm(batch) + + # No per-entry assignment is committed until every entry has packed. + assert first.media is first_media + assert second.media is second_media + with pytest.raises(FileNotFoundError): + shared_memory.SharedMemory(name=str(created_handles[0]["name"])) diff --git a/tests/diffusion/test_video_transport_defaults_are_inert.py b/tests/diffusion/test_video_transport_defaults_are_inert.py new file mode 100644 index 00000000000..fd7346e7606 --- /dev/null +++ b/tests/diffusion/test_video_transport_defaults_are_inert.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Pin the default behavior and the WAN reference integration boundary.""" + +from __future__ import annotations + +import inspect +from pathlib import Path + +import pytest + +from vllm_omni.diffusion.data import OmniDiffusionConfig, VideoOutputTransportConfig +from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def test_transport_defaults() -> None: + transport = VideoOutputTransportConfig() + + assert transport.enable_device_postprocess is False + + +def test_a_default_diffusion_config_carries_a_default_transport() -> None: + config = OmniDiffusionConfig(model=None) + + assert isinstance(config.video_output_transport, VideoOutputTransportConfig) + assert config.video_output_transport.enable_device_postprocess is False + + +def test_wan_uses_the_typed_media_contract() -> None: + import vllm_omni.diffusion.models as models_pkg + + source = (Path(inspect.getfile(models_pkg)).parent / "wan2_2/pipeline_wan2_2.py").read_text() + assert "DiffusionMediaOutput(" in source + assert "reduce_video_to_uint8_frames(" not in source + # The migrated built-in WAN pipeline emits typed media, so it never invokes + # the legacy postprocess hook. The mappings are intentionally retained so + # legacy/out-of-tree WAN replacements (registered with + # post_process_func_name=None) keep the built-in hook instead of exposing raw + # BCTHW tensors. + assert _DIFFUSION_POST_PROCESS_FUNCS["WanPipeline"] == "get_wan22_post_process_func" + assert _DIFFUSION_POST_PROCESS_FUNCS["WanDMDPipeline"] == "get_wan22_post_process_func" + assert _DIFFUSION_POST_PROCESS_FUNCS["WanT2VDMD2Pipeline"] == "get_wan22_post_process_func" + + +def test_non_reference_models_do_not_use_the_runtime_reducer() -> None: + import vllm_omni.diffusion.models as models_pkg + + root = Path(inspect.getfile(models_pkg)).parent + paths = ( + "cosmos3/pipeline_cosmos3.py", + "hunyuan_video/pipeline_hunyuan_video_1_5.py", + "lingbot_video/pipeline_lingbot_video.py", + "ltx2/ltx2_runtime.py", + "minimax_h3/pipeline_minimax_h3.py", + ) + for relative in paths: + source = (root / relative).read_text() + assert "prepare_diffusion_media_for_transport(" not in source + assert "reduce_video_to_uint8_frames(" not in source diff --git a/tests/diffusion/test_wan_device_reduction.py b/tests/diffusion/test_wan_device_reduction.py new file mode 100644 index 00000000000..492aec7a571 --- /dev/null +++ b/tests/diffusion/test_wan_device_reduction.py @@ -0,0 +1,348 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""WAN reference coverage for the typed pre-D2H media contract.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from tests.helpers.mark import hardware_marks +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, VideoOutputTransportConfig +from vllm_omni.diffusion.ipc import pack_diffusion_output_shm, unpack_diffusion_output_shm +from vllm_omni.diffusion.media import ( + DiffusionMediaOutput, + VideoMediaOutput, + VideoTensorEncoding, + VideoTensorLayout, + VideoTensorSpec, + VideoValueRange, +) +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import get_wan22_post_process_func +from vllm_omni.diffusion.postprocess.device_reduction import prepare_diffusion_media_for_transport +from vllm_omni.diffusion.postprocess.media import finalize_diffusion_media +from vllm_omni.entrypoints.openai.video_api_utils import _coerce_video_to_uint8_frames +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.diffusion, + *hardware_marks(res={"cuda": "L4"}, num_cards=1), +] + +_GPU = torch.accelerator.is_available() if hasattr(torch, "accelerator") else torch.cuda.is_available() +requires_gpu = pytest.mark.skipif(not _GPU, reason="device reduction is a GPU path") + + +def _decoded_wan_video(tensor: torch.Tensor) -> DiffusionMediaOutput: + return DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=tensor, + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ), + ) + ) + + +def _config(*, enabled: bool) -> OmniDiffusionConfig: + return OmniDiffusionConfig( + model=None, + video_output_transport=VideoOutputTransportConfig(enable_device_postprocess=enabled), + ) + + +def _encoder_frames(video_payload: np.ndarray) -> list[np.ndarray]: + return [_coerce_video_to_uint8_frames(video_payload[i]) for i in range(video_payload.shape[0])] + + +@requires_gpu +def test_wan_media_reduction_matches_float32_reference_and_bounds_native_path() -> None: + torch.manual_seed(0) + vae_out = torch.rand(2, 3, 8, 64, 96, device="cuda:0", dtype=torch.bfloat16) * 2.4 - 1.2 + sampling = OmniDiffusionSamplingParams(output_type=None, enable_frame_interpolation=False) + + legacy_postprocess = get_wan22_post_process_func(_config(enabled=False)) + widened_out = legacy_postprocess(vae_out.float(), output_type="np", sampling_params=sampling) + widened_frames = _encoder_frames(widened_out["payload"]["video"]) + native_out = legacy_postprocess(vae_out, output_type="np", sampling_params=sampling) + native_frames = _encoder_frames(native_out["payload"]["video"]) + + prepared = prepare_diffusion_media_for_transport( + _decoded_wan_video(vae_out), + od_config=_config(enabled=True), + sampling_params=sampling, + ) + reduced_out = finalize_diffusion_media(prepared, sampling_params=sampling) + assert prepared.video.spec.encoding is VideoTensorEncoding.UINT8_FRAMES + assert reduced_out["payload"]["video"].dtype == np.uint8 + reduced_frames = _encoder_frames(reduced_out["payload"]["video"]) + + assert len(widened_frames) == len(reduced_frames) == 2 + for expected, native, produced in zip(widened_frames, native_frames, reduced_frames, strict=True): + assert produced.shape == expected.shape == (8, 64, 96, 3) + np.testing.assert_array_equal(produced, expected) + deviation = np.abs(produced.astype(np.int16) - native.astype(np.int16)) + assert deviation.max() <= 1 + + +@requires_gpu +def test_wan_uint8_media_finalization_does_not_rescale() -> None: + frames = torch.randint(0, 256, (1, 4, 16, 16, 3), dtype=torch.uint8, device="cuda:0") + media = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=frames, + spec=VideoTensorSpec( + layout=VideoTensorLayout.BTHWC, + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + ), + ), + prepared_for_transport=True, + ) + + out = finalize_diffusion_media(media, sampling_params=OmniDiffusionSamplingParams(output_type="np")) + + np.testing.assert_array_equal(out["payload"]["video"], frames.cpu().numpy()) + + +@requires_gpu +def test_diffusion_output_to_cpu_moves_wan_media_off_device() -> None: + prepared = prepare_diffusion_media_for_transport( + _decoded_wan_video(torch.randn(1, 3, 2, 4, 5, device="cuda:0")), + od_config=_config(enabled=False), + sampling_params=OmniDiffusionSamplingParams(output_type="np"), + ) + output = DiffusionOutput(media=prepared, to_cpu=True) + + assert output.media is not None + assert output.media.video.tensor.device.type == "cpu" + + +@requires_gpu +def test_small_typed_media_is_on_cpu_before_ipc() -> None: + sampling = OmniDiffusionSamplingParams(output_type="np") + prepared = prepare_diffusion_media_for_transport( + _decoded_wan_video(torch.randn(1, 3, 2, 4, 5, device="cuda:0")), + od_config=_config(enabled=True), + sampling_params=sampling, + ) + output = DiffusionOutput(media=prepared) + + pack_diffusion_output_shm(output) + + assert isinstance(output.media, dict) + packed_tensor = output.media["video"]["tensor"] + assert isinstance(packed_tensor, torch.Tensor) + assert packed_tensor.device.type == "cpu" + unpack_diffusion_output_shm(output) + + +@requires_gpu +def test_wan_float_media_finalization_matches_the_legacy_path() -> None: + video = torch.rand(1, 3, 4, 16, 16, device="cuda:0", dtype=torch.bfloat16) * 2 - 1 + sampling = OmniDiffusionSamplingParams(output_type="np") + prepared = prepare_diffusion_media_for_transport( + _decoded_wan_video(video), + od_config=_config(enabled=False), + sampling_params=sampling, + ) + + produced = finalize_diffusion_media(prepared, sampling_params=sampling) + expected = get_wan22_post_process_func(_config(enabled=False))( + video, + output_type="np", + sampling_params=sampling, + ) + + assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT + np.testing.assert_array_equal(produced["payload"]["video"], expected["payload"]["video"]) + + +@requires_gpu +def test_oom_fallback_media_packs_without_device_widening(monkeypatch: pytest.MonkeyPatch) -> None: + # Regression: when the fp32 reduction OOMs, the fallback keeps the bf16 float + # payload. Async IPC must copy it to the host in bf16 and widen on the CPU; + # a second device-side fp32 conversion would OOM again and drop the payload. + from vllm_omni.diffusion.postprocess import device_reduction + + # Large enough (> 1 MB in bf16) that packing routes the tensor through the + # shared-memory D2H path rather than the small-tensor CPU copy. + video = torch.rand(1, 3, 16, 128, 128, device="cuda:0", dtype=torch.bfloat16) * 2 - 1 + + def _raise_oom(*_args: object, **_kwargs: object) -> torch.Tensor: + raise torch.OutOfMemoryError("forced fp32 reduction OOM") + + monkeypatch.setattr(device_reduction, "reduce_video_to_uint8_frames", _raise_oom) + + prepared = prepare_diffusion_media_for_transport( + _decoded_wan_video(video), + od_config=_config(enabled=True), + sampling_params=OmniDiffusionSamplingParams(output_type="np"), + ) + assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT + assert prepared.video.tensor.dtype == torch.bfloat16 + + output = DiffusionOutput(media=prepared) + + pinned_dtypes: list[torch.dtype] = [] + real_empty = torch.empty + + def _spy_empty(*args: object, **kwargs: object) -> torch.Tensor: + tensor = real_empty(*args, **kwargs) + if kwargs.get("pin_memory"): + pinned_dtypes.append(tensor.dtype) + return tensor + + monkeypatch.setattr(torch, "empty", _spy_empty) + + d2h_stream = torch.Stream(device=torch.device("cuda", 0)) + pack_diffusion_output_shm(output, d2h_stream=d2h_stream) + d2h_stream.synchronize() + + assert pinned_dtypes == [torch.bfloat16], ( + "async packing must D2H in the original dtype and widen on the CPU, " + "never widen bf16 to fp32 on the accelerator before the copy" + ) + + unpack_diffusion_output_shm(output) + assert isinstance(output.media, DiffusionMediaOutput) + assert output.media.video.tensor.dtype == torch.bfloat16 + torch.testing.assert_close(output.media.video.tensor, video.cpu()) + + +@requires_gpu +def test_oom_fallback_recovers_the_allocator(monkeypatch: pytest.MonkeyPatch) -> None: + # A real float32 OOM leaves the CUDA caching allocator in a failed state; + # the fallback must empty the cache and drain the device before marking the + # float representation transport-ready, so the worker D2H of the same tensor + # can actually run. Use a genuine allocation failure (per-process cap) rather + # than injecting the exception, which never poisons the allocator. + from vllm_omni.platforms import current_omni_platform + + video = torch.empty(1, 3, 16, 128, 128, device="cuda:0", dtype=torch.bfloat16) + video.uniform_(-1.0, 1.0) + expected = video.cpu() + d2h_stream = torch.Stream(device=video.device) + base_bytes = video.numel() * video.element_size() + _free, total = current_omni_platform.get_device_memory() + + # Cap the process so the fp32 conversion (2x base bytes) cannot fit, while + # retaining the original bf16 tensor. This forces a real allocator OOM at + # the .to(fp32), rather than injecting an exception from a monkeypatch. + torch.cuda.set_per_process_memory_fraction(min(0.99, (1.2 * base_bytes) / total)) + real_empty_cache = current_omni_platform.empty_cache + real_synchronize = current_omni_platform.synchronize + try: + empty_calls: list[int] = [] + sync_calls: list[bool] = [] + + def _empty_cache() -> None: + real_empty_cache() + empty_calls.append(torch.accelerator.current_device_index()) + + def _synchronize() -> None: + real_synchronize() + sync_calls.append(True) + + monkeypatch.setattr(current_omni_platform, "empty_cache", _empty_cache) + monkeypatch.setattr(current_omni_platform, "synchronize", _synchronize) + + prepared = prepare_diffusion_media_for_transport( + _decoded_wan_video(video), + od_config=_config(enabled=True), + sampling_params=OmniDiffusionSamplingParams(output_type="np"), + ) + + # The fallback ran past a genuine OOM and kept the bf16 float payload. + assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT + assert prepared.video.tensor.dtype == torch.bfloat16 + assert empty_calls, "fallback must empty the CUDA allocator cache after a real OOM" + assert sync_calls, "fallback must drain the device after a real OOM" + + # Exercise the next production step: async D2H/SHM packing of the same + # tensor must succeed after allocator recovery and preserve its values. + output = DiffusionOutput(media=prepared) + pack_diffusion_output_shm(output, d2h_stream=d2h_stream) + d2h_stream.synchronize() + unpack_diffusion_output_shm(output) + assert isinstance(output.media, DiffusionMediaOutput) + torch.testing.assert_close(output.media.video.tensor, expected) + finally: + torch.cuda.set_per_process_memory_fraction(1.0) + real_empty_cache() + real_synchronize() + + +@requires_gpu +def test_oom_fallback_releases_late_uint8_traceback_before_clone( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A final uint8 allocation OOM must release float intermediates first.""" + from vllm_omni.diffusion.postprocess import device_reduction + from vllm_omni.platforms import current_omni_platform + + # A split request is a non-compact view, so float fallback must allocate a + # request-owned clone after reduction fails. + backing = torch.empty(2, 3, 16, 128, 128, device="cuda:0", dtype=torch.bfloat16) + backing.uniform_(-1.0, 1.0) + video = backing[1:2].detach() + expected = video.cpu() + _, total = current_omni_platform.get_device_memory() + real_to = torch.Tensor.to + real_empty_cache = current_omni_platform.empty_cache + real_synchronize = current_omni_platform.synchronize + real_ensure_request_owned = device_reduction.ensure_request_owned_tensor + uint8_attempts: list[int] = [] + fallback_allocations: list[int] = [] + + def _record_fallback_allocation(tensor: torch.Tensor) -> torch.Tensor: + fallback_allocations.append(torch.accelerator.memory_allocated(tensor.device)) + return real_ensure_request_owned(tensor) + + def _fail_real_uint8_allocation( + tensor: torch.Tensor, + *args: object, + **kwargs: object, + ) -> torch.Tensor: + dtype = kwargs.get("dtype") + if dtype is None and args and isinstance(args[0], torch.dtype): + dtype = args[0] + if tensor.is_cuda and dtype is torch.uint8: + # Earlier float32 work runs without a cap. Immediately before the + # real final uint8 allocation, cap the allocator below that request + # but above the post-unwind request-owned bf16 clone requirement. + real_empty_cache() + allocated = torch.accelerator.memory_allocated(tensor.device) + uint8_bytes = tensor.numel() + limit_bytes = allocated + max(1, uint8_bytes // 2) + torch.cuda.set_per_process_memory_fraction(min(0.99, limit_bytes / total)) + uint8_attempts.append(allocated) + return real_to(tensor, *args, **kwargs) + + monkeypatch.setattr(device_reduction, "ensure_request_owned_tensor", _record_fallback_allocation) + monkeypatch.setattr(torch.Tensor, "to", _fail_real_uint8_allocation) + try: + prepared = prepare_diffusion_media_for_transport( + _decoded_wan_video(video), + od_config=_config(enabled=True), + sampling_params=OmniDiffusionSamplingParams(output_type="np"), + ) + + assert uint8_attempts, "the real final .to(uint8) allocation was not attempted" + assert fallback_allocations, "float fallback did not request compact storage" + assert fallback_allocations[0] < uint8_attempts[0], ( + "the failed reduction's float intermediates remained live during fallback" + ) + assert prepared.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT + assert prepared.video.tensor.dtype is torch.bfloat16 + assert prepared.video.tensor.storage_offset() == 0 + assert prepared.video.tensor.untyped_storage().data_ptr() != video.untyped_storage().data_ptr() + torch.testing.assert_close(prepared.video.tensor.cpu(), expected) + finally: + torch.cuda.set_per_process_memory_fraction(1.0) + real_empty_cache() + real_synchronize() diff --git a/tests/docs/test_generate_argparse.py b/tests/docs/test_generate_argparse.py new file mode 100644 index 00000000000..16b92365f9a --- /dev/null +++ b/tests/docs/test_generate_argparse.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import importlib.util +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +ROOT_DIR = Path(__file__).parents[2] +HOOK_PATH = ROOT_DIR / "docs/mkdocs/hooks/generate_argparse.py" +SPEC = importlib.util.spec_from_file_location("generate_argparse", HOOK_PATH) +assert SPEC and SPEC.loader +generate_argparse = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(generate_argparse) + + +def test_static_serve_parser_supports_video_output_transport() -> None: + parser = generate_argparse.create_parser_subparser_init(generate_argparse.OmniServeCommand) + + args = parser.parse_args(["--omni", "--video-output-transport", '{"enable_device_postprocess": true}']) + + assert args.video_output_transport == {"enable_device_postprocess": True} diff --git a/tests/e2e/offline_inference/test_device_postprocess_equivalence.py b/tests/e2e/offline_inference/test_device_postprocess_equivalence.py new file mode 100644 index 00000000000..a7fbcafd926 --- /dev/null +++ b/tests/e2e/offline_inference/test_device_postprocess_equivalence.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""End-to-end WAN check for the typed pre-D2H media contract. + +Runs the whole engine twice with the same seed, once per flag state, and compares +the frames a client receives. WAN denormalizes historical bfloat16 output in the +input dtype while device preparation computes in float32, so the paths may differ +by one 255th. + +Run with a WAN2.2 checkpoint:: + + VLLM_OMNI_DEVICE_POSTPROCESS_MODEL=Wan-AI/Wan2.2-TI2V-5B-Diffusers \ + pytest -s tests/e2e/offline_inference/test_device_postprocess_equivalence.py +""" + +from __future__ import annotations + +import os + +import numpy as np +import pytest + +from tests.helpers.mark import hardware_test +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +MODEL = os.environ.get("VLLM_OMNI_DEVICE_POSTPROCESS_MODEL") + +pytestmark = [ + pytest.mark.full_model, + pytest.mark.diffusion, + pytest.mark.gpu, + pytest.mark.skipif( + not MODEL, + reason="Set VLLM_OMNI_DEVICE_POSTPROCESS_MODEL to a WAN2.2 checkpoint to run this.", + ), +] + +_SEED = 1234 +_STEPS = int(os.environ.get("VLLM_OMNI_DEVICE_POSTPROCESS_STEPS", "2")) +_SIZE = int(os.environ.get("VLLM_OMNI_DEVICE_POSTPROCESS_SIZE", "256")) +_FRAMES = int(os.environ.get("VLLM_OMNI_DEVICE_POSTPROCESS_FRAMES", "17")) + + +def _generate(enable_device_postprocess: bool) -> np.ndarray: + from vllm_omni.entrypoints.omni import Omni + + engine = Omni( + model=MODEL, + num_gpus=1, + video_output_transport={"enable_device_postprocess": enable_device_postprocess}, + ) + try: + # output_type must travel with the request: the pipelines read it from the + # sampling params, and the reduction only covers the "np" path. + sampling_params = OmniDiffusionSamplingParams( + output_type="np", + seed=_SEED, + num_inference_steps=_STEPS, + height=_SIZE, + width=_SIZE, + num_frames=_FRAMES, + ) + outputs = engine.generate({"prompt": "a robot waving"}, sampling_params) + finally: + engine.close() + + video = outputs[0].images[0] + return video.numpy() if hasattr(video, "numpy") else np.asarray(video) + + +def _as_uint8_frames(video: np.ndarray) -> np.ndarray: + """Reproduce what the API server does before encoding. + + A float payload is scaled and rounded there, so that is the only fair place to + compare it against a payload the worker already reduced. + """ + if video.dtype == np.uint8: + return video + return np.rint(np.clip(video, 0.0, 1.0) * 255.0).astype(np.uint8) + + +# A dtype difference moves ~20% of pixels by one; dropping the round moves ~50%. +_MAX_DIFFERING_FRACTION = 0.35 + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_device_postprocess_produces_equivalent_frames() -> None: + float_path = _generate(False) + device_path = _generate(True) + + # Guard against the run silently taking the float path twice, which would make + # the comparison meaningless. + assert float_path.dtype != np.uint8, "the float path returned uint8; the gate may be stuck open" + assert device_path.dtype == np.uint8, "the WAN typed-media path did not prepare uint8 frames" + assert float_path.shape == device_path.shape + + expected = _as_uint8_frames(float_path) + deviation = np.abs(device_path.astype(np.int16) - expected.astype(np.int16)) + differing = float((deviation > 0).mean()) + print(f"maxdiff={deviation.max()} differing={differing:.4f} ({int((deviation > 0).sum())}/{deviation.size})") + + assert deviation.max() <= 1, "the reduction moved a pixel by more than one 255th" + assert differing < _MAX_DIFFERING_FRACTION, ( + f"{differing:.1%} of pixels differ, more than a denormalize-precision difference explains" + ) diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index b64758e5d75..0089330b291 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -69,6 +69,30 @@ def test_serve_parser_accepts_deploy_config() -> None: assert args.get_explicit_kwargs_dict()["deploy_config"] == "/tmp/deploy.yaml" +def test_serve_parser_accepts_video_output_transport() -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args( + ["serve", "fake-model", "--omni", "--video-output-transport", '{"enable_device_postprocess": true}'] + ) + + expected = {"enable_device_postprocess": True} + assert args.video_output_transport == expected + assert args.get_explicit_kwargs_dict()["video_output_transport"] == expected + + +@pytest.mark.parametrize("value", ["{not json", "[]"]) +def test_serve_parser_rejects_invalid_video_output_transport(value: str) -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + + with pytest.raises(SystemExit): + parser.parse_args(["serve", "fake-model", "--omni", "--video-output-transport", value]) + + def test_serve_parser_rejects_stage_configs_path() -> None: parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="subcommand") diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 81eb76f0b41..2b016be0d16 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -718,6 +718,7 @@ class keeps the diffusion-only knobs from ``OmniDiffusionConfig`` without cache_strategy: str = "none" cache_backend: str = "none" cache_config: Any = field(default_factory=dict) + video_output_transport: object = field(default_factory=dict) enable_cache_dit_summary: bool = False diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY diffusion_kv_max_rows_per_request: int | None = Field(default=None, ge=1, strict=True) @@ -812,6 +813,7 @@ def __post_init__(self) -> None: AttentionConfig, DiffusionCacheConfig, TransformerConfig, + VideoOutputTransportConfig, build_attention_config, parse_kv_cache_skip_selector, validate_dlo_host_registration_options, @@ -854,6 +856,13 @@ def __post_init__(self) -> None: elif not isinstance(self.cache_config, DiffusionCacheConfig): self.cache_config = DiffusionCacheConfig() + if self.video_output_transport is None: + self.video_output_transport = VideoOutputTransportConfig() + elif isinstance(self.video_output_transport, Mapping): + self.video_output_transport = VideoOutputTransportConfig(**dict(self.video_output_transport)) + elif not isinstance(self.video_output_transport, VideoOutputTransportConfig): + raise TypeError("video_output_transport must be a VideoOutputTransportConfig or mapping") + self._propagate_quantization_from_tf_config(self.tf_model_config) if self.quantization_config is not None: if isinstance(self.quantization_config, QuantizationConfig): diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 2ac1f99dd1c..f7d71489075 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -458,6 +458,7 @@ class StageDeployConfig: fa_deterministic: bool | None = None cache_backend: str | None = None cache_config: dict[str, Any] | None = None + video_output_transport: dict[str, Any] | None = None enable_cache_dit_summary: bool | None = None step_execution: bool | None = None vae_use_slicing: bool | None = None diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index c3052952076..d27cc36e191 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -29,6 +29,7 @@ parse_diffusion_kv_cache_mode, ) from vllm_omni.diffusion.lora.manager import LoRABackend +from vllm_omni.diffusion.media import DiffusionMediaOutput from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata from vllm_omni.diffusion.utils.network_utils import is_port_available from vllm_omni.errors import client_error_metadata @@ -758,6 +759,15 @@ def uses_diffusers_adapter(od_config: object) -> bool: ) +@dataclass +class VideoOutputTransportConfig: + enable_device_postprocess: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.enable_device_postprocess, bool): + raise TypeError("enable_device_postprocess must be a bool") + + @dataclass class OmniDiffusionConfig: # Model and path configuration (for convenience) @@ -793,6 +803,7 @@ class OmniDiffusionConfig: # Cache backend configuration (NEW) cache_backend: str = "none" # "tea_cache", "deep_cache", etc. cache_config: DiffusionCacheConfig | dict[str, Any] = field(default_factory=dict) + video_output_transport: VideoOutputTransportConfig = field(default_factory=VideoOutputTransportConfig) enable_cache_dit_summary: bool = False # Prompt-embedding cache. When enabled, ``DiffusionModelRunner`` wraps the @@ -1258,6 +1269,13 @@ def __post_init__(self): # If it's neither dict nor DiffusionCacheConfig, convert to empty config self.cache_config = DiffusionCacheConfig() + if self.video_output_transport is None: + self.video_output_transport = VideoOutputTransportConfig() + elif isinstance(self.video_output_transport, Mapping): + self.video_output_transport = VideoOutputTransportConfig(**dict(self.video_output_transport)) + elif not isinstance(self.video_output_transport, VideoOutputTransportConfig): + raise TypeError("video_output_transport must be a VideoOutputTransportConfig or mapping") + # Auto-detect quantization from TransformerConfig if not explicitly set. # This covers the case where tf_model_config is passed at construction # time. For late (post-construction) assignment, callers should use @@ -1618,8 +1636,8 @@ class DiffusionOutput: output: torch.Tensor | tuple[Any, ...] | dict[str, Any] | None = None - # Legacy compatibility fields. New pipeline-specific payloads should be - # carried by output["payload"] instead. + # Legacy compatibility fields. Decoded video uses ``media``; other new + # payloads should be carried by output["payload"] instead. trajectory_timesteps: torch.Tensor | dict[str, Any] | None = None trajectory_latents: torch.Tensor | dict[str, Any] | None = None trajectory_log_probs: torch.Tensor | dict[str, Any] | None = None @@ -1658,7 +1676,18 @@ class DiffusionOutput: # mode) and the receiving side must not initialise a stray CUDA context. to_cpu: bool = False + # Typed video-media contract. Declared last so the pre-existing positional + # constructor order (output, trajectory_timesteps, ...) that out-of-tree + # pipelines rely on is preserved. Mutually exclusive with ``output``. + media: DiffusionMediaOutput | None = None + def __post_init__(self) -> None: + if self.media is not None and not isinstance(self.media, DiffusionMediaOutput): + raise TypeError(f"media must be DiffusionMediaOutput, got {type(self.media).__name__}") + if self.media is not None and self.output is not None: + raise ValueError("DiffusionOutput cannot contain both media and legacy output") + if self.media is not None and self.post_process_func is not None: + raise ValueError("Typed diffusion media cannot carry a model-specific post_process_func") if not self.to_cpu: return @@ -1674,6 +1703,10 @@ def _maybe_to_cpu(value: Any) -> Any: return value self.output = _maybe_to_cpu(self.output) + if self.media is not None: + if not self.media.prepared_for_transport: + raise ValueError("Diffusion media must be prepared before to_cpu=True") + self.media = self.media.to_cpu() self.trajectory_timesteps = _maybe_to_cpu(self.trajectory_timesteps) self.trajectory_latents = _maybe_to_cpu(self.trajectory_latents) self.trajectory_log_probs = _maybe_to_cpu(self.trajectory_log_probs) diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index c36ac0459db..417090503f5 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -47,6 +47,7 @@ format_empty_diffusion_outputs, normalize_diffusion_postprocess_output, ) +from vllm_omni.diffusion.postprocess.media import finalize_diffusion_media from vllm_omni.diffusion.registry import ( DiffusionModelRegistry, get_diffusion_post_process_func, @@ -534,26 +535,33 @@ def postprocess_output( raise RuntimeError(output.error) logger.debug("Generation completed successfully.") - if output.output is None: - logger.warning("Output is None, returning empty OmniRequestOutput") - return format_empty_diffusion_outputs(request, finished=output.finished) - - # When CPU offload is enabled, move output to CPU before - # post-processing to avoid device OOM — model weights may still - # reside on the device and leave no headroom for intermediates. - output_data = output.output - if self.od_config.enable_cpu_offload: - output_data = _move_tensor_tree_to_cpu(output_data) - - if self.post_process_func is not None: - # Some video pipelines need request-level controls during - # postprocess (for example worker-side frame interpolation). - postprocess_kwargs: dict[str, object] = {} - if self._post_process_accepts_sampling_params: - postprocess_kwargs["sampling_params"] = request.sampling_params - outputs = self.post_process_func(output_data, **postprocess_kwargs) + if output.media is not None: + if output.output is not None: + raise ValueError("DiffusionOutput cannot contain both media and legacy output") + media = output.media.to_cpu() if self.od_config.enable_cpu_offload else output.media + output_data = media.video.tensor + outputs = finalize_diffusion_media(media, sampling_params=request.sampling_params) else: - outputs = output_data + if output.output is None: + logger.warning("Output is None, returning empty OmniRequestOutput") + return format_empty_diffusion_outputs(request, finished=output.finished) + + # When CPU offload is enabled, move output to CPU before + # post-processing to avoid device OOM — model weights may still + # reside on the device and leave no headroom for intermediates. + output_data = output.output + if self.od_config.enable_cpu_offload: + output_data = _move_tensor_tree_to_cpu(output_data) + + if self.post_process_func is not None: + # Some video pipelines need request-level controls during + # postprocess (for example worker-side frame interpolation). + postprocess_kwargs: dict[str, object] = {} + if self._post_process_accepts_sampling_params: + postprocess_kwargs["sampling_params"] = request.sampling_params + outputs = self.post_process_func(output_data, **postprocess_kwargs) + else: + outputs = output_data postprocess_output = normalize_diffusion_postprocess_output(outputs) diff --git a/vllm_omni/diffusion/ipc.py b/vllm_omni/diffusion/ipc.py index 6eb24c54033..f0e88a2c97e 100644 --- a/vllm_omni/diffusion/ipc.py +++ b/vllm_omni/diffusion/ipc.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 """IPC utilities for transferring large tensors via POSIX shared memory. @@ -11,15 +11,32 @@ from __future__ import annotations +from collections.abc import Iterator from typing import Any import numpy as np import torch from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.media import ( + DiffusionMediaOutput, + FloatVideoConsumer, + VideoColorModel, + VideoMediaOutput, + VideoTensorEncoding, + VideoTensorLayout, + VideoTensorSpec, + VideoTransportConstraints, + VideoValueRange, +) _SHM_TENSOR_THRESHOLD = 1_000_000 # 1 MB DIFFUSION_RPC_RESULT_ENVELOPE = "diffusion_rpc_result" +_DIFFUSION_MEDIA_WIRE_TYPE = "diffusion_media_v1" + +# Sentinel so compute-then-assign packing can tell "field not packed" apart from +# "field packed to None". +_UNSET = object() def _array_to_shm(array: np.ndarray) -> dict[str, Any]: @@ -61,9 +78,39 @@ def _array_from_shm(handle: dict[str, Any]) -> np.ndarray: return array +def _unlink_shm_handle(handle: dict[str, Any]) -> None: + """Best-effort removal of a shared-memory segment created during packing.""" + from multiprocessing import shared_memory + + name = handle.get("name") + if not name: + return + try: + shm = shared_memory.SharedMemory(name=name) + except FileNotFoundError: + return + try: + shm.close() + finally: + try: + shm.unlink() + except FileNotFoundError: + pass + + +def _unlink_shm_handles(handles: list[dict[str, Any]]) -> None: + for handle in handles: + try: + _unlink_shm_handle(handle) + except Exception: + # Cleanup is best-effort; never mask the original packing failure. + pass + + def _tensor_to_shm( tensor: torch.Tensor, d2h_stream: torch.Stream | None = None, + created: list[dict[str, Any]] | None = None, ) -> dict[str, Any]: """Copy a tensor into POSIX shared memory and return a metadata handle. @@ -82,14 +129,17 @@ def _tensor_to_shm( torch.accelerator.set_stream(d2h_stream) try: t = tensor.detach() - if original_dtype == torch.bfloat16: - t = t.to(torch.float32) cpu = torch.empty(t.shape, dtype=t.dtype, pin_memory=True) cpu.copy_(t, non_blocking=True) finally: torch.accelerator.set_stream(old_stream) d2h_stream.synchronize() tensor = cpu + # numpy has no bfloat16, so widen on the host after D2H. Never widen on + # the accelerator before the copy: a payload that already fell back after + # an fp32 OOM would OOM again on a second device-side fp32 allocation. + if original_dtype == torch.bfloat16: + tensor = tensor.to(torch.float32) else: tensor = tensor.detach().cpu().contiguous() if original_dtype == torch.bfloat16: @@ -101,6 +151,8 @@ def _tensor_to_shm( "torch_dtype": str(original_dtype), } ) + if created is not None: + created.append(handle) return handle @@ -120,6 +172,7 @@ def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor: def _pack_tensor_if_large( val: torch.Tensor, d2h_stream: torch.Stream | None = None, + created: list[dict[str, Any]] | None = None, ) -> torch.Tensor | dict: """Replace a tensor with an SHM handle if it exceeds the threshold. @@ -134,14 +187,16 @@ def _pack_tensor_if_large( except Exception: storage_bytes = view_bytes if max(view_bytes, storage_bytes) > _SHM_TENSOR_THRESHOLD: - return _tensor_to_shm(val, d2h_stream=d2h_stream) + return _tensor_to_shm(val, d2h_stream=d2h_stream, created=created) return val -def _ndarray_to_shm(array: np.ndarray) -> dict[str, Any]: +def _ndarray_to_shm(array: np.ndarray, created: list[dict[str, Any]] | None = None) -> dict[str, Any]: """Copy a contiguous NumPy array into POSIX shared memory.""" handle = _array_to_shm(array) handle["__ndarray_shm__"] = True + if created is not None: + created.append(handle) return handle @@ -153,6 +208,7 @@ def _ndarray_from_shm(handle: dict[str, Any]) -> np.ndarray: def _pack_value_if_large( val: object, d2h_stream: torch.Stream | None = None, + created: list[dict[str, Any]] | None = None, ) -> object: """Recursively replace large tensors with SHM handles. @@ -163,15 +219,19 @@ def _pack_value_if_large( the two in sync. """ if isinstance(val, torch.Tensor): - return _pack_tensor_if_large(val, d2h_stream=d2h_stream) + return _pack_tensor_if_large(val, d2h_stream=d2h_stream, created=created) if isinstance(val, np.ndarray): - return _ndarray_to_shm(val) if not val.dtype.hasobject and val.nbytes > _SHM_TENSOR_THRESHOLD else val + return ( + _ndarray_to_shm(val, created=created) + if not val.dtype.hasobject and val.nbytes > _SHM_TENSOR_THRESHOLD + else val + ) if isinstance(val, dict): - return {key: _pack_value_if_large(value, d2h_stream=d2h_stream) for key, value in val.items()} + return {key: _pack_value_if_large(value, d2h_stream=d2h_stream, created=created) for key, value in val.items()} if isinstance(val, list): - return [_pack_value_if_large(item, d2h_stream=d2h_stream) for item in val] + return [_pack_value_if_large(item, d2h_stream=d2h_stream, created=created) for item in val] if isinstance(val, tuple): - return tuple(_pack_value_if_large(item, d2h_stream=d2h_stream) for item in val) + return tuple(_pack_value_if_large(item, d2h_stream=d2h_stream, created=created) for item in val) return val @@ -190,18 +250,116 @@ def _unpack_if_shm_handle(val: object) -> object: return val +def _pack_diffusion_media( + media: DiffusionMediaOutput, + *, + d2h_stream: torch.Stream | None, + created: list[dict[str, Any]] | None = None, +) -> dict[str, Any]: + media.validate() + if not media.prepared_for_transport: + raise ValueError("Diffusion media must be prepared before IPC packing") + video = media.video + packed_tensor = _pack_tensor_if_large(video.tensor, d2h_stream=d2h_stream, created=created) + if isinstance(packed_tensor, torch.Tensor) and packed_tensor.device.type != "cpu": + packed_tensor = packed_tensor.detach().cpu().contiguous() + return { + "__type__": _DIFFUSION_MEDIA_WIRE_TYPE, + "schema_version": video.schema_version, + "prepared_for_transport": True, + "video": { + "tensor": packed_tensor, + "layout": video.spec.layout.value, + "encoding": video.spec.encoding.value, + "value_range": video.spec.value_range.value, + "color_model": video.spec.color_model.value, + "pending_float_consumers": sorted(consumer.value for consumer in video.constraints.pending_float_consumers), + }, + } + + +def _unpack_diffusion_media(packed: dict[str, Any]) -> DiffusionMediaOutput: + if packed.get("__type__") != _DIFFUSION_MEDIA_WIRE_TYPE: + raise ValueError("Invalid diffusion media wire payload") + video_payload = packed.get("video") + if not isinstance(video_payload, dict): + raise ValueError("Diffusion media wire payload is missing video metadata") + tensor = _unpack_if_shm_handle(video_payload.get("tensor")) + if not isinstance(tensor, torch.Tensor): + raise ValueError("Diffusion media wire payload contains an invalid video tensor") + schema_version = packed.get("schema_version") + if schema_version != 1: + raise ValueError(f"Unsupported video media schema version: {schema_version}") + if packed.get("prepared_for_transport") is not True: + raise ValueError("Diffusion media wire payload was not prepared for transport") + media = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=tensor, + spec=VideoTensorSpec( + layout=VideoTensorLayout(video_payload["layout"]), + encoding=VideoTensorEncoding(video_payload["encoding"]), + value_range=VideoValueRange(video_payload["value_range"]), + color_model=VideoColorModel(video_payload["color_model"]), + ), + constraints=VideoTransportConstraints( + pending_float_consumers=frozenset( + FloatVideoConsumer(value) for value in video_payload.get("pending_float_consumers", []) + ) + ), + schema_version=1, + ), + prepared_for_transport=True, + ) + media.validate() + return media + + def _pack_diffusion_fields( output: DiffusionOutput, d2h_stream: torch.Stream | None = None, + created: list[dict[str, Any]] | None = None, + pending_updates: list[tuple[DiffusionOutput, str, object]] | None = None, ) -> DiffusionOutput: + if output.media is not None and output.output is not None: + raise ValueError("DiffusionOutput cannot contain both media and legacy output") + if output.media is not None and not isinstance(output.media, DiffusionMediaOutput): + raise TypeError(f"DiffusionOutput.media must be DiffusionMediaOutput, got {type(output.media).__name__}") + + # Pack into locals first so a mid-payload failure leaves ``output`` unmutated + # (never enqueue a half-packed object); the caller unlinks any SHM segments + # accumulated in ``created`` on failure. + packed_output = _UNSET if output.output is not None: - output.output = _pack_value_if_large(output.output, d2h_stream=d2h_stream) - if output.trajectory_latents is not None and isinstance(output.trajectory_latents, torch.Tensor): - output.trajectory_latents = _pack_tensor_if_large(output.trajectory_latents, d2h_stream=d2h_stream) - if output.trajectory_timesteps is not None and isinstance(output.trajectory_timesteps, torch.Tensor): - output.trajectory_timesteps = _pack_tensor_if_large(output.trajectory_timesteps, d2h_stream=d2h_stream) - if output.trajectory_log_probs is not None and isinstance(output.trajectory_log_probs, torch.Tensor): - output.trajectory_log_probs = _pack_tensor_if_large(output.trajectory_log_probs, d2h_stream=d2h_stream) + packed_output = _pack_value_if_large(output.output, d2h_stream=d2h_stream, created=created) + packed_media = _UNSET + if output.media is not None: + packed_media = _pack_diffusion_media(output.media, d2h_stream=d2h_stream, created=created) + packed_latents = _UNSET + if isinstance(output.trajectory_latents, torch.Tensor): + packed_latents = _pack_tensor_if_large(output.trajectory_latents, d2h_stream=d2h_stream, created=created) + packed_timesteps = _UNSET + if isinstance(output.trajectory_timesteps, torch.Tensor): + packed_timesteps = _pack_tensor_if_large(output.trajectory_timesteps, d2h_stream=d2h_stream, created=created) + packed_log_probs = _UNSET + if isinstance(output.trajectory_log_probs, torch.Tensor): + packed_log_probs = _pack_tensor_if_large(output.trajectory_log_probs, d2h_stream=d2h_stream, created=created) + + updates = [ + (output, field_name, value) + for field_name, value in ( + ("output", packed_output), + ("media", packed_media), + ("trajectory_latents", packed_latents), + ("trajectory_timesteps", packed_timesteps), + ("trajectory_log_probs", packed_log_probs), + ) + if value is not _UNSET + ] + if pending_updates is None: + for target, field_name, value in updates: + setattr(target, field_name, value) + else: + pending_updates.extend(updates) return output @@ -223,35 +381,84 @@ def pack_diffusion_output_shm( If *d2h_stream* is provided, D2H copies use that stream (non-blocking on the default stream). The caller must synchronize *d2h_stream* afterward. + + Packing is failure-atomic: every SHM segment created during the call is + tracked and unlinked if any part of the payload fails to pack, and field + mutations are committed only after the complete payload succeeds. """ + created: list[dict[str, Any]] = [] + pending_updates: list[tuple[DiffusionOutput, str, object]] = [] + try: + packed = _pack_output_shm( + output, + d2h_stream=d2h_stream, + created=created, + pending_updates=pending_updates, + ) + except BaseException: + _unlink_shm_handles(created) + raise + + # Commit mutations only after every field and batch entry packed. Until + # this point, a failure leaves the complete caller-owned payload untouched. + for target, field_name, value in pending_updates: + setattr(target, field_name, value) + return packed + + +def _iter_diffusion_outputs(output: object) -> Iterator[DiffusionOutput]: + """Yield every packable ``DiffusionOutput`` leaf in a worker payload.""" if isinstance(output, DiffusionOutput): - return _pack_diffusion_fields(output, d2h_stream=d2h_stream) + yield output + return # DP multi-concurrency: {"dp_rank": int, "output": DiffusionOutput} if isinstance(output, dict) and "dp_rank" in output and "output" in output: inner = output["output"] if isinstance(inner, DiffusionOutput): - output["output"] = _pack_diffusion_fields(inner, d2h_stream=d2h_stream) - return output + yield inner + return if _is_rpc_result_envelope(output): - result = output.get("result") - output["result"] = pack_diffusion_output_shm(result, d2h_stream=d2h_stream) - return output + yield from _iter_diffusion_outputs(output.get("result")) + return result = getattr(output, "result", None) if isinstance(result, DiffusionOutput): - output.result = _pack_diffusion_fields(result, d2h_stream=d2h_stream) + yield result runner_outputs = getattr(output, "runner_outputs", None) if isinstance(runner_outputs, list): for runner_output in runner_outputs: - pack_diffusion_output_shm(runner_output, d2h_stream=d2h_stream) + yield from _iter_diffusion_outputs(runner_output) + + +def payload_carries_typed_media(output: object) -> bool: + """True if *output* carries any typed ``DiffusionOutput.media`` payload.""" + return any(item.media is not None for item in _iter_diffusion_outputs(output)) + + +def _pack_output_shm( + output: object, + *, + d2h_stream: torch.Stream | None, + created: list[dict[str, Any]], + pending_updates: list[tuple[DiffusionOutput, str, object]], +) -> object: + for item in _iter_diffusion_outputs(output): + _pack_diffusion_fields( + item, + d2h_stream=d2h_stream, + created=created, + pending_updates=pending_updates, + ) return output def _unpack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput: output.output = _unpack_if_shm_handle(output.output) + if isinstance(output.media, dict): + output.media = _unpack_diffusion_media(output.media) output.trajectory_latents = _unpack_if_shm_handle(output.trajectory_latents) output.trajectory_timesteps = _unpack_if_shm_handle(output.trajectory_timesteps) output.trajectory_log_probs = _unpack_if_shm_handle(output.trajectory_log_probs) diff --git a/vllm_omni/diffusion/media.py b/vllm_omni/diffusion/media.py new file mode 100644 index 00000000000..34f4a122455 --- /dev/null +++ b/vllm_omni/diffusion/media.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from dataclasses import dataclass, field, replace +from enum import Enum +from typing import Literal + +import torch + + +class VideoTensorLayout(str, Enum): + BCTHW = "bcthw" + BTHWC = "bthwc" + + +class VideoTensorEncoding(str, Enum): + NORMALIZED_FLOAT = "normalized_float" + UINT8_FRAMES = "uint8_frames" + + +class VideoValueRange(str, Enum): + NEGATIVE_ONE_TO_ONE = "negative_one_to_one" + ZERO_TO_ONE = "zero_to_one" + ZERO_TO_255 = "zero_to_255" + + +class VideoColorModel(str, Enum): + RGB = "rgb" + + +class FloatVideoConsumer(str, Enum): + FRAME_INTERPOLATION = "frame_interpolation" + VIDEO_GUARDRAILS = "video_guardrails" + + +@dataclass(frozen=True, slots=True) +class VideoTensorSpec: + layout: VideoTensorLayout + encoding: VideoTensorEncoding + value_range: VideoValueRange + color_model: VideoColorModel = VideoColorModel.RGB + + +@dataclass(frozen=True, slots=True) +class VideoTransportConstraints: + pending_float_consumers: frozenset[FloatVideoConsumer] = frozenset() + + +@dataclass(frozen=True, slots=True) +class VideoMediaOutput: + tensor: torch.Tensor + spec: VideoTensorSpec + constraints: VideoTransportConstraints = field(default_factory=VideoTransportConstraints) + schema_version: Literal[1] = 1 + + def validate(self) -> None: + if self.schema_version != 1: + raise ValueError(f"Unsupported video media schema version: {self.schema_version}") + if not isinstance(self.spec, VideoTensorSpec): + raise TypeError(f"Video media spec must be a VideoTensorSpec, got {type(self.spec).__name__}") + if not isinstance(self.constraints, VideoTransportConstraints): + raise TypeError( + f"Video media constraints must be VideoTransportConstraints, got {type(self.constraints).__name__}" + ) + if not isinstance(self.constraints.pending_float_consumers, frozenset): + raise TypeError("pending_float_consumers must be a frozenset") + if not isinstance(self.tensor, torch.Tensor): + raise TypeError(f"Video media tensor must be a torch.Tensor, got {type(self.tensor).__name__}") + if self.tensor.ndim != 5: + raise ValueError(f"Video media tensor must be rank 5, got shape {tuple(self.tensor.shape)}") + if any(size <= 0 for size in self.tensor.shape): + raise ValueError(f"Video media dimensions must be positive, got shape {tuple(self.tensor.shape)}") + if self.spec.color_model is not VideoColorModel.RGB: + raise ValueError(f"Unsupported video color model: {self.spec.color_model}") + if any(not isinstance(consumer, FloatVideoConsumer) for consumer in self.constraints.pending_float_consumers): + raise ValueError("pending_float_consumers must contain only FloatVideoConsumer values") + + if self.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT: + self._validate_normalized_float() + return + if self.spec.encoding is VideoTensorEncoding.UINT8_FRAMES: + self._validate_uint8_frames() + return + raise ValueError(f"Unsupported video tensor encoding: {self.spec.encoding}") + + def _validate_normalized_float(self) -> None: + if self.spec.layout is not VideoTensorLayout.BCTHW: + raise ValueError("NORMALIZED_FLOAT video must use BCTHW layout") + if self.tensor.shape[1] != 3: + raise ValueError(f"RGB BCTHW video must have 3 channels, got {self.tensor.shape[1]}") + if self.tensor.dtype not in {torch.float16, torch.bfloat16, torch.float32}: + raise ValueError(f"NORMALIZED_FLOAT video must use fp16, bf16, or fp32, got {self.tensor.dtype}") + if self.spec.value_range not in { + VideoValueRange.NEGATIVE_ONE_TO_ONE, + VideoValueRange.ZERO_TO_ONE, + }: + raise ValueError(f"NORMALIZED_FLOAT video has incompatible value range: {self.spec.value_range}") + + def _validate_uint8_frames(self) -> None: + if self.spec.layout is not VideoTensorLayout.BTHWC: + raise ValueError("UINT8_FRAMES video must use BTHWC layout") + if self.tensor.shape[-1] != 3: + raise ValueError(f"RGB BTHWC video must have 3 channels, got {self.tensor.shape[-1]}") + if self.tensor.dtype is not torch.uint8: + raise ValueError(f"UINT8_FRAMES video must use torch.uint8, got {self.tensor.dtype}") + if self.spec.value_range is not VideoValueRange.ZERO_TO_255: + raise ValueError(f"UINT8_FRAMES video has incompatible value range: {self.spec.value_range}") + if self.constraints.pending_float_consumers: + raise ValueError("UINT8_FRAMES video cannot have pending float consumers") + if not self.tensor.is_contiguous(): + raise ValueError("UINT8_FRAMES video tensor must be contiguous") + + def with_tensor(self, tensor: torch.Tensor) -> VideoMediaOutput: + return replace(self, tensor=tensor) + + def to_cpu(self) -> VideoMediaOutput: + return self.with_tensor(self.tensor.detach().cpu()) + + +@dataclass(frozen=True, slots=True) +class DiffusionMediaOutput: + video: VideoMediaOutput + # Only the model runner may set this before worker output packing. + prepared_for_transport: bool = False + + def validate(self) -> None: + if not isinstance(self.video, VideoMediaOutput): + raise TypeError(f"Diffusion media video must be a VideoMediaOutput, got {type(self.video).__name__}") + if not isinstance(self.prepared_for_transport, bool): + raise TypeError("prepared_for_transport must be a bool") + self.video.validate() + if not self.prepared_for_transport: + return + if not self.video.tensor.is_contiguous(): + raise ValueError("Prepared video media tensor must be contiguous") + if not _owns_compact_storage(self.video.tensor): + raise ValueError("Prepared video media tensor must own request-local storage") + + def with_video(self, video: VideoMediaOutput) -> DiffusionMediaOutput: + return replace(self, video=video) + + def to_cpu(self) -> DiffusionMediaOutput: + self.validate() + moved = self.with_video(self.video.to_cpu()) + moved.validate() + return moved + + +def _owns_compact_storage(tensor: torch.Tensor) -> bool: + """True if *tensor* is contiguous and backed by exactly its own elements. + + ``tensor._base is None`` is not sufficient: ``batch[i : i + 1].detach()`` is + contiguous with ``_base=None`` yet keeps a nonzero storage offset and the + whole batch storage, so serializing it would leak every other request's + pixels. Require a zero offset and a storage sized to the logical elements. + """ + if not tensor.is_contiguous(): + return False + if tensor.storage_offset() != 0: + return False + try: + storage_bytes = tensor.untyped_storage().nbytes() + except Exception: + return False + return storage_bytes == tensor.numel() * tensor.element_size() + + +def ensure_request_owned_tensor(tensor: torch.Tensor) -> torch.Tensor: + if _owns_compact_storage(tensor): + return tensor + # A contiguous view may still retain another request's full batch storage. + return tensor.clone(memory_format=torch.contiguous_format) + + +def slice_diffusion_media_output( + media: DiffusionMediaOutput, + start: int, + stop: int, +) -> DiffusionMediaOutput: + media.validate() + source = media.video.tensor + tensor = source if start == 0 and stop == source.shape[0] else source[start:stop] + sliced = replace(media, video=media.video.with_tensor(tensor), prepared_for_transport=False) + sliced.validate() + return sliced diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index c9ef325c4ee..a285e0009d3 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -26,6 +26,14 @@ from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.forward_context import DenoiseProgressMixin from vllm_omni.diffusion.lora.loader import WanLoraLoaderMixin +from vllm_omni.diffusion.media import ( + DiffusionMediaOutput, + VideoMediaOutput, + VideoTensorEncoding, + VideoTensorLayout, + VideoTensorSpec, + VideoValueRange, +) from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin @@ -907,6 +915,7 @@ def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: if DEBUG_PERF: _t_decode_start = time.perf_counter() + media = None if output_type == "latent": output = latents else: @@ -920,7 +929,27 @@ def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: latents.device, latents.dtype ) latents = latents / latents_std + latents_mean - output = self.vae.decode(latents, return_dict=False)[0] + decoded = self.vae.decode(latents, return_dict=False)[0] + # Distributed VAE decode uses broadcast_result=False, so only the + # output-owning rank receives the full [B, C, T, H, W] video; other + # ranks get an empty placeholder. Emit typed media only from the + # owning rank and keep the placeholder on the legacy output field, so + # the media batch-dimension check in split_diffusion_output_by_request + # does not trip on every non-owner rank. + if decoded.dim() == 5: + output = None + media = DiffusionMediaOutput( + video=VideoMediaOutput( + tensor=decoded, + spec=VideoTensorSpec( + layout=VideoTensorLayout.BCTHW, + encoding=VideoTensorEncoding.NORMALIZED_FLOAT, + value_range=VideoValueRange.NEGATIVE_ONE_TO_ONE, + ), + ) + ) + else: + output = decoded if DEBUG_PERF: current_omni_platform.synchronize() @@ -947,6 +976,7 @@ def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: return split_diffusion_output_by_request( DiffusionOutput( output=output, + media=media, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ), req, diff --git a/vllm_omni/diffusion/postprocess/device_reduction.py b/vllm_omni/diffusion/postprocess/device_reduction.py new file mode 100644 index 00000000000..ec4ef895f39 --- /dev/null +++ b/vllm_omni/diffusion/postprocess/device_reduction.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Device-side reduction of decoded video tensors to uint8 frames.""" + +from __future__ import annotations + +from dataclasses import replace + +import torch +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.media import ( + DiffusionMediaOutput, + FloatVideoConsumer, + VideoMediaOutput, + VideoTensorEncoding, + VideoTensorLayout, + VideoTensorSpec, + VideoTransportConstraints, + VideoValueRange, + ensure_request_owned_tensor, +) +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.platforms import current_omni_platform + +# VaeImageProcessor.denormalize: (x * 0.5 + 0.5).clamp(0, 1) +_DENORM_SCALE = 0.5 +_DENORM_SHIFT = 0.5 + +logger = init_logger(__name__) + + +def _request_float_consumers( + sampling_params: OmniDiffusionSamplingParams | None, +) -> frozenset[FloatVideoConsumer]: + if sampling_params is not None and sampling_params.enable_frame_interpolation: + return frozenset({FloatVideoConsumer.FRAME_INTERPOLATION}) + return frozenset() + + +def _request_supports_uint8_frames(sampling_params: OmniDiffusionSamplingParams | None) -> bool: + return sampling_params is None or (sampling_params.output_type or "np") == "np" + + +def _prepare_float_media_for_transport( + media: DiffusionMediaOutput, + video: VideoMediaOutput, +) -> DiffusionMediaOutput: + prepared = replace( + media, + video=video.with_tensor(ensure_request_owned_tensor(video.tensor)), + prepared_for_transport=True, + ) + prepared.validate() + return prepared + + +def prepare_diffusion_media_for_transport( + media: DiffusionMediaOutput, + *, + od_config: OmniDiffusionConfig, + sampling_params: OmniDiffusionSamplingParams | None = None, +) -> DiffusionMediaOutput: + """Validate and prepare one request-local media output before D2H.""" + media.validate() + if media.prepared_for_transport: + # Only the model runner prepares media, exactly once, on media the + # pipeline emits unprepared. Accepting pre-prepared media here would let + # a pipeline forge runner-owned state and skip request policy (e.g. a + # forged uint8 payload silently bypasses frame interpolation). + raise ValueError( + "Media reached transport preparation already prepared; pipelines must " + "emit unprepared media so request policy is applied by the runner" + ) + if media.video.spec.encoding is not VideoTensorEncoding.NORMALIZED_FLOAT: + raise ValueError("Pipelines must emit unprepared media as NORMALIZED_FLOAT video") + + video = media.video + constraints = VideoTransportConstraints( + pending_float_consumers=(video.constraints.pending_float_consumers | _request_float_consumers(sampling_params)) + ) + constrained_video = replace(video, constraints=constraints) + + enabled = od_config.video_output_transport.enable_device_postprocess + output_supported = _request_supports_uint8_frames(sampling_params) + if not enabled or constraints.pending_float_consumers or not output_supported: + if not enabled: + reason = "disabled" + elif constraints.pending_float_consumers: + reason = ",".join(sorted(consumer.value for consumer in constraints.pending_float_consumers)) + else: + reason = "unsupported_presentation" + logger.debug("Device video preparation kept normalized float: reason=%s", reason) + return _prepare_float_media_for_transport(media, constrained_video) + + try: + frames = reduce_video_to_uint8_frames( + constrained_video.tensor, + do_denormalize=constrained_video.spec.value_range is VideoValueRange.NEGATIVE_ONE_TO_ONE, + ) + except torch.OutOfMemoryError: + logger.warning("Device video preparation ran out of memory; using normalized float transport") + else: + prepared_video = replace( + constrained_video, + tensor=frames, + spec=VideoTensorSpec( + layout=VideoTensorLayout.BTHWC, + encoding=VideoTensorEncoding.UINT8_FRAMES, + value_range=VideoValueRange.ZERO_TO_255, + ), + constraints=VideoTransportConstraints(), + ) + prepared = replace(media, video=prepared_video, prepared_for_transport=True) + prepared.validate() + logger.debug( + "Device video preparation converted normalized float to uint8: shape=%s", + tuple(frames.shape), + ) + return prepared + + # Leave the exception handler before allocator cleanup and fallback + # construction. A late conversion OOM otherwise keeps the helper traceback + # (and its live float32 intermediates) reachable while the fallback may need + # to clone a non-compact request view. + current_omni_platform.empty_cache() + current_omni_platform.synchronize() + return _prepare_float_media_for_transport(media, constrained_video) + + +def reduce_video_to_uint8_frames(video: torch.Tensor, *, do_denormalize: bool = True) -> torch.Tensor: + """Reduce a decoded ``[B, C, F, H, W]`` video to uint8 ``[B, F, H, W, C]`` frames. + + Runs denormalize/clamp/permute/round on the input's device so the following + D2H copy carries uint8 instead of float. The result matches + ``VideoProcessor.postprocess_video(output_type="np")`` then the ``*255`` + rounding done in the API server. Pass ``do_denormalize=False`` for VAEs that + already emit ``[0, 1]``. + """ + if video.dim() != 5: + raise ValueError(f"expected a [B, C, F, H, W] video tensor, got shape {tuple(video.shape)}") + + # Match the numpy path, which promotes to float before scaling. + frames = video.to(torch.float32) + if do_denormalize: + frames = frames.mul(_DENORM_SCALE).add(_DENORM_SHIFT).clamp_(0.0, 1.0) + else: + frames = frames.clamp(0.0, 1.0) + frames = frames.permute(0, 2, 3, 4, 1) + frames = frames.mul_(255.0).round_().clamp_(0.0, 255.0).to(torch.uint8) + return frames.contiguous() diff --git a/vllm_omni/diffusion/postprocess/media.py b/vllm_omni/diffusion/postprocess/media.py new file mode 100644 index 00000000000..fd256443478 --- /dev/null +++ b/vllm_omni/diffusion/postprocess/media.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from dataclasses import replace + +from vllm_omni.diffusion.media import ( + DiffusionMediaOutput, + FloatVideoConsumer, + VideoTensorEncoding, + VideoTensorLayout, + VideoTransportConstraints, + VideoValueRange, +) +from vllm_omni.diffusion.postprocess.rife_interpolator import interpolate_video_tensor +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + +def finalize_diffusion_media( + media: DiffusionMediaOutput, + *, + sampling_params: OmniDiffusionSamplingParams, +) -> dict[str, object]: + """Finalize typed video media without invoking a model postprocessor.""" + media.validate() + if not media.prepared_for_transport: + raise ValueError("Diffusion media reached the engine before transport preparation") + + video = media.video + metadata: dict[str, object] = {} + consumers = video.constraints.pending_float_consumers + if FloatVideoConsumer.FRAME_INTERPOLATION in consumers: + if video.spec.encoding is not VideoTensorEncoding.NORMALIZED_FLOAT: + raise ValueError("Frame interpolation requires NORMALIZED_FLOAT video") + # RIFE guesses [0,1] vs [-1,1] from the tensor's min/max, so relying on + # the sample content misclassifies both an all-nonnegative [-1,1] clip + # (read as [0,1]) and a [0,1] clip whose single bf16 value overshoots 1 + # (read as [-1,1]). Always map to unit range per the declared spec and + # clamp to [0,1] before RIFE, then restore the declared range so + # downstream denormalization is correct regardless of the sample. + to_negative_one = video.spec.value_range is VideoValueRange.NEGATIVE_ONE_TO_ONE + if to_negative_one: + interp_input = video.tensor.mul(0.5).add(0.5).clamp_(0.0, 1.0) + else: + interp_input = video.tensor.clamp(0.0, 1.0) + interpolated, multiplier = interpolate_video_tensor( + interp_input, + exp=sampling_params.frame_interpolation_exp, + scale=sampling_params.frame_interpolation_scale, + model_path=sampling_params.frame_interpolation_model_path, + ) + if to_negative_one: + interpolated = interpolated.mul(2.0).sub(1.0) + consumers = consumers - {FloatVideoConsumer.FRAME_INTERPOLATION} + video = replace( + video, + tensor=interpolated, + constraints=VideoTransportConstraints(pending_float_consumers=frozenset(consumers)), + ) + metadata["video"] = {"video_fps_multiplier": multiplier} + + if consumers: + names = sorted(consumer.value for consumer in consumers) + raise ValueError(f"Video has pending float consumers at finalization: {names}") + + output_type = getattr(sampling_params, "output_type", None) or "np" + if output_type == "latent": + raise ValueError("Latent output cannot be finalized from VideoMediaOutput") + + if video.spec.encoding is VideoTensorEncoding.UINT8_FRAMES: + if video.spec.layout is not VideoTensorLayout.BTHWC: + raise ValueError("UINT8_FRAMES video must use BTHWC layout") + if output_type != "np": + raise ValueError(f"UINT8_FRAMES finalization does not support output_type={output_type!r}") + frames = video.tensor.detach().cpu().numpy() + else: + if video.spec.layout is not VideoTensorLayout.BCTHW: + raise ValueError("NORMALIZED_FLOAT video must use BCTHW layout") + from diffusers.video_processor import VideoProcessor + + do_denormalize = video.spec.value_range is VideoValueRange.NEGATIVE_ONE_TO_ONE + frames = VideoProcessor(vae_scale_factor=8).postprocess_video( + video.tensor, + output_type=output_type, + do_denormalize=[do_denormalize] * video.tensor.shape[2], + ) + + return { + "payload": {"video": frames}, + "metadata": metadata, + } diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 8804785fadc..9dfcd6f3fb2 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -53,6 +53,7 @@ ) from vllm_omni.diffusion.offloader import enable_offload_backend from vllm_omni.diffusion.offloader.config import TEXT_ENCODER_COMPONENT, resolve_offload +from vllm_omni.diffusion.postprocess.device_reduction import prepare_diffusion_media_for_transport from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.interface import ( @@ -680,6 +681,21 @@ def _refresh_cache_for_requests( od_config.cache_backend, ) + def _prepare_output_for_transport( + self, + output: DiffusionOutput, + sampling_params: OmniDiffusionSamplingParams, + ) -> DiffusionOutput: + if output.media is not None: + if output.output is not None: + raise ValueError("DiffusionOutput cannot contain both media and legacy output") + output.media = prepare_diffusion_media_for_transport( + output.media, + od_config=self.od_config, + sampling_params=sampling_params, + ) + return output + def _runner_output_from_outputs( self, reqs: list[OmniDiffusionRequest], @@ -774,6 +790,11 @@ def _execute_request_list( allow_single_output=allow_single_output, pipeline_name=type(self.pipeline).__name__, ) + with record_function("prepare_output_for_transport"): + outputs = [ + self._prepare_output_for_transport(output, req.sampling_params) + for req, output in zip(reqs, outputs, strict=True) + ] if is_primary and outputs and record_output_peak_memory: batch_peak_memory_mb = self._sample_peak_memory_mb() @@ -1289,6 +1310,7 @@ def _execute_stepwise_core( clear_pipeline_stage_durations(self.pipeline) result = self.pipeline.post_decode(req) if result is not None: + result = self._prepare_output_for_transport(result, req.sampling) self._attach_stepwise_metadata( req, result, diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 4443353a008..b9506f102d9 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -64,7 +64,11 @@ model_parallel_is_initialized, ) from vllm_omni.diffusion.forward_context import set_forward_context -from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, pack_diffusion_output_shm +from vllm_omni.diffusion.ipc import ( + DIFFUSION_RPC_RESULT_ENVELOPE, + pack_diffusion_output_shm, + payload_carries_typed_media, +) from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager, LoRABackend from vllm_omni.diffusion.registry import get_diffusion_ir_op_priority_func from vllm_omni.diffusion.request import OmniDiffusionRequest @@ -1131,7 +1135,15 @@ def _return_result(self, output: Any, rpc_id: str | None = None) -> None: # Sync path (original, or async fallback). try: pack_diffusion_output_shm(output) + except (TypeError, ValueError): + raise except Exception as e: + # Typed media that fails to pack is left unprepared/on device (the + # pack is failure-atomic and does not mutate it), so enqueueing it + # would ship a broken payload. Re-raise memory/packing failures here + # instead of swallowing them for typed media. + if payload_carries_typed_media(output): + raise if hasattr(output, "output"): logger.warning("SHM pack failed for model output: %s", e) self._enqueue_result(output) diff --git a/vllm_omni/diffusion/worker/request_batch.py b/vllm_omni/diffusion/worker/request_batch.py index 079a213dbf2..d5f1d5585b2 100644 --- a/vllm_omni/diffusion/worker/request_batch.py +++ b/vllm_omni/diffusion/worker/request_batch.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 """Request-level batch abstraction for diffusion runner.""" from __future__ import annotations @@ -12,6 +12,7 @@ import torch from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.media import DiffusionMediaOutput, slice_diffusion_media_output from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType @@ -35,6 +36,18 @@ def split_diffusion_output_by_request( """Split a batched DiffusionOutput into one output per request.""" if num_outputs_per_prompt <= 0: raise ValueError(f"num_outputs_per_prompt must be positive, got {num_outputs_per_prompt}.") + if result.media is not None and result.output is not None: + raise ValueError("DiffusionOutput cannot contain both media and legacy output") + if result.media is not None and not isinstance(result.media, DiffusionMediaOutput): + raise TypeError(f"DiffusionOutput.media must be DiffusionMediaOutput, got {type(result.media).__name__}") + if result.media is not None: + expected_batch = req.num_reqs * num_outputs_per_prompt + actual_batch = result.media.video.tensor.shape[0] + if actual_batch != expected_batch: + raise ValueError( + f"Video media batch dimension must equal request_count * num_outputs_per_prompt " + f"({expected_batch}), got {actual_batch}" + ) return [ DiffusionOutput( @@ -43,6 +56,15 @@ def split_diffusion_output_by_request( idx * num_outputs_per_prompt, (idx + 1) * num_outputs_per_prompt, ), + media=( + slice_diffusion_media_output( + result.media, + idx * num_outputs_per_prompt, + (idx + 1) * num_outputs_per_prompt, + ) + if result.media is not None + else None + ), error=result.error, finished=result.finished, stage_durations=result.stage_durations, diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 44bd9058c7d..ee555275189 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -566,6 +566,7 @@ class OrchestratorArgs: diffusion_compile_dynamic: bool | None = None cache_backend: str = "none" cache_config: str | None = None + video_output_transport: dict[str, object] | None = None enable_cache_dit_summary: bool = False step_execution: bool = False vae_use_slicing: bool = False diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 5c0204a3345..348e4131643 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -71,6 +71,16 @@ def _nonneg_finite_float(value: str) -> float: return parsed +def _json_object(value: str) -> dict[str, object]: + try: + parsed = json.loads(value) + except json.JSONDecodeError as exc: + raise argparse.ArgumentTypeError(f"must be valid JSON: {exc}") from exc + if not isinstance(parsed, dict): + raise argparse.ArgumentTypeError("must be a JSON object") + return parsed + + def _ensure_vllm_platform(): """Ensure vLLM's current_platform is valid before arg parsing. @@ -669,6 +679,14 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu 'MagCache: \'{"mag_threshold": 0.24, "mag_max_skip_steps": 5, "mag_retention_ratio": 0.1}\'. ' "Calibration mode: add '\"mag_calibrate\": true'", ) + omni_config_group.add_argument( + "--video-output-transport", + type=_json_object, + default=None, + help=( + "JSON object configuring video output preparation, for example '{\"enable_device_postprocess\": true}'." + ), + ) omni_config_group.add_argument( "--enable-cache-dit-summary", action="store_true", From 7576c4ceb98c2219a3f54b295d27be8b5fed3e31 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Thu, 10 Sep 2026 17:56:27 +0800 Subject: [PATCH 53/93] [Bugfix] Bound HWR domain initialization lock waits (#7128) Signed-off-by: Hongsheng Liu --- docs/design/feature/host_weight_runtime.md | 11 +- docs/design/module/host_weight_runtime.md | 12 +- .../test_runtime_resolution.py | 108 ++++++++++++++++++ .../host_weight_runtime/filesystem/store.py | 10 +- vllm_omni/host_weight_runtime/runtime.py | 2 +- 5 files changed, 134 insertions(+), 9 deletions(-) diff --git a/docs/design/feature/host_weight_runtime.md b/docs/design/feature/host_weight_runtime.md index a9bb1b0b838..ac7bc7b3f9d 100644 --- a/docs/design/feature/host_weight_runtime.md +++ b/docs/design/feature/host_weight_runtime.md @@ -313,8 +313,15 @@ One process per exact identity owns a build; other workers wait and then acquire leases for the published artifact. Publication is invisible until all payloads and metadata are validated, hashed, fsynced, and atomically renamed. -`coordination_timeout_seconds` bounds filesystem lock acquisition. It does not -cancel synchronous validation, a producer that has already started, or atomic +`coordination_timeout_seconds` bounds domain-initialization and lookup/build +lock acquisition. Store construction and each later resolution or publication +operation have separate budgets from the same wait policy, rather than one +end-to-end startup deadline. A domain-init timeout follows the retryable domain +failure policy below. After contention ends, a fresh runtime construction can +retry initialization; the timed-out runtime retains its original failure. + +The coordination budget does not cancel filesystem I/O, synchronous validation, +a producer that has already started, or atomic publication. A hung in-process producer therefore blocks its owning process and must be handled by external process supervision. Enforceable producer cancellation requires a future process-isolated producer contract. diff --git a/docs/design/module/host_weight_runtime.md b/docs/design/module/host_weight_runtime.md index ac92711fc4d..c975616671b 100644 --- a/docs/design/module/host_weight_runtime.md +++ b/docs/design/module/host_weight_runtime.md @@ -136,8 +136,16 @@ so only when it is marked retryable; unsupported capabilities, identity collisions, producer failures, and publication failures remain visible as startup failures. -`coordination_timeout_seconds` bounds acquisition of lookup and build locks; it -is not a hard wall-clock deadline for in-process validation, production, or +`coordination_timeout_seconds` bounds acquisition of the domain-initialization +lock and lookup/build locks. Store construction and each later resolution or +publication operation receive separate coordination budgets from the same +`WaitPolicy`; they do not share an end-to-end startup deadline. A domain-init +lock timeout is a retryable domain failure: preferred mode may fall back and +required mode fails. Retrying initialization requires constructing a fresh +runtime after contention ends; a waiter never cancels or unlocks the owner. + +The coordination budget is not a hard wall-clock deadline for filesystem I/O, +in-process validation, production, or atomic publication. Once this V1 implementation becomes the producer, the synchronous producer and publication run to completion. A hung producer blocks its owning process and requires external process supervision; enforceable diff --git a/tests/host_weight_runtime/test_runtime_resolution.py b/tests/host_weight_runtime/test_runtime_resolution.py index 21bab37dee4..cd2de28165d 100644 --- a/tests/host_weight_runtime/test_runtime_resolution.py +++ b/tests/host_weight_runtime/test_runtime_resolution.py @@ -4,6 +4,10 @@ from __future__ import annotations +import multiprocessing as mp +import threading +import time +from multiprocessing.connection import Connection from pathlib import Path from typing import cast @@ -15,6 +19,7 @@ CanonicalJson, ComponentIdentity, FailureCode, + HostWeightError, HostWeightFailure, HostWeightLease, HostWeightLeaseCarrier, @@ -40,12 +45,14 @@ StoreStatus, TensorWriteSpec, ValidationLevel, + WaitPolicy, WeightArtifactIdentity, WeightProductionSpec, WeightRepresentation, WeightSourceIdentity, ) from vllm_omni.host_weight_runtime.filesystem import FilesystemHostWeightStore, detect_storage_class +from vllm_omni.host_weight_runtime.filesystem.locks import FileLock, lock_is_active pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -518,6 +525,107 @@ def test_retryable_store_initialization_failure_obeys_runtime_mode( assert resolution.report.attempts[0].failure.retryable +def _initialize_contended_runtime(config: HostWeightRuntimeConfig, output: Connection) -> None: + try: + output.send("starting") + started = time.monotonic() + runtime = HostWeightRuntime.from_config(config) + resolution = runtime.resolve(_identity()) + try: + output.send((time.monotonic() - started, resolution.report)) + finally: + if resolution.lease is not None: + resolution.lease.close() + finally: + output.close() + + +@pytest.mark.parametrize( + ("mode", "expected"), + [ + (RuntimeMode.PREFERRED, ResolutionOutcome.CANONICAL_FALLBACK), + (RuntimeMode.REQUIRED, ResolutionOutcome.FAILED), + ], +) +def test_domain_lock_timeout_and_fresh_runtime_recovery( + tmp_path: Path, mode: RuntimeMode, expected: ResolutionOutcome +) -> None: + config = HostWeightRuntimeConfig(mode=mode, domain=_domain(tmp_path), wait=WaitPolicy(0.1)) + initial = HostWeightRuntime.from_config(config) + built = initial.resolve(_identity(), producer=CountingProducer(_identity())) + assert built.lease is not None + built.lease.close() + metadata = {name: (tmp_path / name).read_bytes() for name in ("domain.json", "domain-policy.json")} + lock_path = tmp_path / "locks" / "domain-init.lock" + ctx = mp.get_context("spawn") + reader, writer = ctx.Pipe(duplex=False) + process = ctx.Process(target=_initialize_contended_runtime, args=(config, writer)) + + with FileLock(lock_path, exclusive=True, deadline=None): + process.start() + writer.close() + try: + # Import/bootstrap is outside the short coordination budget. + assert reader.poll(60), "waiter did not start" + assert reader.recv() == "starting" + assert reader.poll(5), "domain initialization did not respect the coordination timeout" + elapsed, report = reader.recv() + assert 0.1 <= elapsed < 5 + assert report.outcome is expected + failure = report.attempts[0].failure + assert failure is not None + assert failure.stage is ResolutionStage.DOMAIN + assert failure.code is FailureCode.DOMAIN_UNAVAILABLE + assert failure.retryable + assert "timed out waiting for" in failure.message + assert "domain-init.lock" in failure.message + assert lock_is_active(lock_path), "waiter must not unlock the owner" + process.join(5) + assert process.exitcode == 0 + finally: + if process.is_alive(): + process.kill() + process.join(5) + process.close() + reader.close() + + assert {name: (tmp_path / name).read_bytes() for name in metadata} == metadata + recovered = HostWeightRuntime.from_config(config).resolve(_identity()) + assert recovered.report.outcome is ResolutionOutcome.LOCAL_HIT + assert recovered.lease is not None + recovered.lease.close() + + +def test_direct_store_waits_for_domain_lock_released_before_deadline(tmp_path: Path) -> None: + config = HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path)) + assert config.domain is not None + initial = HostWeightRuntime.from_config(config) + assert initial.store is not None + lock_path = tmp_path / "locks" / "domain-init.lock" + with FileLock(lock_path, exclusive=True, deadline=None) as owner: + release = threading.Timer(0.1, owner.close) + release.start() + try: + store = FilesystemHostWeightStore(config.domain, config.capacity, config.integrity, wait=WaitPolicy(5)) + finally: + release.join(5) + assert store.root == tmp_path + + +def test_direct_store_domain_timeout_is_typed(tmp_path: Path) -> None: + config = HostWeightRuntimeConfig(mode=RuntimeMode.PREFERRED, domain=_domain(tmp_path)) + assert config.domain is not None + HostWeightRuntime.from_config(config) + lock_path = tmp_path / "locks" / "domain-init.lock" + with FileLock(lock_path, exclusive=True, deadline=None): + with pytest.raises(HostWeightError) as caught: + FilesystemHostWeightStore(config.domain, config.capacity, config.integrity, wait=WaitPolicy(0.01)) + assert caught.value.failure.code is FailureCode.DOMAIN_UNAVAILABLE + assert caught.value.failure.retryable + assert lock_is_active(lock_path) + assert not lock_is_active(lock_path) + + @pytest.mark.parametrize( ("stage", "code"), [ diff --git a/vllm_omni/host_weight_runtime/filesystem/store.py b/vllm_omni/host_weight_runtime/filesystem/store.py index 2b58b77e874..ec6b49f3c79 100644 --- a/vllm_omni/host_weight_runtime/filesystem/store.py +++ b/vllm_omni/host_weight_runtime/filesystem/store.py @@ -23,7 +23,7 @@ import torch from safetensors import SafetensorError, safe_open -from ..config import CapacityPolicy, IntegrityPolicy, StorageClass, StorageDomainPolicy, ValidationLevel +from ..config import CapacityPolicy, IntegrityPolicy, StorageClass, StorageDomainPolicy, ValidationLevel, WaitPolicy from ..errors import FailureCode, HostWeightError, HostWeightFailure, ResolutionStage from ..identity import CanonicalJson, WeightArtifactIdentity, canonical_json from ..inspection import ( @@ -429,6 +429,8 @@ def __init__( domain: StorageDomainPolicy, capacity: CapacityPolicy, integrity: IntegrityPolicy, + *, + wait: WaitPolicy = WaitPolicy(), ) -> None: self.domain_policy = domain self.capacity_policy = capacity @@ -464,7 +466,7 @@ def __init__( ) ) try: - self._initialize_domain() + self._initialize_domain(deadline=time.monotonic() + wait.coordination_timeout_seconds) except HostWeightError: raise except OSError as exc: @@ -485,7 +487,7 @@ def __init__( ) ) from exc - def _initialize_domain(self) -> None: + def _initialize_domain(self, *, deadline: float) -> None: try: self.filesystem_type = detect_filesystem_type(self.root) detected_storage_class = _storage_class_for_filesystem_type(self.filesystem_type) @@ -545,7 +547,7 @@ def _initialize_domain(self) -> None: ) ) - with FileLock(self.locks_dir / "domain-init.lock", exclusive=True, deadline=None): + with FileLock(self.locks_dir / "domain-init.lock", exclusive=True, deadline=deadline): domain_path = self.root / _DOMAIN_FILE if domain_path.exists(): domain_metadata = _read_json_file(domain_path) diff --git a/vllm_omni/host_weight_runtime/runtime.py b/vllm_omni/host_weight_runtime/runtime.py index fbc726669cf..17ababce9eb 100644 --- a/vllm_omni/host_weight_runtime/runtime.py +++ b/vllm_omni/host_weight_runtime/runtime.py @@ -70,7 +70,7 @@ def from_config( from .filesystem import FilesystemHostWeightStore # noqa: PLC0415 try: - store = FilesystemHostWeightStore(config.domain, config.capacity, config.integrity) + store = FilesystemHostWeightStore(config.domain, config.capacity, config.integrity, wait=config.wait) except HostWeightError as exc: return cls( config, From ce27ee6038bb188ed874f024ebae431d809d8b12 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Thu, 10 Sep 2026 17:56:44 +0800 Subject: [PATCH 54/93] [Bugfix] Escalate diffusion worker shutdown and retain survivors (#7126) Signed-off-by: Hongsheng Liu --- .../test_multiproc_engine_concurrency.py | 133 +++++++++++++++++- .../diffusion/executor/multiproc_executor.py | 71 +++++++--- 2 files changed, 185 insertions(+), 19 deletions(-) diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index efe6bf9a320..ade8d8fa725 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -4,8 +4,10 @@ import asyncio import multiprocessing as mp import queue +import signal import threading import time +import weakref from types import SimpleNamespace from unittest.mock import MagicMock, Mock @@ -1330,6 +1332,7 @@ def test_worker_joins_share_one_global_deadline(self, monkeypatch): class FakeProcess: def __init__(self, name): self.name = name + self.pid = 123 self.alive = True self.terminated = False self.join_timeouts = [] @@ -1358,6 +1361,132 @@ def terminate(self): assert first.terminated and second.terminated assert not first.is_alive() and not second.is_alive() + @pytest.mark.parametrize("cooperative", [False, True]) + def test_real_workers_exit_and_are_reaped(self, monkeypatch, cooperative): + from vllm_omni.diffusion.executor import multiproc_executor as executor_module + + ctx = mp.get_context("fork") + ready, stop = ctx.Event(), ctx.Event() + + def run_worker(): + signal.signal(signal.SIGTERM, signal.SIG_IGN) + ready.set() + stop.wait() + + dead = ctx.Process(target=lambda: None) + worker = ctx.Process(target=run_worker) + monkeypatch.setattr(executor_module, "_WORKER_SHUTDOWN_GRACE_S", 0.2) + monkeypatch.setattr(executor_module, "_WORKER_TERMINATE_GRACE_S", 0.1) + monkeypatch.setattr(executor_module, "_WORKER_KILL_GRACE_S", 5.0, raising=False) + dead.start() + worker.start() + try: + dead.join(5) + assert not dead.is_alive() + assert ready.wait(5), "worker did not install its signal handler" + mq = Mock() + if cooperative: + mq.enqueue.side_effect = lambda *args, **kwargs: stop.set() + else: + mq.enqueue.side_effect = OSError("shutdown queue unavailable") + cleaner = executor_module._ExecutorShutdownCleaner(mq, 2, [dead, worker]) + + cleaner() + + assert not worker.is_alive() + assert worker.exitcode == (0 if cooperative else -signal.SIGKILL) + assert cleaner.processes == [] + assert cleaner.broadcast_mq is None + cleaner() + finally: + for proc in (dead, worker): + if proc.is_alive(): + proc.kill() + proc.join(5) + proc.close() + + @pytest.mark.parametrize("failed_action", ["terminate", "kill"]) + def test_kill_phase_shares_deadline_and_continues_after_os_errors(self, monkeypatch, failed_action): + from vllm_omni.diffusion.executor import multiproc_executor as executor_module + + first, second = Mock(pid=101), Mock(pid=102) + first.name, second.name = "first", "second" + for proc in (first, second): + proc.is_alive.side_effect = [True, True, True, False] + getattr(first, failed_action).side_effect = OSError("signal failed") + first.join.side_effect = [OSError("join failed"), None, None] + monotonic = Mock(side_effect=[100, 100, 110, 120, 120, 124, 130, 130, 134]) + monkeypatch.setattr(executor_module, "time", SimpleNamespace(monotonic=monotonic)) + cleaner = executor_module._ExecutorShutdownCleaner(processes=[first, second]) + + cleaner() + + first.kill.assert_called_once() + second.kill.assert_called_once() + assert [call.args[0] for call in first.join.call_args_list] == [15, 5, 5] + assert [call.args[0] for call in second.join.call_args_list] == [5, 1, 1] + assert cleaner.processes == [] + + def test_executor_retains_survivor_for_explicit_shutdown_retry(self, monkeypatch): + from vllm_omni.diffusion.executor import multiproc_executor as executor_module + + survivor = Mock(pid=123) + survivor.name = "surviving-worker" + survivor.is_alive.return_value = True + cleaner = executor_module._ExecutorShutdownCleaner(processes=[survivor]) + executor, _, _ = _make_executor() + executor._shutdown_cleaner = cleaner + executor._processes = [survivor] + executor._finalizer = weakref.finalize(executor, cleaner) + executor._pump_stop = threading.Event() + executor._futures_lock = threading.RLock() + executor._rpc_futures = {} + executor._output_futures = {} + executor._batch_split_map = {} + log_error = Mock() + monkeypatch.setattr(executor_module.logger, "error", log_error) + + executor.shutdown() + + assert not executor._finalizer.alive + assert executor._shutdown_cleaner is cleaner + assert executor._processes == [survivor] + log_error.assert_called_once() + assert log_error.call_args.args[1] == [("surviving-worker", 123)] + survivor.is_alive.return_value = False + + executor.shutdown() + executor.shutdown() + + assert executor._shutdown_cleaner is None + assert executor._processes == [] + + def test_concurrent_cleanup_does_not_repeat_process_operations(self): + from vllm_omni.diffusion.executor import multiproc_executor as executor_module + + joining, release = threading.Event(), threading.Event() + proc = Mock(pid=123) + proc.is_alive.return_value = True + + def join(timeout): + joining.set() + assert release.wait(5) + proc.is_alive.return_value = False + + proc.join.side_effect = join + cleaner = executor_module._ExecutorShutdownCleaner(processes=[proc]) + thread = threading.Thread(target=cleaner) + thread.start() + try: + assert joining.wait(5) + cleaner() + proc.join.assert_called_once() + finally: + release.set() + thread.join(5) + assert not thread.is_alive() + assert cleaner.processes == [] + # ───────── monitor thread & death sentinel integration tests ───────── @@ -1406,7 +1535,7 @@ def test_worker_monitor_sets_is_failed_and_calls_callbacks_on_death(self): executor._result_mq = None executor._shutdown_cleaner = None # Use a no-op so shutdown() doesn't crash on None resources. - executor._finalizer = lambda: None + executor._finalizer = weakref.finalize(executor, lambda: None) # ------------------------------------------------------------------ # Attributes added by remove_bubble_v2 (async D2H); shutdown() iterates # over them, so they need to exist even when constructed via __new__. @@ -1441,7 +1570,7 @@ def test_worker_monitor_noop_when_already_closed(self): executor._broadcast_mq = None executor._result_mq = None executor._shutdown_cleaner = None - executor._finalizer = lambda: None + executor._finalizer = weakref.finalize(executor, lambda: None) proc = _make_short_lived_process() executor._processes = [proc] diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index 39e85abcc5c..57b23656cf8 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.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 concurrent.futures @@ -10,7 +13,7 @@ import time import weakref from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import dataclass, field from multiprocessing.synchronize import Event from typing import TYPE_CHECKING, Any, cast @@ -42,6 +45,7 @@ _DLO_DP_WAVE_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_DLO_DP_WAVE_TIMEOUT", 600.0)) _WORKER_SHUTDOWN_GRACE_S = 15.0 _WORKER_TERMINATE_GRACE_S = 5.0 +_WORKER_KILL_GRACE_S = 5.0 _RESULT_PUMP_JOIN_TIMEOUT_S = 2.0 @@ -80,33 +84,61 @@ class _ExecutorShutdownCleaner: broadcast_mq: MessageQueue | None = None num_workers: int = 0 processes: list[mp.Process] | None = None + _lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False) def __call__(self) -> None: """Clean up background resources.""" + # The worker monitor and an explicit shutdown may race. A retry must + # not signal/join the same Process objects while cleanup is in flight. + if not self._lock.acquire(blocking=False): + return + try: + self._cleanup() + finally: + self._lock.release() + + def _cleanup(self) -> None: if self.broadcast_mq is not None: try: for _ in range(self.num_workers): self.broadcast_mq.enqueue(SHUTDOWN_MESSAGE, timeout=1.0) - self.broadcast_mq = None except Exception as exc: logger.warning("Failed to send shutdown signal: %s", exc) + finally: + self.broadcast_mq = None if self.processes: - join_deadline = time.monotonic() + _WORKER_SHUTDOWN_GRACE_S - for proc in self.processes: - if not proc.is_alive(): - continue - proc.join(max(0.0, join_deadline - time.monotonic())) - alive = [proc for proc in self.processes if proc.is_alive()] - for proc in alive: - logger.warning("Terminating diffusion worker %s after timeout", proc.name) - proc.terminate() + for action, grace in ( + (None, _WORKER_SHUTDOWN_GRACE_S), + ("terminate", _WORKER_TERMINATE_GRACE_S), + ("kill", _WORKER_KILL_GRACE_S), + ): + if not alive: + break + if action is not None: + for proc in alive: + try: + logger.warning("Calling %s on diffusion worker %s (pid=%s)", action, proc.name, proc.pid) + getattr(proc, action)() + except OSError: + logger.exception("Failed to %s diffusion worker %s (pid=%s)", action, proc.name, proc.pid) - terminate_deadline = time.monotonic() + _WORKER_TERMINATE_GRACE_S - for proc in alive: - proc.join(max(0.0, terminate_deadline - time.monotonic())) + deadline = time.monotonic() + grace + for proc in alive: + try: + proc.join(max(0.0, deadline - time.monotonic())) + except OSError: + logger.exception("Failed to join diffusion worker %s (pid=%s)", proc.name, proc.pid) + alive = [proc for proc in alive if proc.is_alive()] + + self.processes = alive + if alive: + logger.error( + "Diffusion worker cleanup incomplete after kill: %s; retaining processes for shutdown retry", + [(proc.name, proc.pid) for proc in alive], + ) class MultiprocDiffusionExecutor(DiffusionExecutor): @@ -1006,8 +1038,12 @@ def check_health(self) -> None: def shutdown(self) -> None: self._closed = True self._pump_stop.set() + cleaner = self._shutdown_cleaner try: - self._finalizer() + if self._finalizer.alive: + self._finalizer() + elif cleaner is not None: + cleaner() finally: pump_threads = getattr(self, "_result_pump_threads", []) for thread in pump_threads: @@ -1031,5 +1067,6 @@ def shutdown(self) -> None: self._rpc_futures.clear() self._output_futures.clear() self._batch_split_map.clear() - self._shutdown_cleaner = None - self._processes = [] + self._processes = (cleaner.processes or []) if cleaner is not None else [] + if not self._processes: + self._shutdown_cleaner = None From f11bf1a05d8dff5d9b03d0012fb9d3dca3bc6b05 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Thu, 10 Sep 2026 17:57:20 +0800 Subject: [PATCH 55/93] [Misc] Add standalone safetensors retention diagnostic (#7145) Signed-off-by: Hongsheng Liu --- benchmarks/host_weight_runtime/README.md | 57 +++++++++ .../safetensors_retention.py | 120 ++++++++++++++++++ docs/design/feature/host_weight_runtime.md | 5 + 3 files changed, 182 insertions(+) create mode 100644 benchmarks/host_weight_runtime/README.md create mode 100644 benchmarks/host_weight_runtime/safetensors_retention.py diff --git a/benchmarks/host_weight_runtime/README.md b/benchmarks/host_weight_runtime/README.md new file mode 100644 index 00000000000..8998111518f --- /dev/null +++ b/benchmarks/host_weight_runtime/README.md @@ -0,0 +1,57 @@ +# Host weight dependency memory diagnostic + +`safetensors_retention.py` isolates repeated CPU `get_tensor()` calls from +access to one cached tensor view. It imports neither vLLM nor HWR and uses a +synthetic float32 payload controlled by `--tensor-elements` (default: 64). +Tensor size can affect retention: a negative result for 16 elements does not +rule out growth for 64 elements. The selected size is recorded in the JSON +arguments and used for both file creation and the correctness check. +HWR itself calls `get_tensor()` when +acquiring a lease, then exposes cached views; this probe does not measure HWR +requests or repeated lease acquisition. + +Run one mode per fresh process in the same Python environment and CPU affinity: + +```bash +git rev-parse HEAD > revision.txt +git status --short > working-tree.txt +git diff > working-tree.patch +# Select an allowed CPU from your process affinity; 0 is only an example. +CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \ + benchmarks/host_weight_runtime/safetensors_retention.py \ + --mode get_tensor --iterations 10000 --sample-every 5000 > feasibility.json +``` + +For a size comparison, run 16 and 64 elements twice each in fresh processes, +with cached reuse as a control at both sizes: + +```bash +for repetition in 1 2; do + for elements in 16 64; do + for mode in get_tensor reuse; do + CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \ + benchmarks/host_weight_runtime/safetensors_retention.py \ + --mode "$mode" --tensor-elements "$elements" \ + > "$mode-$elements-$repetition.json" + done + done +done +``` + +The diagnostic sets Torch to one CPU thread, warms up before the baseline, +and reports preparation separately from timed loop work. Checkpoints run GC +and report Linux private/anonymous memory, RSS, and open descriptors. It checks +that transient tensor objects and the final cached view are released. Linux +`/proc/self/smaps_rollup` must be readable; invalid arguments and missing proc +support fail explicitly. Keep raw JSON and the repository snapshot with results. + +Compare the final loop sample with the warmed baseline, then inspect the closed +sample. Private-memory growth after Python tensor objects disappear is evidence +of retention in the dependency/allocator path, not proof of a leak's root cause. +RSS can include file-backed pages; it is not a private-memory metric. Samples +also include small diagnostic bookkeeping allocations. Do not assert a portable +memory threshold or turn this synthetic loop into a per-request service estimate. + +Do not drop shared caches, trim allocators, or add serving-time GC to make the +numbers look smaller. A dependency replacement or version change requires its +own reproducible comparison and correctness/lifecycle validation. diff --git a/benchmarks/host_weight_runtime/safetensors_retention.py b/benchmarks/host_weight_runtime/safetensors_retention.py new file mode 100644 index 00000000000..55795e7ef80 --- /dev/null +++ b/benchmarks/host_weight_runtime/safetensors_retention.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Linux CPU diagnostic: repeated safetensors materialization versus view reuse.""" + +from __future__ import annotations + +import argparse +import gc +import json +import os +import platform +import tempfile +import time +import weakref +from pathlib import Path + + +def positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("must be positive") + return parsed + + +def sample(phase: str, iterations: int, loop_seconds: float) -> dict[str, object]: + gc.collect() + counters = {} + for line in Path("/proc/self/smaps_rollup").read_text().splitlines(): + fields = line.split() + if len(fields) == 3 and fields[2] == "kB": + counters[fields[0].removesuffix(":")] = int(fields[1]) * 1024 + return { + "phase": phase, + "iterations": iterations, + "loop_seconds": loop_seconds, + "rss_bytes": counters["Rss"], + "private_bytes": counters["Private_Clean"] + counters["Private_Dirty"], + "anonymous_bytes": counters["Anonymous"], + "fd_count": len(tuple(Path("/proc/self/fd").iterdir())), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--mode", choices=("get_tensor", "reuse"), required=True) + parser.add_argument( + "--tensor-elements", type=positive_int, default=64, help="Number of float32 elements in the synthetic tensor" + ) + parser.add_argument("--iterations", type=positive_int, default=1_000_000) + parser.add_argument("--sample-every", type=positive_int, default=100_000) + parser.add_argument("--warmup", type=positive_int, default=1000) + args = parser.parse_args() + if not Path("/proc/self/smaps_rollup").is_file(): + parser.error("requires Linux /proc/self/smaps_rollup") + + preparation_start = time.monotonic() + import safetensors + import torch + from safetensors import safe_open + from safetensors.torch import save_file + + torch.set_num_threads(1) + torch.set_num_interop_threads(1) + samples = [sample("before_file", 0, 0.0)] + with tempfile.TemporaryDirectory(prefix="safetensors-retention-") as directory: + path = Path(directory) / "synthetic.safetensors" + save_file({"weight": torch.arange(args.tensor_elements, dtype=torch.float32)}, str(path)) + with safe_open(path, framework="pt", device="cpu") as reader: + cached = reader.get_tensor("weight") + assert torch.equal(cached, torch.arange(args.tensor_elements, dtype=torch.float32)) + for _ in range(args.warmup): + tensor = reader.get_tensor("weight") if args.mode == "get_tensor" else cached + del tensor + samples.append(sample("baseline", 0, 0.0)) + preparation_seconds = time.monotonic() - preparation_start + completed = 0 + loop_seconds = 0.0 + while completed < args.iterations: + count = min(args.sample_every, args.iterations - completed) + start = time.monotonic() + for _ in range(count): + tensor = reader.get_tensor("weight") if args.mode == "get_tensor" else cached + last_tensor = weakref.ref(tensor) + del tensor + loop_seconds += time.monotonic() - start + completed += count + snapshot = sample("loop", completed, loop_seconds) + snapshot["last_tensor_alive"] = last_tensor() is not None + if args.mode == "get_tensor": + assert last_tensor() is None, "diagnostic retained a transient tensor" + samples.append(snapshot) + del cached + del reader + samples.append(sample("closed", completed, loop_seconds)) + assert last_tensor() is None, "last tensor survived reader/view teardown" + + print( + json.dumps( + { + "schema_version": 1, + "arguments": vars(args), + "python": platform.python_version(), + "torch": torch.__version__, + "safetensors": safetensors.__version__, + "kernel": platform.release(), + "hostname": platform.node(), + "pid": os.getpid(), + "cpu_affinity": sorted(os.sched_getaffinity(0)), + "torch_threads": torch.get_num_threads(), + "pythonmalloc": os.environ.get("PYTHONMALLOC", "default"), + "preparation_seconds": preparation_seconds, + "samples": samples, + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/docs/design/feature/host_weight_runtime.md b/docs/design/feature/host_weight_runtime.md index ac7bc7b3f9d..15fce07b80a 100644 --- a/docs/design/feature/host_weight_runtime.md +++ b/docs/design/feature/host_weight_runtime.md @@ -42,6 +42,11 @@ V1 does not include: ## Motivation and use cases +For investigating CPU memory retained by dependency tensor materialization, +see the [standalone safetensors diagnostic](https://github.com/vllm-project/vllm-omni/blob/main/benchmarks/host_weight_runtime/README.md). +It distinguishes repeated dependency calls from reuse of cached views and does +not establish per-request HWR leakage. + Model loading can create the same final host representation repeatedly. This is especially expensive when loading performs checkpoint decoding, tensor renaming, TP slicing, quantization, packing, or scale construction before GPU From d8d4062547899d71b9073649f404e61658136a68 Mon Sep 17 00:00:00 2001 From: andyluo7 <43718156+andyluo7@users.noreply.github.com> Date: Thu, 10 Sep 2026 02:58:24 -0700 Subject: [PATCH 56/93] [CI] Isolate layerwise offload memory measurements (#6938) Signed-off-by: andyluo7 --- .../test_diffusion_layerwise_offload.py | 269 +++++++++++++++--- 1 file changed, 231 insertions(+), 38 deletions(-) diff --git a/tests/diffusion/offloader/test_diffusion_layerwise_offload.py b/tests/diffusion/offloader/test_diffusion_layerwise_offload.py index 8e083605f66..c386871d3b0 100644 --- a/tests/diffusion/offloader/test_diffusion_layerwise_offload.py +++ b/tests/diffusion/offloader/test_diffusion_layerwise_offload.py @@ -1,7 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import concurrent.futures import gc +import multiprocessing as mp +import os from typing import Any import numpy as np @@ -45,32 +48,98 @@ "sampler_params": {"height": 480, "width": 640, "num_frames": 5}, } +_OFFLOAD_STATE_PROBE = "vllm_omni_layerwise_offload_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" +_MEASUREMENT_ORDERS: dict[str, tuple[bool, bool]] = { + "baseline-first": (False, True), + "offload-first": (True, False), +} + + +class OffloadStateProbe: + """Test-only worker extension exposing the configured offloader state.""" + + model_runner: Any + + def get_offload_state_for_test(self) -> dict[str, Any]: + backend = getattr(self.model_runner, "offload_backend", None) + block_groups = getattr(backend, "_blocks", ()) + group_sizes = [len(group) for group in block_groups] if isinstance(block_groups, (list, tuple)) else [] + + return { + "probe": _OFFLOAD_STATE_PROBE, + "backend_type": ( + None if backend is None else f"{backend.__class__.__module__}.{backend.__class__.__qualname__}" + ), + "requested": bool(getattr(self.model_runner.od_config, "enable_layerwise_offload", False)), + "enabled": bool(backend is not None and backend.is_enabled()), + "block_group_count": len(group_sizes), + "block_group_sizes": group_sizes, + "block_count": sum(group_sizes), + } -def check_audio_determinism(audio1, audio2, atol=1e-2): - device = current_omni_platform.device_type - if isinstance(audio1, np.ndarray): - audio1 = torch.from_numpy(audio1).to(device) - if isinstance(audio2, np.ndarray): - audio2 = torch.from_numpy(audio2).to(device) - if not torch.allclose(audio1, audio2, atol=atol): - diff = torch.abs(audio1 - audio2) - print(f"Max difference: {diff.max().item()}") - print(f"Mean difference: {diff.mean().item()}") +def check_audio_determinism(audio1: np.ndarray, audio2: np.ndarray, atol: float = 1e-2) -> bool: + if not np.allclose(audio1, audio2, atol=atol): + diff = np.abs(audio1 - audio2) + print(f"Max difference: {diff.max()}") + print(f"Mean difference: {diff.mean()}") raise AssertionError(f"Audio outputs differ beyond tolerance atol={atol}") 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 + multimodal_output = getattr(output[0], "multimodal_output", None) + if not isinstance(multimodal_output, dict): + return None + audio = multimodal_output.get("audio") + if audio is None: + return None + if isinstance(audio, torch.Tensor): + return audio.detach().cpu().numpy() + return np.asarray(audio) + + +def _collect_offload_states(value: Any) -> list[dict[str, Any]]: + if isinstance(value, dict): + return [value] if value.get("probe") == _OFFLOAD_STATE_PROBE else [] + if isinstance(value, (list, tuple)): + states: list[dict[str, Any]] = [] + for item in value: + states.extend(_collect_offload_states(item)) + return states + return [] + + +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") + try: + return _MEASUREMENT_ORDERS[configured] + except KeyError as exc: + allowed = ", ".join(_MEASUREMENT_ORDERS) + raise ValueError(f"{_MEASUREMENT_ORDER_ENV} must be one of: {allowed}; got {configured!r}") from exc + + def run_inference( model_name: str, layerwise_offload: bool = False, num_inference_steps: int = 3, -) -> tuple[float, Any]: +) -> dict[str, Any]: current_omni_platform.empty_cache() device_index = current_omni_platform.current_device() - with current_omni_platform.device(device_index): - free_bytes, total_bytes = current_omni_platform.mem_get_info() - initial_used_mb = (total_bytes - free_bytes) / (1024**2) + initial_used_mb = _device_used_mb(device_index) if model_name in AUDIO_MODEL: params = AUDIO_MODEL_PARAMS @@ -80,10 +149,17 @@ def run_inference( with OmniRunner( model_name, enable_layerwise_offload=layerwise_offload, + worker_extension_cls=f"{OffloadStateProbe.__module__}.{OffloadStateProbe.__qualname__}", # TODO: we might want to add overlapped feature e2e tests # cache_backend="cache_dit", **params["runner_params"], ) as runner: + offload_states = _collect_offload_states( + runner.omni.engine.collective_rpc(method="get_offload_state_for_test", timeout=60) + ) + 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 @@ -108,20 +184,136 @@ def run_inference( finally: monitor.stop() + audio = _extract_audio(output) + del output + # DeviceMemoryMonitor reports absolute device usage. Subtract this run's - # starting usage so process-wide compiler/workspace caches retained from a - # previous run do not make the second measurement order-dependent. - peak = max(0.0, monitor.peak_used_mb - initial_used_mb) + # 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, + "audio": audio, + "offload_states": offload_states, + } + + +def run_inference_isolated( + model_name: str, + layerwise_offload: bool = False, + num_inference_steps: int = 3, +) -> dict[str, Any]: + """Run one measurement in a fresh interpreter and accelerator context.""" + with concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=mp.get_context("spawn")) as executor: + return executor.submit(run_inference, model_name, layerwise_offload, num_inference_steps).result() + + +def _assert_offload_state(measurement: dict[str, Any], *, expected_enabled: bool) -> None: + states = measurement["offload_states"] + assert states, "Expected at least one offload-state result" + for state in states: + assert state["requested"] is expected_enabled, f"Unexpected offload request state: {state}" + assert state["enabled"] is expected_enabled, f"Unexpected offload state: {state}" + if expected_enabled: + assert state["backend_type"], f"Enabled offloader has no backend type: {state}" + assert state["block_group_count"] > 0, f"Enabled offloader has no block groups: {state}" + assert state["block_count"] > 0, f"Enabled offloader has no managed blocks: {state}" + else: + assert state["block_group_count"] == 0, f"Disabled offloader retained block groups: {state}" + assert state["block_count"] == 0, f"Disabled offloader retained managed blocks: {state}" + + +def _assert_audio_outputs( + model_name: str, + audio_no_offload: np.ndarray | None, + audio_offload: np.ndarray | None, +) -> None: + if model_name not in AUDIO_MODEL: + return + + assert audio_no_offload is not None, "Baseline Stable Audio inference returned no audio" + assert audio_offload is not None, "Layerwise-offloaded Stable Audio inference returned no audio" + # Match the sibling cpu-offload test's tolerance: layerwise offload moves + # blocks across the PCIe bus on a side stream, which can perturb cuBLAS + # algorithm selection and produce ~ULP-level drift larger than 1e-3. + check_audio_determinism(audio_offload, audio_no_offload, atol=1e-2) + + +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} offload state: {measurement['offload_states']}") - return peak, output + +@pytest.mark.diffusion +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize( + ("configured", "expected"), + [ + (None, (False, True)), + ("baseline-first", (False, True)), + ("offload-first", (True, False)), + ], +) +def test_measurement_order(monkeypatch, configured: str | None, expected: tuple[bool, bool]) -> None: + if configured is None: + monkeypatch.delenv(_MEASUREMENT_ORDER_ENV, raising=False) + else: + monkeypatch.setenv(_MEASUREMENT_ORDER_ENV, configured) + + assert _measurement_order() == expected + + +@pytest.mark.diffusion +@pytest.mark.core_model +@pytest.mark.cpu +def test_measurement_order_rejects_unknown_value(monkeypatch) -> None: + monkeypatch.setenv(_MEASUREMENT_ORDER_ENV, "unknown") + + with pytest.raises(ValueError, match=_MEASUREMENT_ORDER_ENV): + _measurement_order() + + +@pytest.mark.diffusion +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize( + ("audio_no_offload", "audio_offload", "message"), + [ + (None, None, "Baseline Stable Audio inference returned no audio"), + (None, np.zeros(1), "Baseline Stable Audio inference returned no audio"), + (np.zeros(1), None, "Layerwise-offloaded Stable Audio inference returned no audio"), + ], +) +def test_stable_audio_requires_both_outputs(audio_no_offload, audio_offload, message) -> None: + with pytest.raises(AssertionError, match=message): + _assert_audio_outputs("stabilityai/stable-audio-open-1.0", audio_no_offload, audio_offload) @pytest.mark.diffusion @hardware_test(res={"cuda": "L4", "rocm": "MI325"}) -@pytest.mark.parametrize("model_name", list(MODELS.keys())) +@pytest.mark.parametrize( + "model_name", + [pytest.param(name, marks=MODEL_MARKS[name]) for name in MODELS], +) def test_layerwise_offload_diffusion_model(model_name: str): """Test that layerwise offloading reduces GPU memory usage. @@ -130,14 +322,15 @@ def test_layerwise_offload_diffusion_model(model_name: str): offloader keeps only a single transformer block on GPU at a time, with prefetching for compute-memory overlap. """ + measurements: dict[bool, dict[str, Any]] = {} try: - # Run without layerwise offloading (baseline) - no_offload_peak_memory, output_no_offload = run_inference(model_name, layerwise_offload=False) - cleanup_dist_env_and_memory() - - # Run with layerwise offloading (1 layer on device) - layerwise_offload_peak_memory, output_offload = run_inference(model_name, layerwise_offload=True) - cleanup_dist_env_and_memory() + measurement_order = _measurement_order() + print( + "Measurement order: " + + " -> ".join("layerwise-offload" if enabled else "baseline" for enabled in measurement_order) + ) + for enabled in measurement_order: + measurements[enabled] = run_inference_isolated(model_name, layerwise_offload=enabled) except ValueError as exc: # omni_snapshot_download wraps GatedRepoError in a ValueError; skip instead of failing. if "Access to model" in str(exc) and "is restricted" in str(exc): @@ -145,20 +338,18 @@ def test_layerwise_offload_diffusion_model(model_name: str): f"Skipping: gated HF repo {model_name!r} inaccessible " f"({exc}). See docs/contributing/ci/hf_credentials.md." ) - pytest.fail(f"Inference failed: {exc}") - except Exception: - pytest.fail("Inference failed") + raise - print(f"Layerwise offload peak memory (1 GPU layer): {layerwise_offload_peak_memory} MB") - print(f"No offload peak memory: {no_offload_peak_memory} MB") + no_offload = measurements[False] + layerwise_offload = measurements[True] + _assert_offload_state(no_offload, expected_enabled=False) + _assert_offload_state(layerwise_offload, expected_enabled=True) + _print_measurement("No offload", no_offload) + _print_measurement("Layerwise offload", layerwise_offload) - if model_name == "stabilityai/stable-audio-open-1.0": - audio_offload = output_offload[0].multimodal_output.get("audio") - audio_no_offload = output_no_offload[0].multimodal_output.get("audio") - # Match the sibling cpu-offload test's tolerance: layerwise offload moves - # blocks across the PCIe bus on a side stream, which can perturb cuBLAS - # algorithm selection and produce ~ULP-level drift larger than 1e-3. - check_audio_determinism(audio_offload, audio_no_offload, atol=1e-2) + audio_no_offload = no_offload["audio"] + audio_offload = layerwise_offload["audio"] + _assert_audio_outputs(model_name, audio_no_offload, audio_offload) is_rocm = torch.version.hip is not None platform = "rocm" if is_rocm else "cuda" @@ -170,6 +361,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"] 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 43a476c1e5817cedfdf706a6d62c867a43cd5065 Mon Sep 17 00:00:00 2001 From: wkutak Date: Thu, 10 Sep 2026 13:51:09 +0200 Subject: [PATCH 57/93] [Model] Add Cosmos3 mixed W8A8/W8A16 and W4A4/W4A16 denoising (#6560) Signed-off-by: Rahul Steiger Signed-off-by: Wojciech Kutak Co-authored-by: Rahul Steiger --- docs/user_guide/quantization/modelopt.md | 85 +++ recipes/cosmos3/Cosmos3-Nano.md | 32 + .../test_cosmos3_mixed_precision.py | 690 ++++++++++++++++++ .../cosmos3/mixed_precision/__init__.py | 14 + .../cosmos3/mixed_precision/checkpoint.py | 176 +++++ .../models/cosmos3/mixed_precision/config.py | 106 +++ .../models/cosmos3/mixed_precision/runtime.py | 116 +++ .../cosmos3/mixed_precision/strategy.py | 218 ++++++ .../models/cosmos3/pipeline_cosmos3.py | 24 +- .../models/cosmos3/transformer_cosmos3.py | 52 ++ 10 files changed, 1509 insertions(+), 4 deletions(-) create mode 100644 tests/diffusion/quantization/test_cosmos3_mixed_precision.py create mode 100644 vllm_omni/diffusion/models/cosmos3/mixed_precision/__init__.py create mode 100644 vllm_omni/diffusion/models/cosmos3/mixed_precision/checkpoint.py create mode 100644 vllm_omni/diffusion/models/cosmos3/mixed_precision/config.py create mode 100644 vllm_omni/diffusion/models/cosmos3/mixed_precision/runtime.py create mode 100644 vllm_omni/diffusion/models/cosmos3/mixed_precision/strategy.py diff --git a/docs/user_guide/quantization/modelopt.md b/docs/user_guide/quantization/modelopt.md index d8da4e1a92f..07d05730e1f 100644 --- a/docs/user_guide/quantization/modelopt.md +++ b/docs/user_guide/quantization/modelopt.md @@ -204,6 +204,91 @@ omni = Omni( | `--linear-backend cutlass` | str | auto | Select the validated CUTLASS linear backend for supported ModelOpt NVFP4 or mixed FP8/NVFP4 diffusion stages | | `--moe-backend cutlass` | str | auto | Select the validated CUTLASS MoE backend for supported ModelOpt mixed MoE checkpoints | +## Cosmos3 Mixed-Precision Schedule + +The experimental Cosmos3 schedule can use the checkpoint-native quantized GEMM for middle denoising +steps and dense 16-bit activation execution for selected first and last steps. +The schedule supports serialized ModelOpt FP8 and NVFP4 checkpoints: + +| Detected weight format | Native middle steps | A16 boundary steps | +|------------------------|---------------------|--------------------| +| FP8 | W8A8 | W8A16 | +| NVFP4 | W4A4 | W4A16 | + +The native and A16 paths share one quantized weight representation. The A16 +path dequantizes the live backend weight and calls `F.linear`; it does not keep +a second checkpoint-weight snapshot. Scheduled FP8 currently requires +serialized tensorwise scales and a backend that retains canonical FP8 weights. +Scheduled NVFP4 requires a supported CUTLASS-compatible native or FlashInfer +layout. SmoothQuant, Marlin-repacked FP8, and per-channel/per-token FP8 +schedules fail closed. + +The reasoner uses dense A16 execution by default when its weights are FP8 or +NVFP4. Set the nested `reasoner` field to `native` to retain the checkpoint-native +W8A8/W4A4 path. A BF16 reasoner remains BF16 in either mode. The schedule +currently requires tensor parallel size 1 and one active request per worker, +does not support HSDP, and does not support block-scaled FP8. Its live-weight +design is compatible with model-level and standard layer-wise offload, but +those paths still require GPU end-to-end validation. Distributed layer-wise +offload is rejected because its direct loader bypasses ModelOpt post-load +transformations. + +A checkpoint can opt in by adding a versioned policy to its authoritative +`transformer/config.json`. The same policy schema is used for FP8 and NVFP4: + +```json +{ + "quantization_config": { + "quant_method": "modelopt", + "quant_algo": "FP8", + "runtime": { + "diffusion_step_policy": { + "schema_version": 1, + "type": "first_last_n", + "index_space": "denoising_loop_iteration", + "scope": ["transformer"], + "default_mode": "native", + "first_steps": {"count": 3, "mode": "a16"}, + "last_steps": {"count": 3, "mode": "a16"}, + "overlap": "a16", + "reasoner": "a16" + } + } + } +} +``` + +The example preserves the checkpoint's existing `quant_method` and +`quant_algo`; only `runtime` is added. + +When present and valid, the checkpoint policy is the default. An explicit +runtime object overrides it. To retain ordinary checkpoint-native execution +for every step, disable only the schedule: + +```bash +vllm serve /path/to/Cosmos3-Nano-modelopt \ + --omni \ + --additional-config '{"cosmos3_mixed_precision":{"enabled":false}}' +``` + +This does not disable checkpoint quantization; it prevents installation of the +mixed-step runtime. + +```bash +vllm serve /path/to/Cosmos3-Nano-modelopt \ + --omni \ + --additional-config \ + '{"cosmos3_mixed_precision":{"first_steps":3,"last_steps":3,"reasoner":"a16"}}' +``` + +The presence of `cosmos3_mixed_precision` supplies an explicit runtime +override; an empty object uses the defaults shown above. +The schedule uses vLLM's selected native backend and fails during loading if +its live weight layout cannot support A16 dequantization. +The runtime selects FP8 or NVFP4 from the checkpoint's ModelOpt configuration. +BF16 linears are untouched. Mixed FP8/NVFP4 checkpoints are not supported by +this schedule. + ## Validation and Notes 1. Compare the ModelOpt checkpoint against the BF16 baseline with the same diff --git a/recipes/cosmos3/Cosmos3-Nano.md b/recipes/cosmos3/Cosmos3-Nano.md index 490f5b1f4cc..f81e5e887f6 100644 --- a/recipes/cosmos3/Cosmos3-Nano.md +++ b/recipes/cosmos3/Cosmos3-Nano.md @@ -115,6 +115,38 @@ shift at the same seed). The pipeline auto-resolves from `model_index.json`; pass `--model-class-name Cosmos3OmniDiffusersPipeline` to force it explicitly. +For a serialized ModelOpt FP8 or NVFP4 checkpoint, an experimental +mixed-precision schedule can use native W8A8/W4A4 in middle denoising steps +and dense W8A16/W4A16 in the first and last steps: + +```bash +vllm serve /path/to/Cosmos3-Nano-modelopt \ + --omni \ + --additional-config \ + '{"cosmos3_mixed_precision":{"first_steps":3,"last_steps":3,"reasoner":"a16"}}' +``` + +The nested object's presence supplies an explicit runtime override; use an +empty object for these defaults. A compatible checkpoint can instead declare +the versioned `runtime.diffusion_step_policy` in its authoritative +`transformer/config.json`. The checkpoint policy is used only when no runtime +override is present. Use +`--additional-config '{"cosmos3_mixed_precision":{"enabled":false}}'` to +disable the checkpoint schedule without disabling checkpoint quantization. +This path currently requires tensor parallel size 1 and does not support HSDP +or block-scaled FP8. It keeps one live quantized weight representation: +scheduled FP8 requires serialized tensorwise scales and a canonical backend, +while scheduled NVFP4 requires a supported CUTLASS-compatible native or +FlashInfer layout. Its dequantization is a correctness baseline; no speedup is +implied. Quantized reasoner weights use A16 by default; set +`"reasoner":"native"` in the nested object to keep W8A8/W4A4. +The checkpoint's ModelOpt configuration selects FP8 or NVFP4; mixed FP8/NVFP4 +checkpoints are not supported by this schedule. BF16 linears are left unchanged. +The schedule currently supports one active request per worker. Model-level and +layer-wise offload are designed to work with this live-weight path but still +require GPU validation. Distributed layer-wise offload and incompatible native +weight layouts fail during loading. + #### Verification Best quality uses the JSON-upsampled prompts from `assets/` (download with diff --git a/tests/diffusion/quantization/test_cosmos3_mixed_precision.py b/tests/diffusion/quantization/test_cosmos3_mixed_precision.py new file mode 100644 index 00000000000..026090fe99c --- /dev/null +++ b/tests/diffusion/quantization/test_cosmos3_mixed_precision.py @@ -0,0 +1,690 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.diffusion.models.cosmos3.mixed_precision import ( + Cosmos3MixedPrecisionConfig, + Cosmos3MixedPrecisionRuntime, + resolve_mixed_precision_config, +) +from vllm_omni.diffusion.models.cosmos3.mixed_precision import runtime as runtime_impl +from vllm_omni.diffusion.models.cosmos3.mixed_precision.checkpoint import read_checkpoint_policy +from vllm_omni.diffusion.models.cosmos3.mixed_precision.runtime import ( + Cosmos3MixedPrecisionLinearMethod, +) +from vllm_omni.diffusion.models.cosmos3.mixed_precision.strategy import ( + Fp8W8A8W8A16Strategy, + Nvfp4W4A4W4A16Strategy, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _policy( + *, + first_steps: int = 3, + last_steps: int = 3, + reasoner: str = "a16", + scope: list[str] | None = None, +) -> dict[str, object]: + return { + "schema_version": 1, + "type": "first_last_n", + "index_space": "denoising_loop_iteration", + "scope": scope or ["transformer"], + "default_mode": "native", + "first_steps": {"count": first_steps, "mode": "a16"}, + "last_steps": {"count": last_steps, "mode": "a16"}, + "overlap": "a16", + "reasoner": reasoner, + } + + +def _fake_quant_config(name: str, *, serialized: bool = True): + if name == "modelopt": + return SimpleNamespace( + get_name=lambda: name, + is_checkpoint_fp8_serialized=serialized, + ) + if name == "modelopt_fp4": + return SimpleNamespace( + get_name=lambda: name, + quant_method="NVFP4", + is_checkpoint_nvfp4_serialized=serialized, + ) + return SimpleNamespace(get_name=lambda: name) + + +def _checkpoint_od_config( + quant_config, + *, + policy: object | None = None, + additional_config: dict[str, object] | None = None, +): + quant_config_name = quant_config.get_name() + quant_algo = { + "modelopt": "FP8", + "modelopt_fp4": "NVFP4", + }.get(quant_config_name, "FP8") + disk_quant_config: dict[str, object] = { + "quant_method": quant_config_name, + "quant_algo": quant_algo, + } + if policy is not None: + disk_quant_config["runtime"] = {"diffusion_step_policy": policy} + return SimpleNamespace( + quantization_config=quant_config, + tf_model_config=SimpleNamespace( + params={"quantization_config": disk_quant_config}, + quant_config=quant_config, + ), + additional_config=additional_config, + ) + + +def _swizzle_blockscale_cpu(scale: torch.Tensor) -> torch.Tensor: + rows, cols = scale.shape + padded_rows = ((rows + 127) // 128) * 128 + padded_cols = ((cols + 3) // 4) * 4 + padded = torch.zeros( + (1, padded_rows, padded_cols), + dtype=scale.dtype, + device=scale.device, + ) + padded[0, :rows, :cols] = scale + padded = padded.reshape(1, padded_rows // 128, 4, 32, padded_cols // 4, 4) + return ( + padded.permute(0, 1, 4, 3, 2, 5) + .contiguous() + .reshape( + padded_rows, + padded_cols, + ) + ) + + +def test_config_parses_nested_schedule_and_reasoner_policy() -> None: + config = Cosmos3MixedPrecisionConfig.from_additional_config( + { + "cosmos3_mixed_precision": { + "first_steps": 2, + "last_steps": 4, + "reasoner": "native", + } + } + ) + + assert config is not None + assert config.reasoner == "native" + assert [index for index in range(10) if config.use_high_precision(index, 10)] == [0, 1, 6, 7, 8, 9] + + +@pytest.mark.parametrize( + ("first_steps", "last_steps", "selected"), + [ + (0, 2, [5, 6]), + (2, 0, [0, 1]), + (0, 0, []), + (4, 4, list(range(7))), + ], +) +def test_schedule_boundaries_and_overlap( + first_steps: int, + last_steps: int, + selected: list[int], +) -> None: + config = Cosmos3MixedPrecisionConfig( + first_steps=first_steps, + last_steps=last_steps, + ) + assert [index for index in range(7) if config.use_high_precision(index, 7)] == selected + + +def test_one_step_request_honors_boundary_precision() -> None: + config = Cosmos3MixedPrecisionConfig(first_steps=1, last_steps=1) + assert config.use_high_precision(0, 1) + + +@pytest.mark.parametrize( + ("values", "message"), + [ + ({"cosmos3_mixed_precision": True}, "must be a mapping"), + ({"cosmos3_mixed_precision": {"first_steps": -1}}, "non-negative"), + ({"cosmos3_mixed_precision": {"reasoner": "fp16"}}, "must be one of"), + ({"cosmos3_mixed_precision": {"enabled": 0}}, "must be a boolean"), + ( + {"cosmos3_mixed_precision": {"enabled": False, "first_steps": 1}}, + "cannot be combined", + ), + ({"cosmos3_mixed_precision": {"unknown": 1}}, "Unknown"), + ], +) +def test_config_rejects_invalid_values(values: dict, message: str) -> None: + with pytest.raises((TypeError, ValueError), match=message): + Cosmos3MixedPrecisionConfig.from_additional_config(values) + + +def test_config_presence_enables_defaults() -> None: + assert Cosmos3MixedPrecisionConfig.from_additional_config({}) is None + config = Cosmos3MixedPrecisionConfig.from_additional_config({"cosmos3_mixed_precision": {}}) + assert config == Cosmos3MixedPrecisionConfig() + + +def test_noop_config_does_not_install_runtime() -> None: + assert ( + Cosmos3MixedPrecisionConfig.from_additional_config( + { + "cosmos3_mixed_precision": { + "first_steps": 0, + "last_steps": 0, + "reasoner": "native", + } + } + ) + is None + ) + + +def test_additional_config_distinguishes_absent_and_disabled() -> None: + assert Cosmos3MixedPrecisionConfig.resolve_additional_config({}) == (False, None) + assert Cosmos3MixedPrecisionConfig.resolve_additional_config({"cosmos3_mixed_precision": {"enabled": False}}) == ( + True, + None, + ) + assert Cosmos3MixedPrecisionConfig.resolve_additional_config( + { + "cosmos3_mixed_precision": { + "first_steps": 0, + "last_steps": 0, + "reasoner": "native", + } + } + ) == (True, None) + + +def test_checkpoint_policy_round_trip_through_transformer_config() -> None: + from vllm_omni.diffusion.data import TransformerConfig + + disk = { + "_class_name": "Cosmos3VFMTransformer", + "quantization_config": { + "quant_method": "modelopt", + "quant_algo": "FP8", + "runtime": {"diffusion_step_policy": _policy(first_steps=2, last_steps=4)}, + }, + } + + tf_config = TransformerConfig.from_dict(disk) + assert tf_config.quant_config is not None + assert tf_config.quant_config.get_name() == "modelopt" + assert tf_config.to_dict()["quantization_config"]["runtime"] == disk["quantization_config"]["runtime"] + assert read_checkpoint_policy(SimpleNamespace(tf_model_config=tf_config)) == Cosmos3MixedPrecisionConfig( + first_steps=2, + last_steps=4, + reasoner="a16", + ) + + +@pytest.mark.parametrize("name", ["modelopt", "modelopt_fp4"]) +def test_checkpoint_policy_supports_serialized_modelopt_formats(name: str) -> None: + od_config = _checkpoint_od_config(_fake_quant_config(name), policy=_policy()) + assert read_checkpoint_policy(od_config) == Cosmos3MixedPrecisionConfig() + + +@pytest.mark.parametrize("name", ["modelopt", "modelopt_fp4"]) +def test_checkpoint_policy_rejects_nonserialized_modelopt_formats(name: str) -> None: + od_config = _checkpoint_od_config( + _fake_quant_config(name, serialized=False), + policy=_policy(), + ) + with pytest.raises(ValueError, match="serialized ModelOpt"): + read_checkpoint_policy(od_config) + + +def test_checkpoint_policy_rejects_mixed_modelopt_format() -> None: + with pytest.raises(ValueError, match="serialized ModelOpt FP8 or NVFP4"): + read_checkpoint_policy(_checkpoint_od_config(_fake_quant_config("modelopt_mixed"), policy=_policy())) + + +def test_checkpoint_policy_missing_metadata_preserves_ordinary_path() -> None: + assert read_checkpoint_policy(_checkpoint_od_config(_fake_quant_config("modelopt"))) is None + + +def test_checkpoint_policy_rejects_malformed_runtime_container() -> None: + od_config = _checkpoint_od_config(_fake_quant_config("modelopt")) + od_config.tf_model_config.params["quantization_config"]["runtime"] = True + with pytest.raises(TypeError, match="runtime must be a mapping"): + read_checkpoint_policy(od_config) + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("schema_version", 2, "schema_version"), + ("type", "sigma_ranges", "type"), + ("index_space", "scheduler_timestep", "index_space"), + ("default_mode", "a16", "default_mode"), + ("overlap", "native", "overlap"), + ("reasoner", "fp16", "reasoner"), + ], +) +def test_checkpoint_policy_rejects_invalid_schema(field: str, value: object, message: str) -> None: + policy = _policy() + policy[field] = value + with pytest.raises(ValueError, match=message): + read_checkpoint_policy(_checkpoint_od_config(_fake_quant_config("modelopt"), policy=policy)) + + +@pytest.mark.parametrize( + ("step_range", "message"), + [ + ({"count": -1, "mode": "a16"}, "non-negative"), + ({"count": True, "mode": "a16"}, "non-negative"), + ({"count": 1, "mode": "native"}, "mode"), + ({"mode": "a16"}, "Missing"), + ({"count": 1, "mode": "a16", "sigma": 1.0}, "Unknown"), + ], +) +def test_checkpoint_policy_rejects_invalid_step_range(step_range: object, message: str) -> None: + policy = _policy() + policy["first_steps"] = step_range + with pytest.raises((TypeError, ValueError), match=message): + read_checkpoint_policy(_checkpoint_od_config(_fake_quant_config("modelopt"), policy=policy)) + + +def test_checkpoint_policy_rejects_unsupported_quantization_config() -> None: + with pytest.raises(ValueError, match="serialized ModelOpt FP8 or NVFP4"): + read_checkpoint_policy(_checkpoint_od_config(_fake_quant_config("compressed-tensors"), policy=_policy())) + + +def test_checkpoint_policy_for_other_scope_is_not_applied() -> None: + od_config = _checkpoint_od_config( + _fake_quant_config("modelopt"), + policy=_policy(scope=["vae"]), + ) + assert read_checkpoint_policy(od_config) is None + + +def test_additional_config_overrides_and_disables_checkpoint_policy() -> None: + checkpoint = _checkpoint_od_config( + _fake_quant_config("modelopt"), + policy=_policy(first_steps=1, last_steps=1), + ) + assert resolve_mixed_precision_config(checkpoint) == ( + Cosmos3MixedPrecisionConfig(first_steps=1, last_steps=1), + "checkpoint", + ) + + checkpoint.additional_config = { + "cosmos3_mixed_precision": { + "first_steps": 2, + "last_steps": 0, + "reasoner": "native", + } + } + assert resolve_mixed_precision_config(checkpoint) == ( + Cosmos3MixedPrecisionConfig(first_steps=2, last_steps=0, reasoner="native"), + "additional_config", + ) + + checkpoint.additional_config = {"cosmos3_mixed_precision": {"enabled": False}} + assert resolve_mixed_precision_config(checkpoint) == (None, "additional_config_disabled") + + +@pytest.mark.parametrize("name", ["modelopt", "modelopt_fp4"]) +def test_additional_config_supports_serialized_modelopt_formats_without_checkpoint_policy(name: str) -> None: + od_config = _checkpoint_od_config( + _fake_quant_config(name), + additional_config={"cosmos3_mixed_precision": {}}, + ) + assert resolve_mixed_precision_config(od_config) == (Cosmos3MixedPrecisionConfig(), "additional_config") + + +@pytest.mark.parametrize("has_checkpoint_policy", [False, True]) +@pytest.mark.parametrize( + ("name", "serialized"), + [ + ("modelopt_mixed", True), + ("mxfp8", True), + ("compressed-tensors", True), + ("modelopt", False), + ("modelopt_fp4", False), + ], +) +def test_additional_config_rejects_incompatible_quantization( + name: str, serialized: bool, has_checkpoint_policy: bool +) -> None: + od_config = _checkpoint_od_config( + _fake_quant_config(name, serialized=serialized), + policy=_policy() if has_checkpoint_policy else None, + additional_config={"cosmos3_mixed_precision": {}}, + ) + with pytest.raises(ValueError, match="serialized ModelOpt"): + resolve_mixed_precision_config(od_config) + + +@pytest.mark.parametrize("additional_config", [None, {"cosmos3_mixed_precision": {}}]) +def test_policy_validates_effective_quantization_config(additional_config: dict[str, object] | None) -> None: + od_config = _checkpoint_od_config( + _fake_quant_config("modelopt"), + policy=_policy(), + additional_config=additional_config, + ) + od_config.quantization_config = _fake_quant_config("modelopt_mixed") + with pytest.raises(ValueError, match="serialized ModelOpt FP8 or NVFP4"): + resolve_mixed_precision_config(od_config) + + +def test_additional_config_rejects_missing_quantization_config() -> None: + od_config = SimpleNamespace(additional_config={"cosmos3_mixed_precision": {}}) + with pytest.raises(ValueError, match="serialized ModelOpt FP8 or NVFP4"): + resolve_mixed_precision_config(od_config) + + +@pytest.mark.parametrize("policy", [_policy(), True]) +def test_additional_config_disable_bypasses_incompatible_checkpoint(policy: object) -> None: + od_config = _checkpoint_od_config( + _fake_quant_config("modelopt_mixed"), + policy=policy, + additional_config={"cosmos3_mixed_precision": {"enabled": False}}, + ) + od_config.quantization_config = _fake_quant_config("modelopt_mixed") + assert resolve_mixed_precision_config(od_config) == (None, "additional_config_disabled") + + +def test_runtime_allows_standard_offload_and_rejects_distributed(monkeypatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 + + monkeypatch.setattr( + transformer_cosmos3, + "get_tensor_model_parallel_world_size", + lambda: 1, + ) + od_config = SimpleNamespace( + enable_cpu_offload=True, + enable_layerwise_offload=True, + enable_distributed_layerwise_offload=False, + parallel_config=SimpleNamespace(use_hsdp=False), + ) + transformer_cosmos3._validate_mixed_precision_runtime( + Cosmos3MixedPrecisionConfig(), + od_config, + ) + od_config.enable_distributed_layerwise_offload = True + with pytest.raises(ValueError, match="distributed layer-wise offload"): + transformer_cosmos3._validate_mixed_precision_runtime( + Cosmos3MixedPrecisionConfig(), + od_config, + ) + od_config.enable_distributed_layerwise_offload = False + od_config.max_num_seqs = 2 + with pytest.raises(ValueError, match="one active request"): + transformer_cosmos3._validate_mixed_precision_runtime( + Cosmos3MixedPrecisionConfig(), + od_config, + ) + + +class _Layer(torch.nn.Module): + def __init__(self, output_size: int, input_size: int) -> None: + super().__init__() + self.output_size_per_partition = output_size + self.input_size_per_partition = input_size + + +class _BaseMethod: + def __init__(self, *, nvfp4: bool = False) -> None: + self.processed = False + self.apply_calls = 0 + self.fp8_linear = object() + if nvfp4: + from vllm.model_executor.kernels.linear.nvfp4.cutlass import ( + CutlassNvFp4LinearKernel, + ) + + self.kernel = object.__new__(CutlassNvFp4LinearKernel) + else: + self.kernel = None + + def create_weights(self, *args, **kwargs) -> None: + pass + + def process_weights_after_loading(self, layer) -> None: + self.processed = True + if layer.weight.dtype == torch.float8_e4m3fn: + layer.weight = layer.weight.t().contiguous() + layer.weight_scale = layer.weight_scale.max().reshape(1) + elif layer.weight.dtype == torch.uint8: + layer.weight_scale = _swizzle_blockscale_cpu(layer.weight_scale) + layer.weight_global_scale = layer.weight_scale_2.max().float().reshape(1) + del layer.weight_scale_2 + + def apply(self, layer, x, bias=None): + del bias + self.apply_calls += 1 + return torch.full( + (*x.shape[:-1], layer.output_size_per_partition), + 17, + dtype=x.dtype, + device=x.device, + ) + + +def _fp8_layer(output_size: int = 2, input_size: int = 4) -> _Layer: + layer = _Layer(output_size, input_size) + layer.weight = torch.ones(output_size, input_size, dtype=torch.float8_e4m3fn) + layer.weight_scale = torch.tensor([0.5], dtype=torch.float32) + return layer + + +def _nvfp4_layer(output_size: int = 128, input_size: int = 64) -> _Layer: + layer = _Layer(output_size, input_size) + layer.weight = torch.full((output_size, input_size // 2), 0x22, dtype=torch.uint8) + layer.weight_scale = torch.ones(output_size, input_size // 16, dtype=torch.float8_e4m3fn) + layer.weight_scale_2 = torch.ones(1) + return layer + + +def _runtime_and_method( + strategy, + layer: _Layer, + *, + path: str = "generation", + reasoner: str = "native", +): + config = Cosmos3MixedPrecisionConfig( + first_steps=1, + last_steps=1, + reasoner=reasoner, + ) + runtime = Cosmos3MixedPrecisionRuntime(config) + base = _BaseMethod(nvfp4=isinstance(strategy, Nvfp4W4A4W4A16Strategy)) + method = Cosmos3MixedPrecisionLinearMethod( + base, + strategy, + runtime, + f"{path}.linear", + path, + ) + method.process_weights_after_loading(layer) + return runtime, base, method + + +@pytest.mark.parametrize( + ("strategy", "layer"), + [ + (Fp8W8A8W8A16Strategy(), _fp8_layer()), + (Nvfp4W4A4W4A16Strategy(), _nvfp4_layer()), + ], +) +def test_generation_dispatches_native_middle_and_a16_edges(strategy, layer) -> None: + runtime, base, method = _runtime_and_method(strategy, layer) + x = torch.ones(1, layer.input_size_per_partition, dtype=torch.bfloat16) + + runtime.set_step(1, 3) + assert torch.equal( + method.apply(layer, x), + torch.full((1, layer.output_size_per_partition), 17, dtype=x.dtype), + ) + runtime.set_step(0, 3) + assert torch.equal( + method.apply(layer, x), + torch.nn.functional.linear(x, strategy.materialize(layer)), + ) + assert base.apply_calls == 1 + + +@pytest.mark.parametrize( + ("strategy", "layer"), + [ + (Fp8W8A8W8A16Strategy(), _fp8_layer()), + (Nvfp4W4A4W4A16Strategy(), _nvfp4_layer()), + ], +) +def test_a16_uses_only_live_native_weights(strategy, layer) -> None: + runtime, base, method = _runtime_and_method(strategy, layer) + assert base.processed + assert not hasattr(layer, "_cosmos3_precision_weight") + assert not hasattr(layer, "_cosmos3_precision_weight_scale") + + runtime.set_step(0, 3) + output = method.apply( + layer, + torch.ones(1, layer.input_size_per_partition, dtype=torch.bfloat16), + ) + assert output.shape == (1, layer.output_size_per_partition) + + +def test_a16_follows_live_weight_storage_after_offload_rebinding() -> None: + layer = _fp8_layer() + strategy = Fp8W8A8W8A16Strategy() + runtime, _, method = _runtime_and_method(strategy, layer) + runtime.set_step(0, 3) + x = torch.ones(1, 4, dtype=torch.bfloat16) + before = method.apply(layer, x) + + layer.weight = torch.zeros_like(layer.weight) + after = method.apply(layer, x) + assert not torch.equal(before, after) + assert torch.count_nonzero(after) == 0 + + +def test_fp8_rejects_smoothquant() -> None: + layer = _fp8_layer() + layer.pre_quant_scale = torch.ones(1) + with pytest.raises(ValueError, match="SmoothQuant"): + Fp8W8A8W8A16Strategy().validate_before_processing( + _BaseMethod(), + layer, + "gen.linear", + ) + + +def test_fp8_rejects_non_tensorwise_scales() -> None: + layer = _fp8_layer(output_size=3, input_size=2) + layer.weight_scale = torch.tensor([[1.0], [2.0], [3.0]]) + with pytest.raises(ValueError, match="tensorwise FP8"): + Fp8W8A8W8A16Strategy().validate_before_processing( + _BaseMethod(), + layer, + "gen.linear", + ) + + +def test_nvfp4_reference_materializes_live_cutlass_weights() -> None: + layer = _nvfp4_layer() + strategy = Nvfp4W4A4W4A16Strategy() + _runtime_and_method(strategy, layer) + output = strategy.materialize(layer) + assert output.shape == (128, 64) + assert torch.equal(output, torch.ones_like(output)) + + +def test_nvfp4_rejects_fused_global_scales() -> None: + layer = _nvfp4_layer() + layer.weight_scale_2 = torch.tensor([1.0, 2.0]) + strategy = Nvfp4W4A4W4A16Strategy() + with pytest.raises(ValueError, match="NVFP4 scales"): + strategy.validate_before_processing( + _BaseMethod(nvfp4=True), + layer, + "gen.qkv", + ) + + +@pytest.mark.parametrize( + ("policy", "uses_native"), + [("a16", False), ("native", True)], +) +def test_reasoner_policy_is_independent_of_generation_step( + policy: str, + uses_native: bool, +) -> None: + layer = _fp8_layer() + runtime, base, method = _runtime_and_method( + Fp8W8A8W8A16Strategy(), + layer, + path="reasoner", + reasoner=policy, + ) + runtime.set_step(1, 3) + method.apply(layer, torch.ones(1, 4, dtype=torch.bfloat16)) + assert (base.apply_calls == 1) is uses_native + + +def test_reset_clears_generation_state() -> None: + runtime = Cosmos3MixedPrecisionRuntime(Cosmos3MixedPrecisionConfig(first_steps=1, last_steps=1)) + runtime.set_step(0, 5) + assert runtime.use_high_precision("generation") + runtime.reset() + assert not runtime.use_high_precision("generation") + + +def test_install_discovers_generation_and_opt_in_reasoner(monkeypatch) -> None: + class _FakeLinear(torch.nn.Module): + def __init__(self, prefix: str) -> None: + super().__init__() + self.prefix = prefix + self.input_size_per_partition = 4 + self.output_size_per_partition = 2 + self.quant_method = SimpleNamespace(name="fp8") + + monkeypatch.setattr(runtime_impl, "LinearBase", _FakeLinear) + strategy = Fp8W8A8W8A16Strategy() + monkeypatch.setattr(strategy, "accepts", lambda method: getattr(method, "name", None) == "fp8") + monkeypatch.setattr(runtime_impl, "_STRATEGIES", (strategy,)) + runtime = Cosmos3MixedPrecisionRuntime(Cosmos3MixedPrecisionConfig(reasoner="a16")) + reasoner = _FakeLinear("reasoner.q_proj") + generation = [ + _FakeLinear("generation.q_proj"), + _FakeLinear("generation.out_proj"), + ] + transformer = SimpleNamespace( + language_model=SimpleNamespace(layers=torch.nn.Sequential(reasoner)), + gen_layers=torch.nn.Sequential(*generation), + ) + runtime.install(transformer) + assert isinstance(reasoner.quant_method, Cosmos3MixedPrecisionLinearMethod) + assert all(isinstance(layer.quant_method, Cosmos3MixedPrecisionLinearMethod) for layer in generation) + + +def test_pipeline_helpers_forward_and_reset_schedule() -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + Cosmos3OmniDiffusersPipeline, + ) + + calls = [] + pipeline = object.__new__(Cosmos3OmniDiffusersPipeline) + pipeline.transformer = SimpleNamespace( + set_mixed_precision_step=lambda step, count: calls.append((step, count)), + reset_mixed_precision=lambda: calls.append("reset"), + ) + pipeline._set_mixed_precision_step(2, 7) + pipeline._reset_mixed_precision() + assert calls == [(2, 7), "reset"] diff --git a/vllm_omni/diffusion/models/cosmos3/mixed_precision/__init__.py b/vllm_omni/diffusion/models/cosmos3/mixed_precision/__init__.py new file mode 100644 index 00000000000..1d5ee972baf --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/mixed_precision/__init__.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Scheduled A16 execution for quantized Cosmos3 checkpoints.""" + +from .checkpoint import resolve_mixed_precision_config +from .config import Cosmos3MixedPrecisionConfig +from .runtime import Cosmos3MixedPrecisionRuntime + +__all__ = [ + "Cosmos3MixedPrecisionConfig", + "Cosmos3MixedPrecisionRuntime", + "resolve_mixed_precision_config", +] diff --git a/vllm_omni/diffusion/models/cosmos3/mixed_precision/checkpoint.py b/vllm_omni/diffusion/models/cosmos3/mixed_precision/checkpoint.py new file mode 100644 index 00000000000..c8a62f142a0 --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/mixed_precision/checkpoint.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Checkpoint discovery for the Cosmos3 diffusion-step precision policy.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Literal, cast + +from .config import Cosmos3MixedPrecisionConfig, ReasonerPolicy, _non_negative_int + +PolicySource = Literal[ + "none", + "checkpoint", + "additional_config", + "additional_config_disabled", +] + +_COMPONENT = "transformer" +_POLICY_FIELDS = frozenset( + { + "schema_version", + "type", + "index_space", + "scope", + "default_mode", + "first_steps", + "last_steps", + "overlap", + "reasoner", + } +) +_STEP_RANGE_FIELDS = frozenset({"count", "mode"}) +_SUPPORTED_MODEL_OPT_CONFIGS = frozenset( + { + "modelopt", + "modelopt_fp4", + } +) + + +def resolve_mixed_precision_config( + od_config: object, +) -> tuple[Cosmos3MixedPrecisionConfig | None, PolicySource]: + """Resolve an explicit runtime override before the checkpoint default.""" + override_present, override = Cosmos3MixedPrecisionConfig.resolve_additional_config( + getattr(od_config, "additional_config", None) + ) + if override_present: + if override is not None: + _validate_checkpoint_quantization(od_config) + source: PolicySource = "additional_config" if override is not None else "additional_config_disabled" + return override, source + + checkpoint = read_checkpoint_policy(od_config) + if checkpoint is not None: + return checkpoint, "checkpoint" + return None, "none" + + +def read_checkpoint_policy(od_config: object) -> Cosmos3MixedPrecisionConfig | None: + """Read and validate a policy from ``transformer/config.json``. + + Missing metadata preserves ordinary checkpoint behavior. Metadata that is + present but malformed or incompatible fails closed. + """ + tf_config = getattr(od_config, "tf_model_config", None) + params = getattr(tf_config, "params", None) + if not isinstance(params, Mapping): + return None + + disk_quant_config = params.get("quantization_config") + if not isinstance(disk_quant_config, Mapping): + return None + + if "runtime" not in disk_quant_config: + return None + runtime_config = disk_quant_config["runtime"] + if not isinstance(runtime_config, Mapping): + raise TypeError("quantization_config.runtime must be a mapping") + + if "diffusion_step_policy" not in runtime_config: + return None + raw_policy = runtime_config["diffusion_step_policy"] + if not isinstance(raw_policy, Mapping): + raise TypeError("quantization_config.runtime.diffusion_step_policy must be a mapping") + + policy = _parse_policy(raw_policy) + if policy is None: + return None + _validate_checkpoint_quantization(od_config) + return policy + + +def _parse_policy(policy: Mapping[str, object]) -> Cosmos3MixedPrecisionConfig | None: + fields = set(policy) + unknown = fields - _POLICY_FIELDS + if unknown: + raise ValueError(f"Unknown diffusion_step_policy fields: {sorted(unknown)}") + missing = _POLICY_FIELDS - fields + if missing: + raise ValueError(f"Missing diffusion_step_policy fields: {sorted(missing)}") + + schema_version = policy["schema_version"] + if not isinstance(schema_version, int) or isinstance(schema_version, bool) or schema_version != 1: + raise ValueError("diffusion_step_policy.schema_version must be the integer 1") + if policy["type"] != "first_last_n": + raise ValueError("diffusion_step_policy.type must be 'first_last_n'") + if policy["index_space"] != "denoising_loop_iteration": + raise ValueError("diffusion_step_policy.index_space must be 'denoising_loop_iteration'") + if policy["default_mode"] != "native": + raise ValueError("diffusion_step_policy.default_mode must be 'native'") + if policy["overlap"] != "a16": + raise ValueError("diffusion_step_policy.overlap must be 'a16'") + + scope = policy["scope"] + if not isinstance(scope, list) or not scope or not all(isinstance(item, str) for item in scope): + raise TypeError("diffusion_step_policy.scope must be a non-empty list of strings") + + first_steps = _parse_step_range(policy["first_steps"], "first_steps") + last_steps = _parse_step_range(policy["last_steps"], "last_steps") + + reasoner = policy["reasoner"] + if reasoner not in {"native", "a16"}: + raise ValueError("diffusion_step_policy.reasoner must be 'native' or 'a16'") + + if _COMPONENT not in scope: + return None + return Cosmos3MixedPrecisionConfig( + first_steps=first_steps, + last_steps=last_steps, + reasoner=cast(ReasonerPolicy, reasoner), + ) + + +def _parse_step_range(value: object, name: str) -> int: + if not isinstance(value, Mapping): + raise TypeError(f"diffusion_step_policy.{name} must be a mapping") + fields = set(value) + unknown = fields - _STEP_RANGE_FIELDS + if unknown: + raise ValueError(f"Unknown diffusion_step_policy.{name} fields: {sorted(unknown)}") + missing = _STEP_RANGE_FIELDS - fields + if missing: + raise ValueError(f"Missing diffusion_step_policy.{name} fields: {sorted(missing)}") + if value["mode"] != "a16": + raise ValueError(f"diffusion_step_policy.{name}.mode must be 'a16'") + return _non_negative_int(value["count"], f"diffusion_step_policy.{name}.count") + + +def _validate_checkpoint_quantization(od_config: object) -> None: + # Match the effective quantization used to construct the transformer linears. + quant_config = getattr(od_config, "quantization_config", None) + if quant_config is None: + tf_config = getattr(od_config, "tf_model_config", None) + quant_config = getattr(tf_config, "quant_config", None) + + get_name = getattr(quant_config, "get_name", None) + name = get_name() if callable(get_name) else None + if name not in _SUPPORTED_MODEL_OPT_CONFIGS: + raise ValueError(f"diffusion_step_policy requires a serialized ModelOpt FP8 or NVFP4 checkpoint, got {name!r}") + + if name == "modelopt": + _require_serialized(quant_config, "is_checkpoint_fp8_serialized", "FP8") + return + if name == "modelopt_fp4": + if getattr(quant_config, "quant_method", "NVFP4") != "NVFP4": + raise ValueError("diffusion_step_policy requires checkpoint-native NVFP4 W4A4 linears") + _require_serialized(quant_config, "is_checkpoint_nvfp4_serialized", "NVFP4") + return + + +def _require_serialized(quant_config: object | None, flag: str, format_name: str) -> None: + if not bool(getattr(quant_config, flag, False)): + raise ValueError(f"diffusion_step_policy requires serialized ModelOpt {format_name} weights") diff --git a/vllm_omni/diffusion/models/cosmos3/mixed_precision/config.py b/vllm_omni/diffusion/models/cosmos3/mixed_precision/config.py new file mode 100644 index 00000000000..13748c8fd7b --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/mixed_precision/config.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Configuration parsing and denoising-step policy for mixed precision.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Literal, cast + +ReasonerPolicy = Literal["native", "a16"] + +_REASONER_POLICIES = frozenset({"native", "a16"}) +_CONFIG_FIELDS = frozenset({"enabled", "first_steps", "last_steps", "reasoner"}) + + +@dataclass(frozen=True) +class Cosmos3MixedPrecisionConfig: + """Validated precision policy shared by every quantization strategy.""" + + first_steps: int = 3 + last_steps: int = 3 + reasoner: ReasonerPolicy = "a16" + + @classmethod + def from_additional_config( + cls, + additional_config: Mapping[str, object] | None, + ) -> Cosmos3MixedPrecisionConfig | None: + """Parse Cosmos3 fields from vLLM-Omni's additional configuration.""" + _, config = cls.resolve_additional_config(additional_config) + return config + + @classmethod + def resolve_additional_config( + cls, + additional_config: Mapping[str, object] | None, + ) -> tuple[bool, Cosmos3MixedPrecisionConfig | None]: + """Return whether an override was supplied and its effective policy. + + The separate presence bit distinguishes an absent runtime override from + an explicit disable. That distinction matters when a checkpoint owns a + default policy. + """ + if additional_config is not None and not isinstance(additional_config, Mapping): + raise TypeError("additional_config must be a mapping") + values = additional_config or {} + if "cosmos3_mixed_precision" not in values: + return False, None + raw_config = values["cosmos3_mixed_precision"] + if not isinstance(raw_config, Mapping): + raise TypeError("cosmos3_mixed_precision must be a mapping") + unknown = set(raw_config) - _CONFIG_FIELDS + if unknown: + raise ValueError(f"Unknown cosmos3_mixed_precision fields: {sorted(unknown)}") + + enabled = raw_config.get("enabled", True) + if not isinstance(enabled, bool): + raise TypeError("cosmos3_mixed_precision.enabled must be a boolean") + if not enabled: + conflicting = set(raw_config) - {"enabled"} + if conflicting: + raise ValueError( + "cosmos3_mixed_precision.enabled=false cannot be combined with " + f"schedule fields: {sorted(conflicting)}" + ) + return True, None + + first_steps = _non_negative_int( + raw_config.get("first_steps", 3), + "cosmos3_mixed_precision.first_steps", + ) + last_steps = _non_negative_int( + raw_config.get("last_steps", 3), + "cosmos3_mixed_precision.last_steps", + ) + + reasoner = str(raw_config.get("reasoner", "a16")).lower() + if reasoner not in _REASONER_POLICIES: + raise ValueError( + f"cosmos3_mixed_precision.reasoner must be one of {sorted(_REASONER_POLICIES)}, got {reasoner!r}" + ) + config = cls( + first_steps=first_steps, + last_steps=last_steps, + reasoner=cast(ReasonerPolicy, reasoner), + ) + if first_steps == 0 and last_steps == 0 and reasoner == "native": + return True, None + return True, config + + def use_high_precision(self, step_index: int, num_steps: int) -> bool: + """Select the high path for an actual scheduler-step index.""" + if num_steps <= 0: + raise ValueError(f"num_steps must be positive, got {num_steps}") + if step_index < 0 or step_index >= num_steps: + raise IndexError(f"step_index must be in [0, {num_steps}), got {step_index}") + return step_index < self.first_steps or step_index >= num_steps - self.last_steps + + +def _non_negative_int(value: object, name: str) -> int: + """Validate an integer configuration field without accepting booleans.""" + if not isinstance(value, int) or isinstance(value, bool) or value < 0: + raise TypeError(f"{name} must be a non-negative integer") + return value diff --git a/vllm_omni/diffusion/models/cosmos3/mixed_precision/runtime.py b/vllm_omni/diffusion/models/cosmos3/mixed_precision/runtime.py new file mode 100644 index 00000000000..f528262e532 --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/mixed_precision/runtime.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Format-agnostic schedule state, linear discovery, and dispatch.""" + +from __future__ import annotations + +from typing import Literal + +import torch +from vllm.model_executor.layers.linear import LinearBase, LinearMethodBase + +from .config import Cosmos3MixedPrecisionConfig +from .strategy import ( + Cosmos3PrecisionStrategy, + Fp8W8A8W8A16Strategy, + Nvfp4W4A4W4A16Strategy, +) + +PrecisionPath = Literal["reasoner", "generation"] +_STRATEGIES: tuple[Cosmos3PrecisionStrategy, ...] = ( + Fp8W8A8W8A16Strategy(), + Nvfp4W4A4W4A16Strategy(), +) + + +def _strategy_for(method: object | None) -> Cosmos3PrecisionStrategy | None: + return next((strategy for strategy in _STRATEGIES if strategy.accepts(method)), None) + + +class Cosmos3MixedPrecisionLinearMethod(LinearMethodBase): + """Dispatch between a checkpoint-native method and dense A16.""" + + def __init__( + self, + base_method: LinearMethodBase, + strategy: Cosmos3PrecisionStrategy, + runtime: Cosmos3MixedPrecisionRuntime, + module_name: str, + path: PrecisionPath, + ) -> None: + self.base_method = base_method + self.strategy = strategy + self.runtime = runtime + self.module_name = module_name + self.path = path + + def create_weights(self, *args, **kwargs) -> None: + self.base_method.create_weights(*args, **kwargs) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + self.strategy.validate_before_processing( + self.base_method, + layer, + self.module_name, + ) + self.base_method.process_weights_after_loading(layer) + self.strategy.validate_after_processing(layer, self.module_name) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + if not self.runtime.use_high_precision(self.path): + return self.base_method.apply(layer, x, bias) + return self.strategy.apply_high(layer, x, bias) + + +class Cosmos3MixedPrecisionRuntime: + """Own one transformer's schedule state and wrapped linear inventory.""" + + def __init__(self, config: Cosmos3MixedPrecisionConfig) -> None: + self.config = config + # TODO: Make this request-local before Cosmos3 supports interleaved + # denoising requests on one transformer instance. + self._generation_high_precision = False + + def install(self, transformer: torch.nn.Module) -> None: + components: dict[PrecisionPath, torch.nn.Module] = { + "generation": transformer.gen_layers, + } + if self.config.reasoner == "a16": + components["reasoner"] = transformer.language_model.layers + + wrapped_count = 0 + for path, component in components.items(): + for local_name, layer in component.named_modules(): + if not isinstance(layer, LinearBase): + continue + base_method = getattr(layer, "quant_method", None) + strategy = _strategy_for(base_method) + if strategy is None: + continue + module_name = getattr(layer, "prefix", None) or f"{path}.{local_name}" + layer.quant_method = Cosmos3MixedPrecisionLinearMethod( + base_method, + strategy, + self, + module_name, + path, + ) + wrapped_count += 1 + + if not wrapped_count: + raise ValueError("Cosmos3 mixed precision found no compatible FP8 or NVFP4 ModelOpt linears") + + def use_high_precision(self, path: PrecisionPath) -> bool: + return self.config.reasoner == "a16" if path == "reasoner" else self._generation_high_precision + + def set_step(self, step_index: int, num_steps: int) -> None: + self._generation_high_precision = self.config.use_high_precision(step_index, num_steps) + + def reset(self) -> None: + self._generation_high_precision = False diff --git a/vllm_omni/diffusion/models/cosmos3/mixed_precision/strategy.py b/vllm_omni/diffusion/models/cosmos3/mixed_precision/strategy.py new file mode 100644 index 00000000000..8208b76222a --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/mixed_precision/strategy.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Format-specific live-weight validation and reference A16 materialization.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptFp8LinearMethod, + ModelOptNvFp4LinearMethod, +) +from vllm.model_executor.layers.quantization.utils import nvfp4_emulation_utils + +_NVFP4_BLOCK_SIZE = 16 + + +class Cosmos3PrecisionStrategy: + """Quantization-format boundary behind the shared precision schedule.""" + + def accepts(self, method: object | None) -> bool: + """Return whether this strategy owns a native linear method.""" + raise NotImplementedError + + def validate_before_processing( + self, + method: object, + layer: torch.nn.Module, + module_name: str, + ) -> None: + """Reject native transforms that cannot support live dequantization.""" + + def validate_after_processing( + self, + layer: torch.nn.Module, + module_name: str, + ) -> None: + """Validate the live native weight representation.""" + raise NotImplementedError + + def materialize(self, layer: torch.nn.Module) -> torch.Tensor: + """Materialize the live quantized weight as a dense BF16 matrix.""" + raise NotImplementedError + + def apply_high( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None, + ) -> torch.Tensor: + """Apply the reference dense A16 path.""" + input_size = int(layer.input_size_per_partition) + output_size = int(layer.output_size_per_partition) + weight = self.materialize(layer) + expected_shape = (output_size, input_size) + if tuple(weight.shape) != expected_shape: + raise RuntimeError(f"Dense weight shape {tuple(weight.shape)} does not match {expected_shape}") + x_2d = x.reshape(-1, x.shape[-1]) + if x_2d.shape[1] != input_size: + raise ValueError(f"Expected activation width {input_size}, got {x_2d.shape[1]}") + weight = weight.to(dtype=x.dtype) + if bias is not None and bias.dtype != x.dtype: + bias = bias.to(dtype=x.dtype) + output = F.linear(x_2d, weight, bias) + return output.view(*x.shape[:-1], output_size) + + +class Fp8W8A8W8A16Strategy(Cosmos3PrecisionStrategy): + """Use native ModelOpt W8A8 or reference dense W8A16.""" + + def accepts(self, method: object | None) -> bool: + return isinstance(method, ModelOptFp8LinearMethod) + + def validate_before_processing( + self, + method: object, + layer: torch.nn.Module, + module_name: str, + ) -> None: + if hasattr(layer, "pre_quant_scale"): + raise ValueError(f"{module_name} uses unsupported SmoothQuant pre_quant_scale") + weight_scale = getattr(layer, "weight_scale", None) + if not isinstance(weight_scale, torch.Tensor) or weight_scale.numel() != 1: + raise ValueError(f"{module_name} requires serialized tensorwise FP8 weights") + fp8_kernel = getattr(method, "fp8_linear", None) + from vllm.model_executor.kernels.linear.scaled_mm.marlin import ( + MarlinFP8ScaledMMLinearKernel, + ) + + if isinstance(fp8_kernel, MarlinFP8ScaledMMLinearKernel): + raise ValueError( + f"{module_name} selected Marlin FP8, which repacks the only weight copy; " + "use a backend that retains canonical FP8 weights" + ) + + def validate_after_processing( + self, + layer: torch.nn.Module, + module_name: str, + ) -> None: + weight = getattr(layer, "weight", None) + scale = getattr(layer, "weight_scale", None) + if not isinstance(weight, torch.Tensor) or weight.dtype != torch.float8_e4m3fn: + dtype = getattr(weight, "dtype", None) + raise TypeError(f"{module_name} requires a canonical FP8 weight, got {dtype}") + if weight.ndim != 2: + raise ValueError(f"{module_name} requires a rank-2 FP8 weight, got {tuple(weight.shape)}") + if not isinstance(scale, torch.Tensor): + raise ValueError(f"{module_name} has an FP8 weight but no weight_scale") + if scale.numel() != 1: + raise ValueError(f"{module_name} requires one live FP8 weight scale") + input_size = int(layer.input_size_per_partition) + output_size = int(layer.output_size_per_partition) + if weight.shape[0] < input_size or weight.shape[1] < output_size: + raise ValueError( + f"{module_name} live FP8 weight shape {tuple(weight.shape)} does not cover " + f"(K, N)=({input_size}, {output_size})" + ) + _validate_positive_finite_scale(scale, module_name) + + def materialize(self, layer: torch.nn.Module) -> torch.Tensor: + input_size = int(layer.input_size_per_partition) + output_size = int(layer.output_size_per_partition) + weight = layer.weight[:input_size, :output_size] + scale = layer.weight_scale.reshape(1).to(device=weight.device, dtype=torch.float32) + return (weight.to(torch.float32) * scale).t().to(torch.bfloat16) + + +class Nvfp4W4A4W4A16Strategy(Cosmos3PrecisionStrategy): + """Use native ModelOpt W4A4 or reference dense W4A16.""" + + def accepts(self, method: object | None) -> bool: + return isinstance(method, ModelOptNvFp4LinearMethod) + + def validate_before_processing( + self, + method: object, + layer: torch.nn.Module, + module_name: str, + ) -> None: + kernel = getattr(method, "kernel", None) + from vllm.model_executor.kernels.linear.nvfp4.cutlass import ( + CutlassNvFp4LinearKernel, + ) + from vllm.model_executor.kernels.linear.nvfp4.flashinfer import ( + FlashInferCuteDslNvFp4LinearKernel, + FlashInferCutlassNvFp4LinearKernel, + ) + + compatible_kernels = ( + CutlassNvFp4LinearKernel, + FlashInferCuteDslNvFp4LinearKernel, + FlashInferCutlassNvFp4LinearKernel, + ) + if not isinstance(kernel, compatible_kernels): + kernel_name = type(kernel).__name__ if kernel is not None else "none" + raise ValueError( + f"{module_name} selected unsupported NVFP4 backend {kernel_name}; " + "live W4A16 dequantization requires a CUTLASS-compatible layout" + ) + global_scale = _nvfp4_global_scale(layer) + if global_scale is None or global_scale.numel() != 1: + raise ValueError(f"{module_name} is missing ModelOpt NVFP4 scales") + + def validate_after_processing( + self, + layer: torch.nn.Module, + module_name: str, + ) -> None: + packed = getattr(layer, "weight", None) + scale = getattr(layer, "weight_scale", None) + global_scale = getattr(layer, "weight_global_scale", None) + if not isinstance(packed, torch.Tensor) or packed.dtype != torch.uint8: + dtype = getattr(packed, "dtype", None) + raise TypeError(f"{module_name} requires a live packed NVFP4 weight, got {dtype}") + if not isinstance(scale, torch.Tensor) or not isinstance(global_scale, torch.Tensor): + raise ValueError(f"{module_name} is missing live NVFP4 scales") + output_size = int(layer.output_size_per_partition) + input_size = int(layer.input_size_per_partition) + if input_size % _NVFP4_BLOCK_SIZE != 0: + raise ValueError(f"{module_name} input size {input_size} is not divisible by {_NVFP4_BLOCK_SIZE}") + expected_weight = (output_size, input_size // 2) + if packed.shape[0] < expected_weight[0] or packed.shape[1] < expected_weight[1]: + raise ValueError( + f"{module_name} live NVFP4 weight shape {tuple(packed.shape)} does not cover {expected_weight}" + ) + _validate_positive_finite_scale(global_scale, module_name) + lut_handle = nvfp4_emulation_utils.kE2M1ToFloat_handle + if lut_handle.val.device != packed.device: + lut_handle.val = lut_handle.val.to(packed.device) + + def materialize(self, layer: torch.nn.Module) -> torch.Tensor: + output_size = int(layer.output_size_per_partition) + input_size = int(layer.input_size_per_partition) + packed = layer.weight[:output_size, : input_size // 2] + return nvfp4_emulation_utils.dequantize_to_dtype( + packed, + layer.weight_scale, + layer.weight_global_scale, + torch.bfloat16, + _NVFP4_BLOCK_SIZE, + True, + ) + + +def _validate_positive_finite_scale(scale: torch.Tensor, module_name: str) -> None: + values = scale.detach().float() + if not torch.isfinite(values).all() or not (values > 0).all(): + raise ValueError(f"{module_name} has a non-finite or non-positive weight scale") + + +def _nvfp4_global_scale(layer: torch.nn.Module) -> torch.Tensor | None: + for name in ("weight_scale_2", "weight_global_scale"): + value = getattr(layer, name, None) + if isinstance(value, torch.Tensor): + return value + return None diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 49a4e40274f..c60817afb18 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -1822,6 +1822,16 @@ def do_classifier_free_guidance(self): def num_timesteps(self): return self._num_timesteps + def _set_mixed_precision_step(self, step_index: int, num_steps: int) -> None: + setter = getattr(self.transformer, "set_mixed_precision_step", None) + if setter is not None: + setter(step_index, num_steps) + + def _reset_mixed_precision(self) -> None: + resetter = getattr(self.transformer, "reset_mixed_precision", None) + if resetter is not None: + resetter() + @staticmethod def _distilled_unsupported_error(detail: str) -> ValueError: return ValueError( @@ -2819,7 +2829,8 @@ def _assign_step_out(step_out: torch.Tensor | tuple[torch.Tensor, ...]) -> None: # Each CFG-parallel rank runs exactly one branch (rank 0 -> cond, # else uncond), so session keying loads/stores only this rank's branch. cfg_rank_is_negative = get_classifier_free_guidance_rank() != 0 - for t in self.progress_bar(timesteps): + for step_index, t in enumerate(self.progress_bar(timesteps)): + self._set_mixed_precision_step(step_index, len(timesteps)) timestep = t.unsqueeze(0) # Outside the interval, scale=1 makes the combined output equal # the cond branch. Every rank remains on the same iteration and @@ -2858,7 +2869,8 @@ def _assign_step_out(step_out: torch.Tensor | tuple[torch.Tensor, ...]) -> None: uncond_cache: tuple = (None, None) keep_uncond_for_cache = self._cache_requires_paired_cfg() - for t in self.progress_bar(timesteps): + for step_index, t in enumerate(self.progress_bar(timesteps)): + self._set_mixed_precision_step(step_index, len(timesteps)) timestep = t.unsqueeze(0) cfg_active = _cfg_active_at(t) @@ -2912,7 +2924,8 @@ def _assign_step_out(step_out: torch.Tensor | tuple[torch.Tensor, ...]) -> None: else: # No CFG: a single cond branch per step. Bespoke (state None) keeps # using the transformer-instance cache exactly as before. - for t in self.progress_bar(timesteps): + for step_index, t in enumerate(self.progress_bar(timesteps)): + self._set_mixed_precision_step(step_index, len(timesteps)) timestep = t.unsqueeze(0) self._kv_load_und(kv_state, is_negative=False) noise_pred = self.transformer( @@ -2928,6 +2941,7 @@ def _assign_step_out(step_out: torch.Tensor | tuple[torch.Tensor, ...]) -> None: self._kv_capture_und(kv_state, is_negative=False) _assign_step_out(_step(noise_pred, t, latents, action_latents, sound_latents)) finally: + self._reset_mixed_precision() # Cosmos3 currently receives a unique request_id rather than a # reusable rollout session id. Retaining its state would only pin # K/V buffers on device after this generation finishes. @@ -3112,7 +3126,8 @@ def _active_at(t: torch.Tensor, interval: tuple[float, float] | None) -> bool: self.transformer.reset_cache() self._cosmos3_branch_caches = {} try: - for t in self.progress_bar(timesteps): + for step_index, t in enumerate(self.progress_bar(timesteps)): + self._set_mixed_precision_step(step_index, len(timesteps)) timestep = t.unsqueeze(0) step_guidance = guidance_scale if _active_at(t, guidance_interval) else 1.0 step_control = control_guidance if _active_at(t, control_guidance_interval) else 1.0 @@ -3224,6 +3239,7 @@ def _active_at(t: torch.Tensor, interval: tuple[float, float] | None) -> bool: )[0] latents = velocity_mask * latents + (1.0 - velocity_mask) * condition_latents finally: + self._reset_mixed_precision() self._cosmos3_branch_caches = None self.transformer.reset_cache() return latents diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index fe7a34c8b6e..1ded3db2f0d 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -40,6 +40,12 @@ from vllm_omni.diffusion.layers.norm import RMSNorm as _VllmRMSNorm from vllm_omni.platforms import current_omni_platform +from .mixed_precision import ( + Cosmos3MixedPrecisionConfig, + Cosmos3MixedPrecisionRuntime, + resolve_mixed_precision_config, +) + if TYPE_CHECKING: from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook @@ -136,6 +142,26 @@ def _od_config_get(od_config: Any, key: str, default: Any = None) -> Any: return default if value is None else value +def _validate_mixed_precision_runtime( + config: Cosmos3MixedPrecisionConfig | None, + od_config: OmniDiffusionConfig, +) -> None: + if config is None: + return + if get_tensor_model_parallel_world_size() != 1: + raise ValueError("Cosmos3 mixed precision currently supports tensor parallel size 1 only") + if int(getattr(od_config, "max_num_seqs", 1)) != 1: + raise ValueError("Cosmos3 mixed precision currently supports one active request per worker") + if bool(getattr(od_config, "enable_distributed_layerwise_offload", False)): + raise ValueError( + "Cosmos3 mixed precision does not support distributed layer-wise offload " + "because its direct loader bypasses ModelOpt post-load transformations" + ) + parallel_config = getattr(od_config, "parallel_config", None) + if bool(getattr(parallel_config, "use_hsdp", False)): + raise ValueError("Cosmos3 mixed precision has not validated live backend weights under HSDP") + + def _as_bool(value: Any) -> bool: if isinstance(value, str): return value.strip().lower() in {"1", "true", "yes", "on"} @@ -1241,6 +1267,19 @@ def __init__( dtype = od_config.dtype quant_config = getattr(od_config, "quantization_config", None) if od_config else None + 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": + logger.info("Cosmos3 checkpoint mixed-precision policy disabled by additional_config") + else: + logger.info( + "Cosmos3 mixed precision active (source=%s): first_steps=%d last_steps=%d reasoner=%s", + mixed_precision_source, + mixed_precision_config.first_steps, + mixed_precision_config.last_steps, + mixed_precision_config.reasoner, + ) + _validate_mixed_precision_runtime(mixed_precision_config, od_config) self.language_model = self._language_model_cls( hidden_size=self.hidden_size, @@ -1300,6 +1339,11 @@ def __init__( ] ) + self.mixed_precision_runtime: Cosmos3MixedPrecisionRuntime | None = None + if mixed_precision_config is not None: + self.mixed_precision_runtime = Cosmos3MixedPrecisionRuntime(mixed_precision_config) + self.mixed_precision_runtime.install(self) + self.norm_moe_gen = RMSNorm(self.hidden_size, eps=self.rms_norm_eps) self.gen_sp_prepare = Cosmos3GenSPPrepare() self.gen_sp_gather = nn.Identity() @@ -1958,3 +2002,11 @@ def forward( def post_load_weights(self) -> None: """Post-load processing: ensure correct dtypes.""" self.time_embedder.to(torch.float32) + + def set_mixed_precision_step(self, step_index: int, num_steps: int) -> None: + if self.mixed_precision_runtime is not None: + self.mixed_precision_runtime.set_step(step_index, num_steps) + + def reset_mixed_precision(self) -> None: + if self.mixed_precision_runtime is not None: + self.mixed_precision_runtime.reset() From e540dfb72f0439be08c26fe61912d8df572fdc7c Mon Sep 17 00:00:00 2001 From: tlysanhuo <166924864+tlysanhuo@users.noreply.github.com> Date: Thu, 10 Sep 2026 21:44:28 +0800 Subject: [PATCH 58/93] [Test] Use public render_jinja_template in MiniCPM-o native template test (#7362) Signed-off-by: tly <2200895168@qq.com> --- .../models/minicpmo_4_5/test_chat_template.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/tests/model_executor/models/minicpmo_4_5/test_chat_template.py b/tests/model_executor/models/minicpmo_4_5/test_chat_template.py index 8f1068392c2..2894b4a1c8f 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_chat_template.py +++ b/tests/model_executor/models/minicpmo_4_5/test_chat_template.py @@ -6,7 +6,7 @@ import pytest from jinja2 import TemplateError -from transformers.utils.chat_template_utils import _compile_jinja_template +from transformers.utils.chat_template_utils import render_jinja_template pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -26,12 +26,15 @@ ], ) def test_native_content_concatenation(parts, expected): - rendered = _compile_jinja_template(TEMPLATE.read_text()).render( - messages=[{"role": "user", "content": parts}], + # render_jinja_template takes a batch of conversations and returns + # (rendered_conversations, continuation_chunks), so [0][0] is the string. + rendered = render_jinja_template( + [[{"role": "user", "content": parts}]], + chat_template=TEMPLATE.read_text(), add_generation_prompt=True, enable_thinking=False, use_tts_template=True, - ) + )[0][0] assert rendered == ( f"<|im_start|>user\n{expected}<|im_end|>\n<|im_start|>assistant\n\n\n\n\n<|tts_bos|>" ) @@ -39,7 +42,8 @@ def test_native_content_concatenation(parts, expected): def test_unsupported_part_is_not_silently_discarded(): with pytest.raises(TemplateError, match="Unsupported MiniCPM-o content part"): - _compile_jinja_template(TEMPLATE.read_text()).render( - messages=[{"role": "user", "content": [{"type": "unknown"}]}], + render_jinja_template( + [[{"role": "user", "content": [{"type": "unknown"}]}]], + chat_template=TEMPLATE.read_text(), add_generation_prompt=True, ) From c8bef008ebad9913cc7990eb420df5cbd322d20e Mon Sep 17 00:00:00 2001 From: psv666 <150513104+psv666@users.noreply.github.com> Date: Thu, 10 Sep 2026 22:16:57 +0800 Subject: [PATCH 59/93] [Bugfix] Fix video prewarm cache retention and cancel-restart delay (#7363) Signed-off-by: psv666 <2693925048@qq.com> --- .../openai_api/test_serving_video_stream.py | 156 +++++++++++++++++- .../openai/serving_video_stream.py | 6 +- .../entrypoints/openai/video_stream_base.py | 38 +++-- 3 files changed, 182 insertions(+), 18 deletions(-) diff --git a/tests/entrypoints/openai_api/test_serving_video_stream.py b/tests/entrypoints/openai_api/test_serving_video_stream.py index bc80fe0dbcf..b534d6c6706 100644 --- a/tests/entrypoints/openai_api/test_serving_video_stream.py +++ b/tests/entrypoints/openai_api/test_serving_video_stream.py @@ -1,14 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for the serving-layer streaming video WebSocket handler.""" from __future__ import annotations import asyncio import base64 +import gc import io import json import threading +import weakref from typing import Any import pytest @@ -38,10 +40,10 @@ def _b64(data: bytes) -> str: def _text_result(text: str) -> OmniRequestOutput: class Output: - pass + text: str class RequestOutput: - pass + outputs: list[Output] output = Output() output.text = text @@ -52,10 +54,10 @@ class RequestOutput: def _audio_result(audio_data: Any) -> OmniRequestOutput: class Output: - pass + multimodal_output: dict[str, Any] class RequestOutput: - pass + outputs: list[Output] output = Output() output.multimodal_output = {"audio": audio_data} @@ -513,6 +515,88 @@ async def _process_query(self, *args, **kwargs): assert "session.done" in ws.sent_types() +@pytest.mark.asyncio +@pytest.mark.parametrize("delay_abort", [False, True], ids=["immediate-ack", "delayed-ack"]) +async def test_interrupted_queries_wait_for_abort_without_fixed_delay(monkeypatch, delay_abort): + started: asyncio.Queue[str] = asyncio.Queue() + abort_started: asyncio.Queue[str] = asyncio.Queue() + allow_abort = asyncio.Event() + if not delay_abort: + allow_abort.set() + request_ids: list[str] = [] + closed: list[str] = [] + aborted: list[str] = [] + delays: list[float] = [] + original_sleep = asyncio.sleep + + async def record_sleep(delay, result=None): + # Observe requested delays without a machine-dependent latency limit. + delays.append(delay) + return await original_sleep(0, result) + + monkeypatch.setattr(video_stream_base.asyncio, "sleep", record_sleep) + + class BlockingEngine: + async def generate(self, *, request_id, **kwargs): + if request_ids: + assert aborted[-1] == request_ids[-1] + request_ids.append(request_id) + started.put_nowait(request_id) + try: + if len(request_ids) <= 3: + yield _text_result("partial") + await asyncio.Event().wait() + else: + yield _text_result("final") + finally: + closed.append(request_id) + + async def abort(self, request_id): + assert request_id in closed + abort_started.put_nowait(request_id) + await allow_abort.wait() + aborted.append(request_id) + + class PreprocessedHandler(QwenOmniStreamingVideoHandler): + async def _preprocess_to_engine_prompt(self, request): + return {"prompt_token_ids": [1]} + + ws = TimedWebSocket() + handler = PreprocessedHandler(chat_service=object(), engine_client=BlockingEngine(), idle_timeout=5.0) + task = asyncio.create_task(handler.handle_session(ws)) + try: + ws.put({"type": "session.config", "modalities": ["text"], "enable_frame_filter": False}) + ws.put({"type": "video.frame", "data": _b64(_make_jpeg())}) + ws.put({"type": "video.query", "text": "first"}) + previous_id = await asyncio.wait_for(started.get(), timeout=2.0) + + for _ in range(3): + ws.put({"type": "video.query", "text": "interrupt"}) + assert await asyncio.wait_for(abort_started.get(), timeout=2.0) == previous_id + if delay_abort: + assert previous_id not in aborted + assert started.empty() + allow_abort.set() + previous_id = await asyncio.wait_for(started.get(), timeout=2.0) + if delay_abort: + allow_abort.clear() + + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + + assert aborted == request_ids[:-1] + assert len(set(request_ids)) == 4 + assert [msg["text"] for msg in ws.sent if msg["type"] == "response.text.done"] == ["final"] + assert "error" not in ws.sent_types() + assert "session.done" in ws.sent_types() + assert not [delay for delay in delays if delay > 0], "Restart must not add a timed grace period after abort" + finally: + allow_abort.set() + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + @pytest.mark.asyncio async def test_video_done_waits_for_in_flight_query(): query_started = asyncio.Event() @@ -588,6 +672,68 @@ async def _process_query(self, *args, **kwargs): assert "session.done" in ws.sent_types() +@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): + frame_a = _make_jpeg(255, 0, 0) + frame_b = _make_jpeg(0, 255, 0) + decode_started = asyncio.Event() + release_decode = asyncio.Event() + frame_b_accepted = asyncio.Event() + decoded_images: dict[bytes, weakref.ReferenceType[Image.Image]] = {} + prewarm_task: asyncio.Task | None = None + original_to_thread = asyncio.to_thread + + async def controlled_to_thread(function, *args, **kwargs): + nonlocal prewarm_task + if function is video_stream_base._decode_frame_bytes: + if args[0] == frame_a: + prewarm_task = asyncio.current_task() + decode_started.set() + await release_decode.wait() + image = await original_to_thread(function, *args, **kwargs) + decoded_images[args[0]] = weakref.ref(image) + return image + return await original_to_thread(function, *args, **kwargs) + + class AckWebSocket(TimedWebSocket): + async def send_json(self, data): + await super().send_json(data) + if data.get("type") == "video.frame.ack" and data.get("frame_id") == "B": + frame_b_accepted.set() + + monkeypatch.setattr(video_stream_base.asyncio, "to_thread", controlled_to_thread) + ws = AckWebSocket() + handler = QwenOmniStreamingVideoHandler(chat_service=object(), idle_timeout=5.0) + session_task = asyncio.create_task(handler.handle_session(ws)) + ws.put({"type": "session.config", "max_frames": max_frames, "enable_frame_filter": False}) + try: + ws.put({"type": "video.frame", "frame_id": "A", "data": _b64(frame_a)}) + await asyncio.wait_for(decode_started.wait(), timeout=5.0) + ws.put({"type": "video.frame", "frame_id": "B", "data": _b64(frame_b)}) + await asyncio.wait_for(frame_b_accepted.wait(), timeout=5.0) + ack = next(message for message in ws.sent if message.get("frame_id") == "B") + assert ack["accepted"] is True + assert ack.get("dropped_frame_id") == ("A" if max_frames == 1 else None) + + # Finish A's real decode only after B has either evicted A or joined it. + release_decode.set() + assert prewarm_task is not None + await asyncio.wait_for(asyncio.shield(prewarm_task), timeout=5.0) + gc.collect() + assert not session_task.done() + # A finished task must not keep an evicted PIL image alive for the session. + assert (decoded_images[frame_a]() is not None) == (max_frames == 2) + finally: + release_decode.set() + ws.put({"type": "video.done"}) + await asyncio.wait_for(session_task, timeout=5.0) + + gc.collect() + assert decoded_images[frame_a]() is None + assert not any(message.get("type") == "error" for message in ws.sent) + + @pytest.mark.asyncio async def test_client_cannot_send_internal_frame_decode_failed_message(): captured_frames: list[list[str]] = [] diff --git a/vllm_omni/entrypoints/openai/serving_video_stream.py b/vllm_omni/entrypoints/openai/serving_video_stream.py index af0083fcb70..e05a40eeffa 100644 --- a/vllm_omni/entrypoints/openai/serving_video_stream.py +++ b/vllm_omni/entrypoints/openai/serving_video_stream.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 """Qwen-Omni streaming video WebSocket handler. Accepts video frames incrementally via WebSocket, buffers them, and @@ -26,12 +26,14 @@ from __future__ import annotations +from collections.abc import Mapping from typing import Any from vllm_omni.entrypoints.openai.video_stream_base import ( _BAD_FRAME, _DEFAULT_CONFIG_TIMEOUT, _DEFAULT_IDLE_TIMEOUT, + PrewarmedFrame, StreamingVideoSessionConfig, VideoStreamTurnTrigger, ) @@ -59,7 +61,7 @@ def build_engine_prompt( audio_buffer: bytearray, message_history: list[dict[str, Any]], query_text: str, - prewarmed_frames: dict[str, tuple[Any, str]], + prewarmed_frames: Mapping[str, PrewarmedFrame], ) -> tuple[list[dict[str, Any]], dict[str, Any]]: n_buf = len(frame_buffer) if n_buf <= config.num_frames: diff --git a/vllm_omni/entrypoints/openai/video_stream_base.py b/vllm_omni/entrypoints/openai/video_stream_base.py index f975adbef65..ab5d51c33a5 100644 --- a/vllm_omni/entrypoints/openai/video_stream_base.py +++ b/vllm_omni/entrypoints/openai/video_stream_base.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 """Base WebSocket handler for streaming video input understanding. Shared session loop, frame/audio buffering, EVS pre-filter, prewarm, @@ -37,7 +37,8 @@ import wave from collections.abc import Mapping from dataclasses import dataclass -from typing import Any, Protocol, runtime_checkable +from enum import Enum +from typing import Any, Final, Protocol, TypeAlias, runtime_checkable import torch from fastapi import WebSocket, WebSocketDisconnect @@ -61,7 +62,14 @@ _MAX_AUDIO_BUFFER_BYTES = 4 * 1024 * 1024 _MAX_MSG_QUEUE = 200 _CODEC_FRAME_SAMPLES = 1920 # CausalConv leading-edge artifact length -_BAD_FRAME = object() + + +class _FrameStatus(Enum): + BAD = "bad" + + +_BAD_FRAME: Final = _FrameStatus.BAD +PrewarmedFrame: TypeAlias = tuple[Any, str] | _FrameStatus def _decode_frame_bytes(raw_bytes: bytes) -> Any: @@ -83,7 +91,7 @@ def build_engine_prompt( audio_buffer: bytearray, message_history: list[dict[str, Any]], query_text: str, - prewarmed_frames: dict[str, tuple[Any, str]], + prewarmed_frames: Mapping[str, PrewarmedFrame], ) -> tuple[list[dict[str, Any]], dict[str, Any]]: """Build OpenAI-style messages and the current user message.""" ... @@ -169,7 +177,7 @@ def build_engine_prompt( audio_buffer: bytearray, message_history: list[dict[str, Any]], query_text: str, - prewarmed_frames: dict[str, tuple[Any, str]], + prewarmed_frames: Mapping[str, PrewarmedFrame], ) -> tuple[list[dict[str, Any]], dict[str, Any]]: raise NotImplementedError @@ -219,7 +227,7 @@ async def handle_session(self, websocket: WebSocket) -> None: frame_buffer: list[str] = [] # base64-encoded JPEG frames frame_metadata: list[dict[str, Any]] = [] # Per-frame PIL cache + uuid for mm_hash reuse. Aligned with frame_buffer by index. - frame_pil_cache: dict[str, tuple[Any, str] | object] = {} # b64 -> (PIL.Image, uuid) or _BAD_FRAME + frame_pil_cache: dict[str, PrewarmedFrame] = {} # b64 -> (PIL.Image, uuid) or _BAD_FRAME frame_filter = ( FrameSimilarityFilter(threshold=config.frame_filter_threshold) if config.enable_frame_filter else None ) @@ -306,7 +314,6 @@ async def _start_query_turn(*, query_text: str) -> None: await self._engine_client.abort(prev_request_id) except Exception: pass - await asyncio.sleep(0.1) prev_was_interrupted = False request_id = f"video-{uuid.uuid4().hex[:12]}" @@ -429,8 +436,12 @@ async def _processor() -> None: async def _prewarm(b64: str, b: bytes, u: str) -> None: try: pil = await asyncio.to_thread(_decode_frame_bytes, b) - frame_pil_cache[b64] = (pil, u) + # The frame may have been evicted while decoding. + if b64 in frame_buffer: + frame_pil_cache[b64] = (pil, u) except Exception: + if b64 not in frame_buffer: + return frame_pil_cache[b64] = _BAD_FRAME logger.warning("prewarm decode failed for frame (len=%d)", len(b)) try: @@ -577,7 +588,7 @@ async def _process_query( query_text: str, request_id: str, interrupt_event: asyncio.Event, - prewarmed_frames: dict[str, tuple[Any, str]], + prewarmed_frames: Mapping[str, PrewarmedFrame], frame_metadata: list[dict[str, Any]] | None = None, ) -> None: """Build prompt, run inference, stream text + audio response.""" @@ -616,10 +627,13 @@ async def _process_query_engine( query_text: str, request_id: str, interrupt_event: asyncio.Event, - prewarmed_frames: dict[str, tuple[Any, str]], + prewarmed_frames: Mapping[str, PrewarmedFrame], frame_metadata: list[dict[str, Any]] | None = None, ) -> None: """Direct engine_client.generate() path for async_chunk audio.""" + engine_client = self._engine_client + assert engine_client is not None, "_process_query must validate the engine client" + from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ) @@ -686,7 +700,7 @@ async def _process_query_engine( audio_tail_tensors: list[Any] = [] try: - result_gen = self._engine_client.generate( + result_gen = engine_client.generate( prompt=engine_prompt, request_id=request_id, output_modalities=config.modalities, @@ -989,6 +1003,8 @@ def _encode_audio_wav_b64(audio_np) -> str: ) mixin = AudioMixin() resp = mixin.create_audio(audio_obj) + # base64_encode=True selects the string result of create_audio(). + assert isinstance(resp.audio_data, str) return resp.audio_data @staticmethod From 05ab5092695169ec711c2a0462e90e1b194968ad Mon Sep 17 00:00:00 2001 From: MaciejBalaNV Date: Thu, 10 Sep 2026 16:32:50 +0200 Subject: [PATCH 60/93] Cosmos3 action policy improvements (#6460) Signed-off-by: Maciej Bala Signed-off-by: MaciejBalaNV Co-authored-by: Hongsheng Liu --- recipes/cosmos3/Cosmos3-Nano.md | 96 ++++++++---- tests/config/test_config_factory.py | 71 +++++++++ .../cosmos3/test_cosmos3_framework_compat.py | 144 ++++++++++++++++++ .../models/cosmos3/test_cosmos3_pipeline.py | 78 ++++++++++ .../openai_api/test_openpi_connection.py | 90 +++++++++++ .../openai_api/test_openpi_serving.py | 78 ++++++++++ tests/entrypoints/test_serve.py | 20 +++ tests/model_extras/test_model_extras.py | 1 + vllm_omni/config/pipeline_registry.py | 6 + vllm_omni/config/stage_config.py | 16 +- vllm_omni/deploy/cosmos3_policy_droid.yaml | 49 ++++++ vllm_omni/diffusion/models/cosmos3/action.py | 11 ++ .../models/cosmos3/pipeline_cosmos3.py | 29 +++- vllm_omni/diffusion/models/cosmos3/utils.py | 116 +++++++++----- vllm_omni/entrypoints/cli/serve.py | 11 ++ vllm_omni/entrypoints/openai/api_server.py | 19 ++- vllm_omni/entrypoints/openpi/connection.py | 63 +++++--- vllm_omni/entrypoints/openpi/serving.py | 30 +++- .../model_executor/models/cosmos3/__init__.py | 2 + .../model_executor/models/cosmos3/pipeline.py | 41 +++++ vllm_omni/model_extras/cosmos3.py | 1 + 21 files changed, 865 insertions(+), 107 deletions(-) create mode 100644 tests/diffusion/models/cosmos3/test_cosmos3_framework_compat.py create mode 100644 vllm_omni/deploy/cosmos3_policy_droid.yaml create mode 100644 vllm_omni/model_executor/models/cosmos3/__init__.py create mode 100644 vllm_omni/model_executor/models/cosmos3/pipeline.py diff --git a/recipes/cosmos3/Cosmos3-Nano.md b/recipes/cosmos3/Cosmos3-Nano.md index f81e5e887f6..6ef24bcf49d 100644 --- a/recipes/cosmos3/Cosmos3-Nano.md +++ b/recipes/cosmos3/Cosmos3-Nano.md @@ -50,9 +50,10 @@ mode is selected per request: **`nvidia/Cosmos3-Nano-Policy-DROID`** is served the same way (`domain_name=droid_lerobot`). -- **DROID OpenPI policy server** — serve `nvidia/Cosmos3-Nano-Policy-DROID` and - connect an OpenPI-compatible websocket client to `/v1/realtime/robot/openpi`. - This path returns action chunks directly instead of an mp4. +- **DROID policy server** — serve `nvidia/Cosmos3-Nano-Policy-DROID` and connect + RoboLab or another OpenPI-compatible client to + `/v1/realtime/robot/openpi`. This endpoint returns action chunks directly + instead of an mp4. Action requests can use `input_reference` or `video_reference` for video input. `policy` and `forward_dynamics` can also use an image reference; `inverse_dynamics` @@ -314,40 +315,71 @@ VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \ curl -sS "http://localhost:8000/v1/videos/$VIDEO_ID" | jq '.action | {shape, dtype, raw_action_dim, domain_id}' curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_inverse_dynamics.mp4 -# DROID OpenPI policy server (websocket action serving). -# Requires cosmos_framework on PYTHONPATH because the pipeline reuses the -# reference RoboLab action transforms. If your checkpoint config already -# includes policy_server_config, omit the stage_overrides file and flag. -cat > cosmos3_droid_openpi_stage_overrides.json <<'JSON' -{ - "0": { - "model_config": { - "policy_server_config": { - "image_resolution": [540, 640], - "n_external_cameras": 2, - "needs_wrist_camera": true, - "needs_stereo_camera": false, - "needs_session_id": true, - "action_space": "joint_position" - } - } - } -} -JSON - +# DROID websocket policy server. Use the tested cosmos-framework revision +# directly from source; installing the package can introduce dependency conflicts. +export COSMOS_FRAMEWORK_ROOT=/path/to/cosmos-framework +git -C "$COSMOS_FRAMEWORK_ROOT" checkout c14617c2bc93dacbf69674fb964eec93182933d9 +export PYTHONPATH="$COSMOS_FRAMEWORK_ROOT" + +# Validate every cosmos-framework symbol used by the action-policy pipeline +# before allocating the model. +python - <<'PY' +from vllm_omni.diffusion.models.cosmos3.utils import ( + get_robolab_domain_id, + preflight_cosmos3_action_framework_imports, +) + +preflight_cosmos3_action_framework_imports() +print("cosmos-framework action imports OK; DROID domain:", get_robolab_domain_id("droid_lerobot")) +PY + +# The bundled deploy config (vllm_omni/deploy/cosmos3_policy_droid.yaml) +# selects the registered cosmos3_policy pipeline and carries the DROID +# checkpoint's serving defaults: JSON prompt formatting (a property of this +# checkpoint's training recipe, not a generic server default), the +# model-specific OpenPI binary handshake metadata (policy_server_config), and +# guardrails off (policy serving emits robot actions; the guardrail stack is +# not part of it, so --no-guardrails is implied). +# Policy checkpoints cannot be auto-detected — they share their HF metadata +# with the T2I/video Cosmos3 checkpoints — so --deploy-config is required. +# Per-checkpoint model_config tweaks can be layered on top with +# --stage-overrides; dict overrides deep-merge with the deploy yaml. +export VLLM_OMNI_ROOT=/absolute/path/to/vllm-omni vllm serve nvidia/Cosmos3-Nano-Policy-DROID \ --omni \ --host 0.0.0.0 --port 8000 \ - --model-class-name Cosmos3OmniDiffusersPipeline \ - --no-guardrails \ - --stage-overrides "$(cat cosmos3_droid_openpi_stage_overrides.json)" - -# Point an OpenPI websocket client at: -# ws://localhost:8000/v1/realtime/robot/openpi -# The first server message is policy_server_config. Each infer request sends a -# msgpack-numpy observation dict and receives a writable float32 action array. + --deploy-config "$VLLM_OMNI_ROOT/vllm_omni/deploy/cosmos3_policy_droid.yaml" \ + --robot-openpi-idle-timeout 0 + +# From a RoboLab checkout: +python policies/cosmos3/run.py \ + --remote-uri ws://localhost:8000/v1/realtime/robot/openpi \ + --task BananaInBowlTask + +# For a server started with "--api-key $POLICY_API_KEY", pass the same value as +# a Bearer token. --remote-token takes precedence over COSMOS3_API_TOKEN. +COSMOS3_API_TOKEN="$POLICY_API_KEY" python policies/cosmos3/run.py \ + --remote-uri wss://policy.example/v1/realtime/robot/openpi \ + --task BananaInBowlTask ``` +The endpoint sends `policy_server_config` as its initial binary MsgPack +message and returns the action array directly. RoboLab accepts that direct +array as well as proxy responses shaped as `{"action": ...}` or +`{"actions": ...}`, and turns structured server errors into client exceptions. + +RoboLab supplies `session_id=robolab-episode--env-`, so +parallel environments have independent policy state even though they share +one WebSocket. The server tracks interleaved session IDs independently. +The default receive-idle timeout is 30 seconds; set +`--robot-openpi-idle-timeout 0` when simulator steps between replans can take +longer, or set another non-negative timeout in seconds. + +If vLLM-Omni is started with `--api-key` or `VLLM_API_KEY`, the standard +OpenPI route uses the normal API authentication middleware. RoboLab's +`--remote-token` and `COSMOS3_API_TOKEN` send the required +`Authorization: Bearer ...` header. + #### Notes - **Measured latency (1x B300, bf16, guardrails off):** diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 51c7d549e28..af898ae501e 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -197,6 +197,25 @@ def test_to_omegaconf_max_num_seqs_in_engine_args(self): omega_config = config.to_omegaconf() assert omega_config.engine_args.max_num_seqs == 32 + def test_to_omegaconf_dict_override_deep_merges_yaml_dict(self): + """A partial dict CLI override (e.g. --no-guardrails riding on + model_config) must layer onto the YAML dict, not clobber siblings.""" + config = StageConfig( + stage_id=0, + model_stage="diffusion", + stage_type=StageType.DIFFUSION, + yaml_engine_args={ + "model_config": { + "guardrails": True, + "policy_server_config": {"action_space": "joint_position"}, + } + }, + runtime_overrides={"model_config": {"guardrails": False}}, + ) + omega_config = config.to_omegaconf() + assert omega_config.engine_args.model_config.guardrails is False + assert omega_config.engine_args.model_config.policy_server_config.action_space == "joint_position" + def test_to_omegaconf_diffusion_parallel_overrides_replace_nested_values(self): config = StageConfig( stage_id=1, @@ -623,6 +642,58 @@ def test_pipeline_config_supports_hf_architectures(self): assert p.hf_architectures == ("SomeCollidingArch",) +class TestCosmos3PolicyPipeline: + """Cosmos3 policy serving resolves via deploy yaml selection (π0 precedent).""" + + def test_registered_without_capturing_base_cosmos3_checkpoints(self): + assert "cosmos3_policy" in OMNI_PIPELINES + # T2I/video Cosmos3 checkpoints report model_type=cosmos3_omni and the + # same model_index.json _class_name as policy checkpoints; they must + # keep resolving through the single-stage diffusion fallback, so the + # policy pipeline must not be reachable by auto-detection. + assert "cosmos3_omni" not in OMNI_PIPELINES + pipeline = OMNI_PIPELINES["cosmos3_policy"] + assert pipeline.hf_architectures == () + assert pipeline.diffusers_class_name is None + + def test_droid_deploy_yaml_carries_policy_server_config(self): + deploy = load_deploy_config(get_deploy_config_path("cosmos3_policy_droid.yaml")) + assert deploy.pipeline == "cosmos3_policy" + + stages = merge_pipeline_deploy(OMNI_PIPELINES["cosmos3_policy"], deploy) + assert len(stages) == 1 + stage = stages[0].to_omegaconf() + + assert stage.stage_type == "diffusion" + assert stage.final_output_type == "action" + assert stage.engine_args.model_class_name == "Cosmos3OmniDiffusersPipeline" + # OpenPI websocket handshake metadata must reach the serving layer via + # engine_args.model_config (ServingRealtimeRobotOpenPI reads it there). + psc = stage.engine_args.model_config.policy_server_config + assert list(psc.image_resolution) == [540, 640] + assert psc.n_external_cameras == 2 + assert psc.needs_wrist_camera is True + assert psc.needs_stereo_camera is False + assert psc.needs_session_id is True + assert psc.action_space == "joint_position" + # The DROID checkpoint's training recipe expects JSON-formatted prompts. + assert stage.default_sampling_params.extra_args.format_prompt_as_json is True + + def test_no_guardrails_cli_flag_keeps_policy_server_config(self): + """--no-guardrails becomes a partial model_config CLI override; it must + deep-merge with the deploy yaml's model_config, not clobber the + policy_server_config the OpenPI handshake depends on.""" + deploy = load_deploy_config(get_deploy_config_path("cosmos3_policy_droid.yaml")) + stages, _ = StageConfigFactory._create_legacy_from_registry( + OMNI_PIPELINES["cosmos3_policy"], + {"model_config": {"guardrails": False}}, + user_deploy_config=deploy, + ) + stage = stages[0].to_omegaconf() + assert stage.engine_args.model_config.guardrails is False + assert stage.engine_args.model_config.policy_server_config.action_space == "joint_position" + + class TestStagePipelineConfig: def test_frozen(self): s = StagePipelineConfig(stage_id=0, model_stage="a") diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_framework_compat.py b/tests/diffusion/models/cosmos3/test_cosmos3_framework_compat.py new file mode 100644 index 00000000000..25f06b7a99e --- /dev/null +++ b/tests/diffusion/models/cosmos3/test_cosmos3_framework_compat.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from importlib.util import find_spec + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + +_FRAMEWORK_AVAILABLE = find_spec("cosmos_framework") is not None +requires_cosmos_framework = pytest.mark.skipif( + not _FRAMEWORK_AVAILABLE, + reason="cosmos_framework source is not present on PYTHONPATH", +) + + +def test_framework_preflight_checks_every_policy_symbol(monkeypatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import utils as cosmos3_utils + + imported = [] + + def record_import(module_name: str, symbol_name: str, error_message: str): + imported.append((module_name, symbol_name, error_message)) + return object() + + monkeypatch.setattr(cosmos3_utils, "lazy_import", record_import) + + cosmos3_utils.preflight_cosmos3_action_framework_imports() + + expected_symbols = { + "ActionTransformPipeline", + "FlowUniPCMultistepScheduler", + "build_abs_pose_from_components", + "convert_rotation", + "get_domain_id", + "pose_abs_to_rel", + "pose_rel_to_abs", + } + assert {symbol_name for _, symbol_name, _ in imported} == expected_symbols + assert [(module_name, symbol_name) for module_name, symbol_name, _ in imported] == list( + cosmos3_utils._ROBOLAB_FRAMEWORK_IMPORTS.values() + ) + + +def test_lazy_import_reports_framework_symbol_skew(monkeypatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import utils as cosmos3_utils + + monkeypatch.setattr(cosmos3_utils, "import_module", lambda _module_name: object()) + + with pytest.raises(ImportError, match="MissingSymbol") as exc_info: + cosmos3_utils.lazy_import( + "cosmos_framework.fake_module", + "MissingSymbol", + "Cosmos Framework compatibility error.", + ) + + assert isinstance(exc_info.value.__cause__, AttributeError) + + +@requires_cosmos_framework +def test_action_domain_table_matches_current_framework_and_legacy_aliases() -> None: + from cosmos_framework.data.generator.action.utils.domain_utils import ( + EMBODIMENT_TO_DOMAIN_ID as FRAMEWORK_DOMAIN_IDS, + ) + + from vllm_omni.diffusion.models.cosmos3.action import EMBODIMENT_TO_DOMAIN_ID as VLLM_DOMAIN_IDS + from vllm_omni.diffusion.models.cosmos3.utils import _LEGACY_EMBODIMENT_ALIASES + + canonical_vllm_domain_ids = { + name: domain_id for name, domain_id in VLLM_DOMAIN_IDS.items() if name not in _LEGACY_EMBODIMENT_ALIASES + } + + assert canonical_vllm_domain_ids == FRAMEWORK_DOMAIN_IDS + assert {alias: VLLM_DOMAIN_IDS[alias] for alias in _LEGACY_EMBODIMENT_ALIASES} == { + alias: FRAMEWORK_DOMAIN_IDS[canonical_name] for alias, canonical_name in _LEGACY_EMBODIMENT_ALIASES.items() + } + + +@requires_cosmos_framework +def test_current_cosmos_framework_policy_modules_import_and_instantiate() -> None: + from cosmos_framework.data.generator.action.utils.domain_utils import get_domain_id + from cosmos_framework.data.generator.action.utils.pose_utils import ( + build_abs_pose_from_components, + convert_rotation, + pose_abs_to_rel, + pose_rel_to_abs, + ) + from cosmos_framework.data.generator.action.utils.transforms import ActionTransformPipeline + from cosmos_framework.model.generator.diffusion.samplers.fm_solvers_unipc import ( + FlowUniPCMultistepScheduler, + ) + + transform = ActionTransformPipeline( + max_action_dim=64, + cfg_dropout_rate=0.0, + format_prompt_as_json=True, + ) + scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=1000, + shift=1.0, + use_dynamic_shifting=False, + ) + scheduler.set_timesteps(4, device=torch.device("cpu"), shift=5.0) + + assert transform.prompt_json_formatter is not None + assert len(scheduler.timesteps) == 4 + assert get_domain_id("droid_lerobot") == 8 + assert convert_rotation([0.0, 0.0, 0.0, 1.0], "quat_xyzw", "rot6d").shape == (6,) + assert callable(pose_abs_to_rel) + assert callable(pose_rel_to_abs) + assert callable(build_abs_pose_from_components) + + +@requires_cosmos_framework +def test_real_action_transform_pipeline_matches_droid_policy_contract() -> None: + from cosmos_framework.data.generator.action.utils.transforms import ActionTransformPipeline + + transform = ActionTransformPipeline( + max_action_dim=64, + cfg_dropout_rate=0.0, + format_prompt_as_json=True, + ) + sample = { + "ai_caption": "Pick up the cube.", + "video": torch.zeros((3, 3, 192, 320), dtype=torch.uint8), + "action": torch.zeros((3, 8), dtype=torch.float32), + "conditioning_fps": torch.tensor(15, dtype=torch.long), + "mode": "wam", + "domain_id": torch.tensor(8, dtype=torch.long), + "viewpoint": "concat_view", + "additional_view_description": "Wrist view above two exterior views.", + } + + transformed = transform(sample, "256") + + assert isinstance(transformed["ai_caption"], dict) + assert transformed["action"].shape == (3, 64) + assert int(transformed["raw_action_dim"].item()) == 8 + assert transformed["sequence_plan"].has_action is True + assert transformed["mode"] == "wam" + assert int(transformed["domain_id"].item()) == 8 diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index 72e4ab9622d..302409aebaa 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -3,6 +3,7 @@ from __future__ import annotations +import json import sys import types from dataclasses import dataclass @@ -21,6 +22,20 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] +@pytest.mark.parametrize( + ("alias", "canonical_name"), + [ + ("galbot", "embodiment_b"), + ("agibot_gear_gripper", "embodiment_c_gripper"), + ("agibot_gear_gripper_ext", "embodiment_c_gripper_ext"), + ], +) +def test_action_domain_table_preserves_legacy_aliases(alias: str, canonical_name: str) -> None: + from vllm_omni.diffusion.models.cosmos3.action import resolve_domain_id + + assert resolve_domain_id(domain_name=alias) == resolve_domain_id(domain_name=canonical_name) + + def test_pipeline_declares_layerwise_offload_components() -> None: from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline @@ -479,6 +494,69 @@ def fake_forward_robolab(sp, inputs, pipeline_start, session_id=None): assert captured["session_id"] == "robolab-request-7" +@pytest.mark.parametrize("format_prompt_as_json", [False, True]) +def test_robolab_input_builder_threads_prompt_format_and_uses_wam( + make_cosmos3_pipeline, + monkeypatch: pytest.MonkeyPatch, + format_prompt_as_json: bool, +) -> None: + from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3 + + pipeline = make_cosmos3_pipeline() + pipeline.transformer = StubCosmos3Transformer(action_gen=True, action_dim=64) + captured: dict[str, Any] = {} + + def fake_transform(sample, resolution): + captured["sample_mode"] = sample["mode"] + captured["resolution"] = resolution + sample["sequence_plan"] = SimpleNamespace( + condition_frame_indexes_action=[0], + action_start_frame_offset=1, + ) + sample["raw_action_dim"] = torch.tensor(8) + sample["image_size"] = torch.tensor([16, 16, 16, 16]) + if format_prompt_as_json: + sample["ai_caption"] = {"actions": {"instruction": sample["ai_caption"]}} + return sample + + def fake_get_transform(*, format_prompt_as_json: bool): + captured["format_prompt_as_json"] = format_prompt_as_json + return fake_transform + + pipeline._get_robolab_transform = fake_get_transform + monkeypatch.setattr(pipeline_cosmos3, "get_robolab_domain_id", lambda name: 8) + obs = { + "prompt": "Pick up the cube.", + "observation/image": np.zeros((16, 16, 3), dtype=np.uint8), + "observation/joint_position": np.zeros(7, dtype=np.float32), + "observation/gripper_position": np.zeros(1, dtype=np.float32), + } + sampling_params = make_sampling_params( + extra_args={ + "robot_obs": obs, + "action_chunk_size": 2, + "image_height": 16, + "image_width": 16, + "format_prompt_as_json": format_prompt_as_json, + } + ) + + inputs = pipeline._build_robolab_policy_inputs(sampling_params, request_id="request-1") + + assert inputs is not None + assert captured == { + "sample_mode": "wam", + "resolution": "480", + "format_prompt_as_json": format_prompt_as_json, + } + assert inputs.domain_id == 8 + assert inputs.raw_action_dim == 8 + if format_prompt_as_json: + assert json.loads(inputs.prompt) == {"actions": {"instruction": "Pick up the cube."}} + else: + assert inputs.prompt == "Pick up the cube." + + @pytest.mark.parametrize( ("prompt", "sampling_params", "message"), [ diff --git a/tests/entrypoints/openai_api/test_openpi_connection.py b/tests/entrypoints/openai_api/test_openpi_connection.py index 19ec23b6de2..a6f12f6612c 100644 --- a/tests/entrypoints/openai_api/test_openpi_connection.py +++ b/tests/entrypoints/openai_api/test_openpi_connection.py @@ -310,6 +310,29 @@ async def never_receives(): serving.infer.assert_not_called() +def test_handle_connection_can_disable_idle_timeout(monkeypatch): + monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj) + + async def fail_wait_for(*_args, **_kwargs): + pytest.fail("asyncio.wait_for should not be used when the idle timeout is disabled") + + monkeypatch.setattr(openpi_connection.asyncio, "wait_for", fail_wait_for) + websocket = FakeWebSocket([{"type": "websocket.disconnect"}]) + serving = _serving_mock() + + asyncio.run( + openpi_connection.RobotRealtimeConnection( + websocket, + serving, + idle_timeout=None, + ).handle_connection() + ) + + assert websocket.accepted is True + assert websocket.closed is False + serving.infer.assert_not_called() + + def test_handle_connection_keeps_session_state_per_websocket(monkeypatch): monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj) requests = { @@ -343,6 +366,73 @@ def test_handle_connection_keeps_session_state_per_websocket(monkeypatch): assert calls[2].kwargs == {"session_id": "session-b", "reset": True} +def test_handle_connection_keeps_interleaved_parallel_session_state(monkeypatch): + monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj) + requests = { + b"a1": {"prompt": "first-a", "session_id": "session-a"}, + b"b1": {"prompt": "first-b", "session_id": "session-b"}, + b"a2": {"prompt": "second-a", "session_id": "session-a"}, + b"b2": {"prompt": "second-b", "session_id": "session-b"}, + } + monkeypatch.setattr(openpi_connection, "_unpack", lambda data: dict(requests[data])) + serving = _serving_mock() + websocket = FakeWebSocket( + [{"type": "websocket.receive", "bytes": key} for key in (b"a1", b"b1", b"a2", b"b2")] + + [{"type": "websocket.disconnect"}] + ) + + asyncio.run(openpi_connection.RobotRealtimeConnection(websocket, serving).handle_connection()) + + assert [call.kwargs for call in serving.infer.await_args_list] == [ + {"session_id": "session-a", "reset": True}, + {"session_id": "session-b", "reset": True}, + {"session_id": "session-a", "reset": False}, + {"session_id": "session-b", "reset": False}, + ] + + +def test_handle_connection_bounds_tracked_sessions_without_resetting_live_ones(monkeypatch): + """A client may hold one socket across many episodes, minting a fresh + session id per episode and never sending ``reset`` (RoboLab does exactly + this). Session tracking must stay bounded, and evicting finished episodes + must not disturb the sessions that are still running. + """ + monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj) + monkeypatch.setattr(openpi_connection, "MAX_TRACKED_SESSIONS", 8) + + # One long-lived session interleaved with far more short episodes than the cap. + episodes = 40 + keys = [] + requests = {} + for episode in range(episodes): + for session_id in (f"episode-{episode}", "live-env"): + key = f"{session_id}@{episode}".encode() + keys.append(key) + requests[key] = {"prompt": "pick", "session_id": session_id} + + monkeypatch.setattr(openpi_connection, "_unpack", lambda data: dict(requests[data])) + serving = _serving_mock() + websocket = FakeWebSocket( + [{"type": "websocket.receive", "bytes": key} for key in keys] + [{"type": "websocket.disconnect"}] + ) + connection = openpi_connection.RobotRealtimeConnection(websocket, serving) + + asyncio.run(connection.handle_connection()) + + assert len(connection._seen_sessions) <= 8 + + resets_by_session: dict[str, list[bool]] = {} + for call in serving.infer.await_args_list: + resets_by_session.setdefault(call.kwargs["session_id"], []).append(call.kwargs["reset"]) + + # Each new episode resets exactly once, on its first observation. + for episode in range(episodes): + assert resets_by_session[f"episode-{episode}"] == [True] + # The continuously active session is never evicted, so it resets only once + # despite ``episodes`` other sessions churning past the cap. + assert resets_by_session["live-env"] == [True] + [False] * (episodes - 1) + + def test_handle_connection_reset_endpoint_resets_next_infer(monkeypatch): monkeypatch.setattr(openpi_connection, "_pack", lambda obj: obj) requests = { diff --git a/tests/entrypoints/openai_api/test_openpi_serving.py b/tests/entrypoints/openai_api/test_openpi_serving.py index 072e734e48f..82ee6e7dd21 100644 --- a/tests/entrypoints/openai_api/test_openpi_serving.py +++ b/tests/entrypoints/openai_api/test_openpi_serving.py @@ -15,6 +15,7 @@ from vllm_omni.entrypoints.openpi import connection as openpi_connection from vllm_omni.entrypoints.openpi import serving as openpi_serving +from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -177,6 +178,22 @@ def test_create_policy_server_returns_none_without_policy_config(): assert serving is None +def test_legacy_policy_request_defaults_key_no_longer_controls_serving(): + od_config = SimpleNamespace( + model_config={ + "policy_server_config": {}, + "policy_request_defaults": "malformed legacy value", + } + ) + + serving = openpi_serving.ServingRealtimeRobotOpenPI.create_policy_server( + engine_client=SimpleNamespace(get_diffusion_od_config=lambda: od_config), + model_name="generic-policy", + ) + + assert serving is not None + + def test_policy_server_config_allows_explicit_empty_config(): serving = openpi_serving.ServingRealtimeRobotOpenPI( engine_client=_engine_with_policy_config(policy_config={}), @@ -221,6 +238,44 @@ def test_build_request_uses_unique_engine_request_id_per_inference(): assert request_a.request_id != request_b.request_id +def test_build_request_clones_stage_defaults_before_protocol_fields(): + default_params = OmniDiffusionSamplingParams( + extra_args={ + "format_prompt_as_json": True, + "session_id": "configured-session-must-not-win", + "reset": False, + "nested": {"values": []}, + } + ) + od_config = SimpleNamespace(model_config={"policy_server_config": {}}) + engine_client = SimpleNamespace( + get_diffusion_od_config=lambda: od_config, + default_sampling_params_list=[default_params], + ) + serving = openpi_serving.ServingRealtimeRobotOpenPI( + engine_client=engine_client, + model_name="custom-policy", + ) + + request_a = serving._build_request( + {"prompt": "pick up the object"}, + session_id="wire-session", + reset=True, + ) + request_b = serving._build_request( + {"prompt": "pick up the object"}, + session_id="wire-session", + reset=False, + ) + + assert request_a.sampling_params.extra_args["format_prompt_as_json"] is True + assert request_a.sampling_params.extra_args["session_id"] == "wire-session" + assert request_a.sampling_params.extra_args["reset"] is True + request_a.sampling_params.extra_args["nested"]["values"].append("request-a") + assert request_b.sampling_params.extra_args["nested"] == {"values": []} + assert default_params.extra_args["nested"] == {"values": []} + + def test_build_request_forwards_seed_to_sampling_params(): """``seed`` in the inference message is the engine-level seed; omitted, the request auto-seeds.""" serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=_engine_with_policy_config()) @@ -257,6 +312,29 @@ async def run_requests(): assert sampling_params_b.extra_args["reset"] is False +def test_infer_yields_event_loop_while_engine_is_running(): + class DelayedEngine(RecordingEngine): + def generate(self, *, prompt, request_id, sampling_params_list): + async def _generate(): + await asyncio.sleep(0.03) + yield SimpleNamespace(multimodal_output={"actions": [0.0]}) + + return _generate() + + serving = openpi_serving.ServingRealtimeRobotOpenPI(engine_client=DelayedEngine()) + + async def run_test(): + infer_task = asyncio.create_task(serving.infer({"prompt": "pick"}, session_id="session-a", reset=True)) + ticks = 0 + while not infer_task.done(): + await asyncio.sleep(0.005) + ticks += 1 + await infer_task + return ticks + + assert asyncio.run(run_test()) >= 2 + + def test_two_websocket_clients_without_session_id_do_not_conflict(monkeypatch): monkeypatch.setattr(openpi_connection, "_pack", _json_pack) monkeypatch.setattr(openpi_connection, "_unpack", _json_unpack) diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 0089330b291..083f96d9f4e 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -112,6 +112,26 @@ def test_serve_parser_accepts_four_way_cfg_parallelism() -> None: assert args.cfg_parallel_size == 4 +def test_serve_parser_accepts_robot_openpi_idle_timeout() -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args(["serve", "fake-model", "--omni", "--robot-openpi-idle-timeout", "0"]) + + assert args.robot_openpi_idle_timeout == 0 + assert args.get_explicit_kwargs_dict()["robot_openpi_idle_timeout"] == 0 + + +def test_serve_parser_rejects_negative_robot_openpi_idle_timeout() -> None: + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + + with pytest.raises(SystemExit): + parser.parse_args(["serve", "fake-model", "--omni", "--robot-openpi-idle-timeout", "-1"]) + + def test_serve_parser_accepts_ulysses_a2a_permute() -> None: parser = TrackingArgumentParser() subparsers = parser.add_subparsers(dest="subcommand") diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index 6bf82333922..e55691e104d 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -177,6 +177,7 @@ def test_cosmos3_extra_registry_declares_request_and_response_params(pipeline_na "resolution", "image_size", "use_state", + "format_prompt_as_json", "observation", "robot_obs", "deterministic_seed", diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index b711f0cfb36..a69fb08c65b 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -46,6 +46,7 @@ BAGEL_SINGLE_STAGE_PIPELINE, BAGEL_THINK_PIPELINE, ) +from vllm_omni.model_executor.models.cosmos3.pipeline import COSMOS3_POLICY_PIPELINE from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE from vllm_omni.model_executor.models.dots_tts.pipeline import DOTS_TTS_PIPELINE @@ -149,6 +150,11 @@ "lingbot_world": LINGBOT_WORLD_PIPELINE, "Gr00tN1d7": GR00T_N1D7_PIPELINE, "pi0": PI0_PIPELINE, + # Cosmos3 policy checkpoints share HF metadata with the T2I/video Cosmos3 + # checkpoints (which stay on the single-stage diffusion fallback), so this + # entry is only reachable through a deploy yaml's ``pipeline:`` key + # (see deploy/cosmos3_policy_droid.yaml). + "cosmos3_policy": COSMOS3_POLICY_PIPELINE, "gepard": GEPARD_PIPELINE, "glm_image": GLM_IMAGE_PIPELINE, "hunyuan_image_3_moe": HUNYUAN_IMAGE3_PIPELINE, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index f7d71489075..0e962b4cc30 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -1130,10 +1130,22 @@ def to_omegaconf(self) -> Any: _apply_diffusion_parallel_runtime_overrides(engine_args, runtime_overrides) reconcile_diffusion_attention_overrides(engine_args, runtime_overrides) - # CLI overrides take precedence over YAML defaults + # CLI overrides take precedence over YAML defaults. Most dict-valued + # overrides are deep-merged so a partial CLI dict (e.g. --no-guardrails + # riding on ``model_config``) layers onto the deploy YAML instead of + # clobbering sibling keys such as ``policy_server_config`` — the same + # rationale as the platform-overlay deep-merge. Legacy atomic mappings + # are handled explicitly below. for key, value in runtime_overrides.items(): if value is not None and key not in ("devices", "max_batch_size", "num_replicas"): - engine_args[key] = value + existing = engine_args.get(key) + # ``omni_kv_config`` is an atomic legacy override: callers use + # a partial mapping to replace the topology-provided transfer + # role, rather than to add fields to it. + if key != "omni_kv_config" and isinstance(existing, dict) and isinstance(value, dict): + engine_args[key] = _get_recursively_merged_dict(existing, value) + else: + engine_args[key] = value # Build runtime config from YAML defaults + CLI overrides runtime: dict[str, Any] = dict(self.yaml_runtime) diff --git a/vllm_omni/deploy/cosmos3_policy_droid.yaml b/vllm_omni/deploy/cosmos3_policy_droid.yaml new file mode 100644 index 00000000000..cbcbcb42c01 --- /dev/null +++ b/vllm_omni/deploy/cosmos3_policy_droid.yaml @@ -0,0 +1,49 @@ +# Cosmos3 policy (DROID) deploy: single diffusion stage serving the OpenPI +# realtime API (/v1/realtime/robot/openpi). +# +# Topology (COSMOS3_POLICY_PIPELINE) is declared in +# vllm_omni/model_executor/models/cosmos3/pipeline.py. Policy checkpoints share +# their HF metadata (model_type=cosmos3_omni, model_index.json _class_name= +# Cosmos3OmniDiffusersPipeline) with the T2I/video Cosmos3 checkpoints, which +# must keep using the default single-stage diffusion fallback — so this +# pipeline cannot be auto-detected and must be selected explicitly: +# +# vllm serve nvidia/Cosmos3-Nano-Policy-DROID --omni \ +# --deploy-config /absolute/path/to/vllm-omni/vllm_omni/deploy/cosmos3_policy_droid.yaml +# +# Per-checkpoint model_config tweaks (e.g. policy_server_config values) can +# still be layered on top with --stage-overrides; dict overrides deep-merge +# with this file. + +pipeline: cosmos3_policy +# async_chunk defaults to True, which the deploy validator rejects for a +# single-stage pipeline (no next-stage processor). Must be false here. +async_chunk: false +distributed_executor_backend: mp + +stages: + - stage_id: 0 + devices: "0" + # OpenPI policy serving is one observation -> one action reply. + max_num_seqs: 1 + model_class_name: Cosmos3OmniDiffusersPipeline + # JSON prompt formatting is a property of this DROID checkpoint's training + # recipe, not a generic server default. + default_sampling_params: + extra_args: + format_prompt_as_json: true + model_config: + # Policy serving emits robot actions, not text/video; the Cosmos3 + # guardrail stack (gated nvidia/Cosmos-1.0-Guardrail repo) is not part + # of it. Equivalent to launching with --no-guardrails. + guardrails: false + # OpenPI handshake metadata advertised to the robot client. Sent verbatim + # (the pipeline never reads these); they describe the camera inputs and + # action space the client must produce for this checkpoint. + policy_server_config: + image_resolution: [540, 640] + n_external_cameras: 2 + needs_wrist_camera: true + needs_stereo_camera: false + needs_session_id: true + action_space: joint_position diff --git a/vllm_omni/diffusion/models/cosmos3/action.py b/vllm_omni/diffusion/models/cosmos3/action.py index dc755b39c39..affba419517 100644 --- a/vllm_omni/diffusion/models/cosmos3/action.py +++ b/vllm_omni/diffusion/models/cosmos3/action.py @@ -38,13 +38,24 @@ "bridge_orig_lerobot": 7, "droid_lerobot": 8, "robomind-franka": 8, + "embodiment_b": 9, + # Backward-compatible alias for the pre-generator Cosmos Framework name. "galbot": 9, "robomind-franka-dual": 12, "robomind-ur": 13, "agibotworld": 15, + "embodiment_c_gripper": 15, + "embodiment_c_gripper_ext": 15, + # Backward-compatible aliases for the pre-generator Cosmos Framework names. "agibot_gear_gripper": 15, "agibot_gear_gripper_ext": 15, + "xdof_yam": 16, + "molmoact2_yam": 16, + "abc_yam": 16, "fractal": 20, + "drawanything": 21, + "behavior1k_lerobot": 22, + "maniparena": 23, } diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index c60817afb18..241eae80577 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -138,6 +138,7 @@ ensure_gripper_array, extract_robolab_image, extract_robolab_prompt_image, + get_robolab_domain_id, lazy_action_transform_pipeline, make_robolab_action_postprocess_inputs, next_robolab_seed, @@ -1075,7 +1076,7 @@ def __init__( self._guidance_scale = None self._num_timesteps = None self._cosmos3_branch_caches: dict[str, tuple[Any, Any]] | None = None - self._robolab_transform = None + self._robolab_transforms: dict[bool, Any] = {} # Set True by ``enable_cache_for_cosmos3`` when cache-dit is enabled on # this pipeline. Tells the sequential-CFG loop to keep paired @@ -1394,11 +1395,18 @@ def _get_sp_param(sp: OmniDiffusionSamplingParams, key: str, default: Any = None return val return default - def _get_robolab_transform(self): - if self._robolab_transform is None: + def _get_robolab_transform(self, *, format_prompt_as_json: bool = False): + transforms = getattr(self, "_robolab_transforms", None) + if transforms is None: + transforms = {} + self._robolab_transforms = transforms + if format_prompt_as_json not in transforms: action_dim = int(getattr(self.transformer, "action_dim", 64)) - self._robolab_transform = lazy_action_transform_pipeline(action_dim) - return self._robolab_transform + transforms[format_prompt_as_json] = lazy_action_transform_pipeline( + action_dim, + format_prompt_as_json=format_prompt_as_json, + ) + return transforms[format_prompt_as_json] def _build_robolab_policy_inputs( self, @@ -1443,7 +1451,8 @@ def extra_param_alias(primary_key: str, alias_key: str, default: Any) -> Any: resolution = str(extra_param("resolution", ROBOLAB_DEFAULT_RESOLUTION)) fps = float(extra_param("conditioning_fps", ROBOLAB_DEFAULT_CONDITIONING_FPS)) domain_name = str(extra_param("domain_name", ROBOLAB_DEFAULT_DOMAIN_NAME)) - domain_id = resolve_domain_id(domain_name=domain_name, require_explicit=True) + domain_id = get_robolab_domain_id(domain_name) + format_prompt_as_json = self._truthy(extra_param("format_prompt_as_json", False)) if use_state and history_length < 1: raise ValueError("RoboLab history_length must be >= 1 when use_state is true.") @@ -1506,7 +1515,9 @@ def extra_param_alias(primary_key: str, alias_key: str, default: Any) -> Any: "action": action, # Cosmos Framework consumes this as an integer conditioning bucket. "conditioning_fps": torch.tensor(fps, dtype=torch.long), - "mode": ACTION_MODE_POLICY, + # Cosmos Framework 1.2 renamed the internal policy transform mode to + # ``wam``. The public vLLM action_mode remains ``policy``. + "mode": "wam", "domain_id": torch.tensor(domain_id, dtype=torch.long), "viewpoint": "concat_view", "additional_view_description": ROBOLAB_CONCAT_VIEW_DESCRIPTION, @@ -1514,7 +1525,9 @@ def extra_param_alias(primary_key: str, alias_key: str, default: Any) -> Any: if history_action is not None: sample["history_action"] = history_action - sample = self._get_robolab_transform()(sample, resolution) + sample = self._get_robolab_transform(format_prompt_as_json=format_prompt_as_json)(sample, resolution) + if isinstance(sample.get("ai_caption"), dict): + sample["ai_caption"] = json.dumps(sample["ai_caption"]) sequence_plan = sample["sequence_plan"] video_tensor = sample["video"].float() / 127.5 - 1.0 raw_action_dim_tensor = sample.get("raw_action_dim") diff --git a/vllm_omni/diffusion/models/cosmos3/utils.py b/vllm_omni/diffusion/models/cosmos3/utils.py index 76df64b3b26..8d104e98792 100644 --- a/vllm_omni/diffusion/models/cosmos3/utils.py +++ b/vllm_omni/diffusion/models/cosmos3/utils.py @@ -78,6 +78,12 @@ "sides, with the robot visible." ) +_LEGACY_EMBODIMENT_ALIASES = { + "galbot": "embodiment_b", + "agibot_gear_gripper": "embodiment_c_gripper", + "agibot_gear_gripper_ext": "embodiment_c_gripper_ext", +} + @dataclass(frozen=True) class RoboLabPolicyInputs: @@ -252,33 +258,83 @@ def extract_robolab_prompt_image(prompt_data: Any | None) -> np.ndarray | None: return ensure_rgb_uint8_image(image, "multi_modal_data.image") +_ROBOLAB_FRAMEWORK_IMPORT_ERROR = ( + "Cosmos3 RoboLab policy serving requires a compatible cosmos_framework source on PYTHONPATH." +) +_ROBOLAB_FRAMEWORK_IMPORTS: dict[str, tuple[str, str]] = { + "ActionTransformPipeline": ( + "cosmos_framework.data.generator.action.utils.transforms", + "ActionTransformPipeline", + ), + "convert_rotation": ( + "cosmos_framework.data.generator.action.utils.pose_utils", + "convert_rotation", + ), + "pose_abs_to_rel": ( + "cosmos_framework.data.generator.action.utils.pose_utils", + "pose_abs_to_rel", + ), + "pose_rel_to_abs": ( + "cosmos_framework.data.generator.action.utils.pose_utils", + "pose_rel_to_abs", + ), + "build_abs_pose_from_components": ( + "cosmos_framework.data.generator.action.utils.pose_utils", + "build_abs_pose_from_components", + ), + "FlowUniPCMultistepScheduler": ( + "cosmos_framework.model.generator.diffusion.samplers.fm_solvers_unipc", + "FlowUniPCMultistepScheduler", + ), + "get_domain_id": ( + "cosmos_framework.data.generator.action.utils.domain_utils", + "get_domain_id", + ), +} + + def lazy_import(module_name: str, symbol_name: str, error_message: str): try: module = import_module(module_name) - except ModuleNotFoundError as exc: - raise ModuleNotFoundError(error_message) from exc - return getattr(module, symbol_name) + return getattr(module, symbol_name) + except (ImportError, AttributeError) as exc: + raise ImportError( + f"{error_message} Could not import required symbol {symbol_name!r} from {module_name!r}." + ) from exc + + +def _load_robolab_framework_symbol(symbol_name: str): + module_name, exported_name = _ROBOLAB_FRAMEWORK_IMPORTS[symbol_name] + return lazy_import( + module_name, + exported_name, + _ROBOLAB_FRAMEWORK_IMPORT_ERROR, + ) -def lazy_action_transform_pipeline(max_action_dim: int): - ActionTransformPipeline = lazy_import( - "cosmos_framework.data.vfm.action.transforms", - "ActionTransformPipeline", - "Cosmos3 RoboLab policy serving requires cosmos_framework on PYTHONPATH so the " - "golden ActionTransformPipeline can be reused.", +def preflight_cosmos3_action_framework_imports() -> None: + """Validate every Cosmos Framework symbol used by action policy serving.""" + for symbol_name in _ROBOLAB_FRAMEWORK_IMPORTS: + _load_robolab_framework_symbol(symbol_name) + + +def lazy_action_transform_pipeline(max_action_dim: int, *, format_prompt_as_json: bool = False): + ActionTransformPipeline = _load_robolab_framework_symbol("ActionTransformPipeline") + return ActionTransformPipeline( + max_action_dim=max_action_dim, + cfg_dropout_rate=0.0, + format_prompt_as_json=format_prompt_as_json, ) - return ActionTransformPipeline(max_action_dim=max_action_dim, cfg_dropout_rate=0.0) + + +def get_robolab_domain_id(domain_name: str) -> int: + get_domain_id = _load_robolab_framework_symbol("get_domain_id") + canonical_name = _LEGACY_EMBODIMENT_ALIASES.get(domain_name.lower().strip(), domain_name) + return int(get_domain_id(canonical_name)) def build_robolab_unipc_scheduler(num_steps: int, shift: float, device: torch.device): - FlowUniPCMultistepScheduler = lazy_import( - "cosmos_framework.model.vfm.diffusion.samplers.fm_solvers_unipc", - "FlowUniPCMultistepScheduler", - ( - "Cosmos3 RoboLab policy serving requires cosmos_framework on PYTHONPATH so the " - "golden FlowUniPCMultistepScheduler can be reused." - ), - ) + FlowUniPCMultistepScheduler = _load_robolab_framework_symbol("FlowUniPCMultistepScheduler") scheduler = FlowUniPCMultistepScheduler( num_train_timesteps=1000, @@ -290,38 +346,22 @@ def build_robolab_unipc_scheduler(num_steps: int, shift: float, device: torch.de def convert_midtrain_rotation(value: Any, src: str, dst: str) -> np.ndarray: - convert_rotation = lazy_import( - "cosmos_framework.data.vfm.action.pose_utils", - "convert_rotation", - "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.", - ) + convert_rotation = _load_robolab_framework_symbol("convert_rotation") return convert_rotation(value, src, dst) def pose_abs_to_rel(*args, **kwargs) -> np.ndarray: - pose_abs_to_rel_func = lazy_import( - "cosmos_framework.data.vfm.action.pose_utils", - "pose_abs_to_rel", - "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.", - ) + pose_abs_to_rel_func = _load_robolab_framework_symbol("pose_abs_to_rel") return pose_abs_to_rel_func(*args, **kwargs) def pose_rel_to_abs(*args, **kwargs) -> np.ndarray: - pose_rel_to_abs_func = lazy_import( - "cosmos_framework.data.vfm.action.pose_utils", - "pose_rel_to_abs", - "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.", - ) + pose_rel_to_abs_func = _load_robolab_framework_symbol("pose_rel_to_abs") return pose_rel_to_abs_func(*args, **kwargs) def build_abs_pose_from_components(*args, **kwargs) -> np.ndarray: - build_abs_pose_from_components_func = lazy_import( - "cosmos_framework.data.vfm.action.pose_utils", - "build_abs_pose_from_components", - "Cosmos3 RoboLab midtrain action serving requires cosmos_framework pose_utils on PYTHONPATH.", - ) + build_abs_pose_from_components_func = _load_robolab_framework_symbol("build_abs_pose_from_components") return build_abs_pose_from_components_func(*args, **kwargs) diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 348e4131643..598cbe89328 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -930,6 +930,15 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu action="store_true", help="Disable Cosmos3 text/video safety guardrails for this server.", ) + omni_config_group.add_argument( + "--robot-openpi-idle-timeout", + type=_nonneg_finite_float, + default=30.0, + help=( + "Seconds the /v1/realtime/robot/openpi endpoint waits for the next request " + "before closing an idle WebSocket (default: 30). Set to 0 to disable the timeout." + ), + ) # Enable diffusion pipeline profiling omni_config_group.add_argument( @@ -1011,6 +1020,8 @@ def run_headless(args: TrackingNamespace) -> None: # Filter down to a dict of things explicitly requested by the user args_dict = args.get_explicit_kwargs_dict() + # This controls only the API-process WebSocket and is not a stage-engine option. + args_dict.pop("robot_openpi_idle_timeout", None) # This CLI-only negative alias is consumed below when selecting the # launcher log_stats value; it is not a per-stage config override. args_dict.pop("disable_log_stats", None) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index adb7d6e9daf..565362312ba 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -485,6 +485,8 @@ async def build_async_omni_from_stage_config( try: kwargs = args.get_explicit_kwargs_dict() + # This controls only the API-process WebSocket and is not an AsyncOmni option. + kwargs.pop("robot_openpi_idle_timeout", None) model = kwargs.pop("model", None) or args.model kwargs.setdefault("log_stats", not args.disable_log_stats) async_omni = AsyncOmni(model=model, **kwargs) @@ -1494,11 +1496,24 @@ async def realtime_robot_openpi(websocket: WebSocket): serving = getattr(websocket.app.state, "openai_serving_realtime_robot", None) if serving is None: + logger.warning( + "Rejecting robot OpenPI WebSocket: policy serving is disabled. " + "The diffusion stage must provide model_config.policy_server_config." + ) await websocket.accept() - await websocket.send_json({"type": "error", "error": "Robot policy not available", "code": "unsupported"}) + await websocket.send_json( + { + "type": "error", + "error": "Robot policy not available", + "code": "unsupported", + } + ) await websocket.close() return - connection = RobotRealtimeConnection(websocket, serving) + state_args = getattr(websocket.app.state, "args", None) + configured_timeout = getattr(state_args, "robot_openpi_idle_timeout", 30.0) + idle_timeout = None if configured_timeout == 0 else configured_timeout + connection = RobotRealtimeConnection(websocket, serving, idle_timeout=idle_timeout) await connection.handle_connection() diff --git a/vllm_omni/entrypoints/openpi/connection.py b/vllm_omni/entrypoints/openpi/connection.py index 9b8bfe0d812..112254fbe5f 100644 --- a/vllm_omni/entrypoints/openpi/connection.py +++ b/vllm_omni/entrypoints/openpi/connection.py @@ -29,6 +29,7 @@ from __future__ import annotations import asyncio +from collections import OrderedDict from typing import Any import msgspec @@ -45,10 +46,12 @@ _DEFAULT_IDLE_TIMEOUT = 30.0 MAX_OPENPI_PAYLOAD_BYTES = 64 * 1024 * 1024 _MISSING = object() +# Upper bound on the per-connection set of seen session ids. +MAX_TRACKED_SESSIONS = 1024 def _pack_numpy(obj: Any) -> Any: - if isinstance(obj, (np.ndarray, np.generic)) and obj.dtype.kind in ("V", "O", "c"): + if isinstance(obj, np.ndarray | np.generic) and obj.dtype.kind in ("V", "O", "c"): raise ValueError(f"Unsupported dtype: {obj.dtype}") if isinstance(obj, np.ndarray): if not obj.flags.c_contiguous: @@ -165,17 +168,38 @@ def __init__( self, websocket: WebSocket, serving: ServingRealtimeRobotOpenPI, - idle_timeout: float = _DEFAULT_IDLE_TIMEOUT, + idle_timeout: float | None = _DEFAULT_IDLE_TIMEOUT, ) -> None: self.websocket = websocket self.serving = serving self._idle_timeout = idle_timeout self._current_session_id: str | None = None - self._call_count = 0 + # Session ids seen on this connection, most-recently-used last. + self._seen_sessions: OrderedDict[str, None] = OrderedDict() def reset(self) -> None: self._current_session_id = None - self._call_count = 0 + self._seen_sessions.clear() + + def _mark_session_seen(self, session_id: str) -> bool: + """Record ``session_id``; return True when it is the first sighting. + + A session's first observation is the one that carries ``reset`` to the + policy, so this needs per-session memory. + """ + if session_id in self._seen_sessions: + self._seen_sessions.move_to_end(session_id) + return False + + self._seen_sessions[session_id] = None + if len(self._seen_sessions) > MAX_TRACKED_SESSIONS: + evicted, _ = self._seen_sessions.popitem(last=False) + logger.debug( + "Robot OpenPI connection is tracking more than %d sessions; dropped least-recently-used %s", + MAX_TRACKED_SESSIONS, + evicted, + ) + return True async def _send_error(self, message: str) -> None: await self.websocket.send_bytes(_pack({"type": "error", "message": message})) @@ -199,18 +223,22 @@ async def handle_connection(self) -> None: await self.websocket.send_bytes(_pack(metadata)) while True: - try: - msg = await asyncio.wait_for( - self.websocket.receive(), - timeout=self._idle_timeout, - ) - except asyncio.TimeoutError: - logger.info("Robot OpenPI connection idle timeout after %.1f seconds", self._idle_timeout) + idle_timeout = self._idle_timeout + if idle_timeout is None: + msg = await self.websocket.receive() + else: try: - await self.websocket.close() - except Exception: - logger.debug("Failed to close idle robot OpenPI websocket", exc_info=True) - return + msg = await asyncio.wait_for( + self.websocket.receive(), + timeout=idle_timeout, + ) + except asyncio.TimeoutError: + logger.info("Robot OpenPI connection idle timeout after %.1f seconds", idle_timeout) + try: + await self.websocket.close() + except Exception: + logger.debug("Failed to close idle robot OpenPI websocket", exc_info=True) + return if msg.get("type") == "websocket.disconnect": break @@ -245,13 +273,12 @@ async def handle_connection(self) -> None: session_id, ) self._current_session_id = session_id - self._call_count = 0 - self._call_count += 1 + reset = self._mark_session_seen(session_id) actions = await self.serving.infer( obs, session_id=session_id, - reset=self._call_count <= 1, + reset=reset, ) await self.websocket.send_bytes(_pack(actions)) except Exception: diff --git a/vllm_omni/entrypoints/openpi/serving.py b/vllm_omni/entrypoints/openpi/serving.py index f2448dfa19d..80a3839abd7 100644 --- a/vllm_omni/entrypoints/openpi/serving.py +++ b/vllm_omni/entrypoints/openpi/serving.py @@ -148,16 +148,32 @@ def _build_request(self, obs: dict, *, session_id: str, reset: bool) -> Any: `AsyncOmni.generate()` and routed to the diffusion stage. """ from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.entrypoints.openai.stage_params import ( + clone_sampling_params, + get_default_sampling_params_list, + ) from vllm_omni.inputs.data import OmniDiffusionSamplingParams - # An optional integer ``seed`` in the inference message becomes the engine request's - # ``sampling_params.seed``; omitted, ``OmniDiffusionRequest`` assigns a random one. + # The engine applies stage default_sampling_params only to requests + # that carry no explicit params; this endpoint always passes explicit + # params, so start from a clone of the diffusion stage's defaults + # (e.g. a policy deploy yaml's ``extra_args``) and layer the OpenPI + # protocol fields on top. seed = obs.pop("seed", None) - extra_args = { - "reset": reset, - "session_id": session_id, - "robot_obs": obs, - } + sampling_params = OmniDiffusionSamplingParams() + for default_params in get_default_sampling_params_list(self.engine_client): + if isinstance(default_params, OmniDiffusionSamplingParams): + sampling_params = clone_sampling_params(default_params) + break + + extra_args = sampling_params.extra_args or {} + extra_args.update( + { + "reset": reset, + "session_id": session_id, + "robot_obs": obs, + } + ) prompt = obs.get("prompt", "") sampling_params = OmniDiffusionSamplingParams( diff --git a/vllm_omni/model_executor/models/cosmos3/__init__.py b/vllm_omni/model_executor/models/cosmos3/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/model_executor/models/cosmos3/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/model_executor/models/cosmos3/pipeline.py b/vllm_omni/model_executor/models/cosmos3/pipeline.py new file mode 100644 index 00000000000..d26b4f46c80 --- /dev/null +++ b/vllm_omni/model_executor/models/cosmos3/pipeline.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cosmos3 policy single-stage topology for online OpenPI serving. + +Cosmos3 policy checkpoints (e.g. ``nvidia/Cosmos3-Nano-Policy-DROID``) are one +diffusion stage: robot observation -> action chunk. They are +registered in ``OMNI_PIPELINES`` because online OpenPI serving +(``/v1/realtime/robot/openpi``) requires a registered pipeline so the deploy +yaml's ``model_config.policy_server_config`` reaches the websocket handshake. + +This pipeline declares neither ``hf_architectures`` nor +``diffusers_class_name``: policy checkpoints share their HF metadata +(``model_type=cosmos3_omni``, ``model_index.json`` ``_class_name= +Cosmos3OmniDiffusersPipeline``) with the T2I/video Cosmos3 checkpoints, which +must keep resolving through the default single-stage diffusion fallback. +Select this pipeline explicitly via a deploy yaml ``pipeline:`` key, e.g. +``vllm serve nvidia/Cosmos3-Nano-Policy-DROID --omni --deploy-config +/absolute/path/to/vllm-omni/vllm_omni/deploy/cosmos3_policy_droid.yaml``. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +COSMOS3_POLICY_PIPELINE = PipelineConfig( + model_type="cosmos3_policy", + model_arch="Cosmos3OmniDiffusersPipeline", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="diffusion", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(), + final_output=True, + final_output_type="action", + model_arch="Cosmos3OmniDiffusersPipeline", + ), + ), +) diff --git a/vllm_omni/model_extras/cosmos3.py b/vllm_omni/model_extras/cosmos3.py index 01cecf94fca..452946b087c 100644 --- a/vllm_omni/model_extras/cosmos3.py +++ b/vllm_omni/model_extras/cosmos3.py @@ -35,6 +35,7 @@ "resolution", "image_size", "use_state", + "format_prompt_as_json", "observation", "robot_obs", "deterministic_seed", From 627630d7e4725d94c84b64e9996f97bdec2044fe Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Fri, 11 Sep 2026 00:11:27 +0800 Subject: [PATCH 61/93] [BugFix][CI] Restore diff-aware source filtering for post-merge L3 (#7371) Signed-off-by: wangyu <410167048@qq.com> --- .../common/ci_source_file_dependencies.yml | 16 +++++++---- .buildkite/common/scripts/upload_pipeline.py | 27 +++++++++---------- .buildkite/cuda/bootstrap-upload-steps.yml | 11 +++----- .buildkite/cuda/test-merge.yml | 3 ++- .buildkite/cuda/test-nightly.yml | 1 + docs/contributing/ci/ci_settings.md | 2 +- tests/buildkite/test_upload_pipeline.py | 24 +++++++++-------- 7 files changed, 45 insertions(+), 39 deletions(-) diff --git a/.buildkite/common/ci_source_file_dependencies.yml b/.buildkite/common/ci_source_file_dependencies.yml index 0455922c8f8..724794c9526 100644 --- a/.buildkite/common/ci_source_file_dependencies.yml +++ b/.buildkite/common/ci_source_file_dependencies.yml @@ -11,6 +11,8 @@ # YAML anchors are nested lists; the loader flattens them. x-qwen3-tts: &qwen3_tts + - vllm_omni/engine/async_omni_engine.py + - vllm_omni/entrypoints/omni_base.py - vllm_omni/worker/gpu_model_runner.py - vllm_omni/model_executor/models/qwen3_tts/ - vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -23,6 +25,8 @@ x-qwen3-tts: &qwen3_tts - vllm_omni/engine/stage_pool.py x-minicpmo-4-5: &minicpmo_4_5 + - vllm_omni/engine/async_omni_engine.py + - vllm_omni/entrypoints/omni_base.py - vllm_omni/model_executor/models/minicpmo_4_5/ - vllm_omni/model_executor/models/step_audio2/ - vllm_omni/model_executor/models/cosyvoice3/code2wav_core/ @@ -43,7 +47,6 @@ x-minicpmo-4-5-duplex: &minicpmo_4_5_duplex - vllm_omni/core/sched/ - vllm_omni/outputs/multimodal_accumulation.py - vllm_omni/engine/stage_init_utils.py - - vllm_omni/engine/async_omni_engine.py - vllm_omni/engine/orchestrator.py - vllm_omni/entrypoints/cli/serve.py - vllm_omni/entrypoints/openai/api_server.py @@ -51,11 +54,15 @@ x-minicpmo-4-5-duplex: &minicpmo_4_5_duplex - vllm_omni/config/stage_config.py x-qwen2-5-omni: &qwen2_5_omni + - vllm_omni/engine/async_omni_engine.py + - vllm_omni/entrypoints/omni_base.py - vllm_omni/model_executor/models/qwen2_5_omni/ - vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py - vllm_omni/deploy/qwen2_5_omni.yaml x-qwen3-omni: &qwen3_omni + - vllm_omni/engine/async_omni_engine.py + - vllm_omni/entrypoints/omni_base.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 @@ -116,11 +123,9 @@ x-lingbot-world: &lingbot_world source_file_dependencies: - # Shared uploader/registry paths. Not a job key — do not attach this to a step. - # Used only when no job-key prefix matches: then a hit keeps every job. + # Hardware presets and shared test helpers. Not a job key — do not attach this + # to a step. Used only when no job-key prefix matches: then a hit keeps every job. source_filter_fallback: - - .buildkite/common/scripts/upload_pipeline.py - - .buildkite/common/ci_source_file_dependencies.yml - .buildkite/common/ci_mirror_hardwares.yml - tests/helpers/ @@ -280,6 +285,7 @@ source_file_dependencies: - *wan22 - tests/e2e/offline_inference/test_wan22_t2v.py - tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py + - tests/e2e/offline_inference/test_device_postprocess_equivalence.py - tests/e2e/online_serving/test_wan22_t2v.py - tests/e2e/online_serving/test_wan22_expansion.py diff --git a/.buildkite/common/scripts/upload_pipeline.py b/.buildkite/common/scripts/upload_pipeline.py index c8587d8c301..71ed751195e 100644 --- a/.buildkite/common/scripts/upload_pipeline.py +++ b/.buildkite/common/scripts/upload_pipeline.py @@ -13,12 +13,14 @@ Test pipeline mode (e.g. test-merge.yml, test-nightly.yml, test-weekly.yml): - Drop steps whose ``source_file_dependencies`` do not match changed files (string/list keys from ci_source_file_dependencies.yml, or inline path - prefixes). Filtering applies on **PR label** uploads only; ``main`` + env - schedule (NIGHTLY/WEEKLY/merge push) keeps every job and still strips the field. - If no job-key prefix matches, a change to the pipeline YAML being uploaded - or to a path under the ``source_filter_fallback`` registry key keeps every - job so command, env, and hardware edits can be validated before merge. - When any job-key prefix already matches, normal filtering wins. + prefixes). Filtering applies on **PR label** uploads and on post-merge + ``main`` L3 uploads. Scheduled ``main`` uploads that pass ``--all`` + (NIGHTLY/WEEKLY) or ``--e2e`` keep every selected job and still strip + the field. If no job-key prefix matches, a change to the pipeline YAML + being uploaded or to a path under the ``source_filter_fallback`` registry + key keeps every job so command, env, and hardware edits can be validated + before merge. When any job-key prefix already matches, normal filtering + wins. - Expand uploader-only ``mirror_hardwares`` into ``agents`` (+ optional ``image`` for NPU) + ``plugins`` (see ci_mirror_hardwares.yml). - Omit ``mirror_hardwares`` to compose ``{chip}_{n}`` from pytest ``-m`` SKU @@ -780,17 +782,14 @@ def _changed_files_for_source_filter( ) -> list[str] | None: """Return the diff to filter against, or None to keep every step. - ``source_file_dependencies`` is label-only: scheduled ``main`` uploads - (NIGHTLY/WEEKLY/post-merge) run the full pipeline. ``--all`` / ``--e2e`` - also disable filtering. Pipeline YAML / ``source_filter_fallback`` matches - are applied later in ``_render_test_pipeline`` only when no job-key prefix - already matched. + ``--all`` / ``--e2e`` disable filtering (scheduled NIGHTLY/WEEKLY full + uploads and weekly E2E sweeps). Post-merge ``main`` L3 and PR-label + uploads filter against the commit/PR diff. Pipeline YAML / + ``source_filter_fallback`` matches are applied later in + ``_render_test_pipeline`` only when no job-key prefix already matched. """ if force_all or e2e_only: return None - if os.environ.get("BUILDKITE_BRANCH", "") == "main": - _log("main branch: keep all jobs (source_file_dependencies is label-only)") - return None return ctx.changed_files diff --git a/.buildkite/cuda/bootstrap-upload-steps.yml b/.buildkite/cuda/bootstrap-upload-steps.yml index 502a506e6d6..3f5aa1cb1f9 100644 --- a/.buildkite/cuda/bootstrap-upload-steps.yml +++ b/.buildkite/cuda/bootstrap-upload-steps.yml @@ -22,19 +22,16 @@ steps: agents: queue: "cpu_queue_premerge" - # L3 Test — PR merge-test label (diff-aware) or main (full pipeline / WEEKLY --e2e) + # L3 Test — PR merge-test label or post-merge on main (both diff-aware); + # main+WEEKLY=1 uploads the E2E group only (--e2e). - label: "Upload Merge Pipeline" depends_on: image-build key: upload-merge-pipeline commands: - | EXTRA="" - if [ "${BUILDKITE_BRANCH}" = "main" ]; then - if [ "${WEEKLY:-}" = "1" ]; then - EXTRA="--e2e" - else - EXTRA="--all" - fi + if [ "${WEEKLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then + EXTRA="--e2e" fi python3 .buildkite/common/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/cuda/test-merge.yml agents: diff --git a/.buildkite/cuda/test-merge.yml b/.buildkite/cuda/test-merge.yml index 1a458bccf68..d600e8fbef8 100644 --- a/.buildkite/cuda/test-merge.yml +++ b/.buildkite/cuda/test-merge.yml @@ -1,5 +1,6 @@ # CUDA CI L3 (`advanced_model` e2e) plus L1 CPU Simple Test (`core_model and cpu`). -# Uploaded by Upload Merge Pipeline: PR label `merge-test`, or post-merge on main. +# Uploaded by Upload Merge Pipeline: PR label `merge-test`, or post-merge on main +# (both diff-aware via source_file_dependencies). env: VLLM_WORKER_MULTIPROC_METHOD: spawn diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index 583f935dc25..cfaca90fe1b 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -390,6 +390,7 @@ steps: mirror_hardwares: h100_2 - label: ":full_moon: Diffusion X2V · Wan2.2 Device-Postprocess Equivalence Test" + source_file_dependencies: diffusion_wan22_function timeout_in_minutes: 90 commands: - | diff --git a/docs/contributing/ci/ci_settings.md b/docs/contributing/ci/ci_settings.md index fb3af0ed3cb..9931a418d07 100644 --- a/docs/contributing/ci/ci_settings.md +++ b/docs/contributing/ci/ci_settings.md @@ -313,7 +313,7 @@ Unit coverage: `tests/buildkite/test_skip_ci.py`. CUDA **L2** (`.buildkite/cuda/test-ready.yml`), **L3** (`.buildkite/cuda/test-merge.yml`), **L4** (`.buildkite/cuda/test-nightly.yml`), **L5** (`.buildkite/cuda/test-weekly.yml`), and NPU **L4** (`.buildkite/npu/test-npu-nightly.yml`). Bootstrap upload entry: `upload_pipeline.py --upload .buildkite/cuda/bootstrap-upload-steps.yml` (or `npu/bootstrap-upload-steps.yml`). -**When filtering applies:** only **PR label** uploads (`ready`, `merge-test`, `nightly-test`, `weekly-test`). On **`main` + env** (`NIGHTLY=1`, post-merge L3, `WEEKLY=1` / `NON_CRITICAL=1`, `WEEKLY=1 --e2e`) the uploader keeps every job and still strips the key. If **no** job-key prefix matches the diff, a change to the pipeline YAML being uploaded, or any path under the `source_filter_fallback` key in `ci_source_file_dependencies.yml`, keeps every job in that upload. `source_filter_fallback` is not a job key — do not attach it to a step. When any job-key prefix already matches, normal filtering wins (the fallback does not expand the selection). Without this fallback, command / env / hardware-only edits would match nothing (NPU drops every step; CUDA nightly keeps only the email aggregator). A different pipeline YAML does not trigger the fallback for this upload. +**When filtering applies:** **PR label** uploads (`ready`, `merge-test`, `nightly-test`, `weekly-test`) and **post-merge `main` L3** (`test-merge.yml`). Scheduled **`main` + env** uploads that pass `--all` (`NIGHTLY=1`, `WEEKLY=1` / `NON_CRITICAL=1`) or `--e2e` (`WEEKLY=1` on L2/L3) keep every selected job and still strip the key. If **no** job-key prefix matches the diff, a change to the pipeline YAML being uploaded, or any path under the `source_filter_fallback` key in `ci_source_file_dependencies.yml`, keeps every job in that upload. `source_filter_fallback` is not a job key — do not attach it to a step. When any job-key prefix already matches, normal filtering wins (the fallback does not expand the selection). **Uploader-only keys** — removed before Buildkite sees the YAML; never used at runtime on agents: diff --git a/tests/buildkite/test_upload_pipeline.py b/tests/buildkite/test_upload_pipeline.py index 06e2b975c3c..a3192dada51 100644 --- a/tests/buildkite/test_upload_pipeline.py +++ b/tests/buildkite/test_upload_pipeline.py @@ -692,19 +692,24 @@ def test_source_filter_strips_deps_and_expands_hardware(monkeypatch: pytest.Monk assert z_image["agents"]["queue"] == "l4-k8s" -def test_source_filter_disabled_on_main_branch(monkeypatch: pytest.MonkeyPatch) -> None: +def test_source_filter_respects_force_all_and_uses_diff_on_main( + monkeypatch: pytest.MonkeyPatch, +) -> None: class _Ctx: changed_files = ["vllm_omni/unrelated.py"] + # Post-merge main L3 must still filter by the commit diff. monkeypatch.setenv("BUILDKITE_BRANCH", "main") - assert _changed_files_for_source_filter(_Ctx(), force_all=False, e2e_only=False) is None + assert _changed_files_for_source_filter(_Ctx(), force_all=False, e2e_only=False) == [ + "vllm_omni/unrelated.py", + ] monkeypatch.setenv("BUILDKITE_BRANCH", "feat/source-filter") assert _changed_files_for_source_filter(_Ctx(), force_all=False, e2e_only=False) == [ "vllm_omni/unrelated.py", ] - monkeypatch.setenv("BUILDKITE_BRANCH", "main") assert _changed_files_for_source_filter(_Ctx(), force_all=True, e2e_only=False) is None + assert _changed_files_for_source_filter(_Ctx(), force_all=False, e2e_only=True) is None def test_source_filter_fallback_is_fallback_when_no_job_key_matches( @@ -715,7 +720,6 @@ def test_source_filter_fallback_is_fallback_when_no_job_key_matches( monkeypatch.setenv("BUILDKITE_BRANCH", "feat/nightly-yaml") pipeline_yaml = Path(".buildkite/npu/test-npu-nightly.yml") shared_paths = _load_source_file_dependencies()["source_filter_fallback"] - assert ".buildkite/common/scripts/upload_pipeline.py" in shared_paths # Synthetic steps only — do not pin live Buildkite job labels. doc = { @@ -742,7 +746,7 @@ def surviving_keys(changed_files: list[str]) -> set[str]: ) return {step["key"] for step in _iter_steps(rendered) if isinstance(step.get("key"), str)} - # Only pipeline YAML / shared uploader → no listed prefix match → keep every step. + # Only pipeline YAML / fallback paths → no listed prefix match → keep every step. for changed in [pipeline_yaml.as_posix(), *shared_paths]: assert surviving_keys([changed]) == {"ungated", "gated_a", "gated_b"}, changed @@ -750,9 +754,7 @@ def surviving_keys(changed_files: list[str]) -> set[str]: assert surviving_keys([".buildkite/cuda/test-nightly.yml"]) == {"ungated"} # A matching source prefix wins over fallback files in the same diff. - assert surviving_keys( - [ - "pkg/model_a/transformer.py", - ".buildkite/common/scripts/upload_pipeline.py", - ], - ) == {"ungated", "gated_a"} + assert surviving_keys(["pkg/model_a/transformer.py", *shared_paths[:1]]) == { + "ungated", + "gated_a", + } From ffcaaa9439823d9b927d201397bb9d4db7218f14 Mon Sep 17 00:00:00 2001 From: Guangjian Dong <163994576+Hiro208@users.noreply.github.com> Date: Fri, 11 Sep 2026 02:45:28 +0800 Subject: [PATCH 62/93] [Bugfix] Fail when a diffusion LoRA adapter binds no layer (#7349) Signed-off-by: Guangjian --- tests/diffusion/lora/test_lora_manager.py | 44 +++++++++++++++++++++++ vllm_omni/diffusion/lora/manager.py | 13 +++++-- 2 files changed, 54 insertions(+), 3 deletions(-) diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py index 0814bf15f1f..10b023bf0ee 100644 --- a/tests/diffusion/lora/test_lora_manager.py +++ b/tests/diffusion/lora/test_lora_manager.py @@ -411,6 +411,50 @@ def adapter(adapter_id: int): assert len(first.set_calls) == first_calls_after_success + 2 +def test_lora_manager_rejects_adapter_that_binds_no_layer(): + """An adapter whose target modules match nothing must fail, not silently no-op.""" + manager = DiffusionLoRAManager( + pipeline=torch.nn.Module(), + device=torch.device("cpu"), + dtype=torch.bfloat16, + max_cached_adapters=1, + ) + layer = _DummyLoRALayer(n_slices=1, output_slices=(2,)) + manager._lora_modules = {"transformer.blocks.0.attn.to_q": layer} + + # Adapter-side name the engine does not expose, e.g. a diffusers-style + # checkpoint against a differently named engine layout. + mismatched = LoRALayerWeights( + module_name="unet.down_blocks.0.attn.to_q", + rank=2, + lora_alpha=2, + lora_a=torch.ones((2, 2)), + lora_b=torch.ones((2, 2)), + ) + manager._registered_adapters = { + 3: type( + "LM", + (), + { + "id": 3, + "loras": {"unet.down_blocks.0.attn.to_q": mismatched}, + "get_lora": lambda self, key: self.loras.get(key), + }, + )() + } + + with pytest.raises(ValueError, match="applies to no layer") as excinfo: + manager._activate_adapter(3, scale=1.0) + + # The message must name what was received so the mismatch is diagnosable. + assert "unet.down_blocks.0.attn.to_q" in str(excinfo.value) + + # Nothing was bound and the adapter must not be left marked active. + assert manager._active_adapter_id is None + assert layer.set_calls == [] + assert layer.reset_calls >= 1 + + def _dummy_lora_request(adapter_id: int) -> LoRARequest: return LoRARequest( lora_name=f"adapter_{adapter_id}", diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 2f6d4778056..11d19353158 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -556,9 +556,9 @@ def _is_active_at_scale(self, adapter_id: int, scale: float) -> bool: def _bind_adapter_weights(self, lora_model: LoRAModel, scale: float) -> None: binding_validator = getattr(self.pipeline, "_validate_diffusion_lora_binding", None) - lora_names_by_id = ( - {id(weights): name for name, weights in lora_model.loras.items()} if callable(binding_validator) else {} - ) + # Track bindings unconditionally. The zero-binding guard below needs this + # bookkeeping on every pipeline, not only the ones that supply a validator. + lora_names_by_id = {id(weights): name for name, weights in lora_model.loras.items()} bound_lora_names: set[str] = set() def _record_bound(weights: LoRALayerWeights | PackedLoRALayerWeights) -> None: @@ -677,6 +677,13 @@ def _record_bound(weights: LoRALayerWeights | PackedLoRALayerWeights) -> None: scale, ) + if not bound_lora_names: + raise ValueError( + f"LoRA adapter {lora_model.id} applies to no layer: expected target modules in " + f"{sorted(self._expected_lora_modules)} but received {sorted(lora_model.loras)}. " + "Activating it would leave the base model unchanged." + ) + if callable(binding_validator): binding_validator( lora_model=lora_model, From 5a386d789f7646f42e7c97fe873848de9dd400aa Mon Sep 17 00:00:00 2001 From: summer <128961079+zhang-keliang@users.noreply.github.com> Date: Fri, 11 Sep 2026 08:04:08 +0800 Subject: [PATCH 63/93] [Bugfix] Fix host-memory leak on aborted /v1/images/generations (#6462) (#6561) Signed-off-by: summer <128961079+zhang-keliang@users.noreply.github.com> Co-authored-by: Hongsheng Liu --- tests/entrypoints/test_async_omni.py | 89 +++++++++++++++++++ tests/metrics/test_emit_calls.py | 3 +- tests/metrics/test_prometheus.py | 2 - vllm_omni/entrypoints/client_request_state.py | 6 ++ vllm_omni/entrypoints/omni_base.py | 18 +--- 5 files changed, 100 insertions(+), 18 deletions(-) diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index 2549889b3de..a69bcce099d 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/tests/entrypoints/test_async_omni.py @@ -329,6 +329,95 @@ async def empty_abort_async(request_ids): asyncio.run(run()) +@pytest.mark.cpu +def test_abort_drops_consumed_metric_message_ids_with_state(): + """Exercise the production leak path from #6462: a metric-bearing + ``OutputMessage`` goes through ``_handle_output_message`` (which populates + the per-request de-dup set), the request is aborted (which, since #6367, + keeps ``request_states`` registered and enqueues a terminal abort output), + and the set must be released with the state by ``generate()``'s normal + cleanup — with no independent request-keyed metric state surviving on the + instance at any point. A reintroduced + class-level ``_consumed_metric_messages``-style map populated by the + production handler fails the sweep below. + """ + import time as _time + from collections.abc import Mapping + + from vllm_omni.engine.messages import OutputMessage + from vllm_omni.entrypoints.client_request_state import ClientRequestState + from vllm_omni.metrics.stats import OrchestratorAggregator + + async def run(): + omni = get_async_omni_instance() + omni.engine.get_stage_metadata = lambda stage_id: SimpleNamespace( + stage_type="llm", + final_output=True, + final_output_type="text", + ) + + rid = "req-1-cccc" + state = ClientRequestState(request_id=rid, external_request_id="req-1") + state.metrics = OrchestratorAggregator( + num_stages=1, + log_stats=False, + wall_start_ts=_time.time(), + final_stage_id_for_e2e=0, + ) + state.input_stream_task = None + omni.request_states[rid] = state + + msg = OutputMessage( + request_id=rid, + stage_id=0, + engine_outputs=SimpleNamespace(final_output_type="text"), + metrics=SimpleNamespace( + num_tokens_in=3, + num_tokens_out=5, + stage_gen_time_ms=1.0, + rx_transfer_bytes=0, + rx_decode_time_ms=0.0, + rx_in_flight_time_ms=0.0, + ), + finished=False, + ) + handled, out_rid, out_stage, out_state = omni._handle_output_message(msg) + assert handled is False and out_state is state + # The production handler recorded the de-dup entry on the state... + assert id(msg) in state.consumed_metric_message_ids + # ...and de-duplicates a replay of the same message object. + omni._handle_output_message(msg) + assert len(state.consumed_metric_message_ids) == 1 + + assert not hasattr(omni, "_consumed_metric_messages") + + # #6367 contract: abort() keeps the state registered and enqueues a + # terminal abort output; generate()'s ``finally`` owns the cleanup. + await omni.abort("req-1") + assert rid in omni.request_states + terminal = state.queue.get_nowait() + assert terminal.finished is True + assert terminal.engine_outputs.outputs[0].finish_reason == "abort" + # The de-dup set still lives only on the (still-registered) state. + for name, value in vars(omni).items(): + if name == "request_states": + continue + assert not (isinstance(value, Mapping) and rid in value), ( + f"request-keyed metric state kept outside request_states in {name!r}" + ) + + # generate()'s normal cleanup releases the state and the set with it. + omni._log_summary_and_cleanup(rid) + assert rid not in omni.request_states + for name, value in vars(omni).items(): + assert not (isinstance(value, Mapping) and rid in value), ( + f"request-keyed state survived cleanup in {name!r}" + ) + assert not hasattr(omni, "_consumed_metric_messages") + + asyncio.run(run()) + + @pytest.mark.cpu def test_generate_accepts_request_after_repeated_cancellations(): async def run_test(): diff --git a/tests/metrics/test_emit_calls.py b/tests/metrics/test_emit_calls.py index b46f4080fa5..02a98fc00e7 100644 --- a/tests/metrics/test_emit_calls.py +++ b/tests/metrics/test_emit_calls.py @@ -41,7 +41,6 @@ def _make_omni_base_with_mock_prom(mocker): obj = object.__new__(OmniBase) obj.prom_metrics = mocker.Mock(spec=OmniPrometheusMetrics) obj.request_states = {} - obj._consumed_metric_messages = {} obj.log_stats = True return obj, obj.prom_metrics @@ -270,7 +269,7 @@ def test_same_finished_image_message_is_observed_exactly_once(self, mocker, peak obj = object.__new__(OmniBase) obj._enable_ar_profiler = False - obj._consumed_metric_messages = {} + obj.request_states = {"req-replay": SimpleNamespace(consumed_metric_message_ids=set())} obj.prom_metrics = mocker.Mock(spec=OmniPrometheusMetrics) obj.mod_metrics = mocker.Mock() obj.engine = SimpleNamespace( diff --git a/tests/metrics/test_prometheus.py b/tests/metrics/test_prometheus.py index e45365cf1ea..ca8dd8d5e1e 100644 --- a/tests/metrics/test_prometheus.py +++ b/tests/metrics/test_prometheus.py @@ -154,7 +154,6 @@ def test_running_and_waiting_zero_after_request_completes(self, registry: Collec obj.engine = SimpleNamespace(_running_counter=OmniRequestCounter()) obj.prom_metrics = OmniPrometheusMetrics(model_name="lifecycle-test") obj.request_states = {} - obj._consumed_metric_messages = {} obj.log_stats = False # Simulate request lifecycle: start (counter 0→1, dict {} → {req}), @@ -183,7 +182,6 @@ def test_gauges_reflect_remaining_requests_after_one_completes(self, registry: C obj.engine = SimpleNamespace(_running_counter=OmniRequestCounter()) obj.prom_metrics = OmniPrometheusMetrics(model_name="lifecycle-test-2") obj.request_states = {} - obj._consumed_metric_messages = {} obj.log_stats = False # Two in flight, one finalizes — running should report 1, waiting 0. diff --git a/vllm_omni/entrypoints/client_request_state.py b/vllm_omni/entrypoints/client_request_state.py index 0aa0c7c5ae6..22f1c6d5e90 100644 --- a/vllm_omni/entrypoints/client_request_state.py +++ b/vllm_omni/entrypoints/client_request_state.py @@ -39,3 +39,9 @@ def __init__( # without re-querying stage_pools. self.audio_emit_stage_id: int | None = None self.audio_emit_replica_id: int | None = None + # De-dup set for metric messages: OmniBase populates this in + # ``_handle_output_message`` / ``_process_single_result`` so the same + # ``id(msg)`` isn't counted twice into per-request metrics. Kept on + # the request state (not a class-level dict) so it is released with + # the state — see #6462 / #6561. + self.consumed_metric_message_ids: set[int] = set() diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 264b37fb527..899ff116d35 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -231,7 +231,6 @@ def __init__( self.async_chunk = bool(getattr(self.engine, "async_chunk", False)) self.request_states: dict[str, ClientRequestState] = {} - self._consumed_metric_messages: dict[str, set[int]] = {} self.mod_metrics = OmniModalityMetrics(model_name=model, log_stats=log_stats) self.default_sampling_params_list = self.engine.default_sampling_params_list @@ -298,13 +297,6 @@ def _stage_has_no_live_replica(self, pool: StagePool) -> bool: """True when a non-empty stage pool has lost all of its replicas.""" return len(pool.clients) > 0 and self._live_replica_count(pool) == 0 - def _consumed_metric_message_ids(self, request_id: str) -> set[int]: - consumed_by_request = getattr(self, "_consumed_metric_messages", None) - if consumed_by_request is None: - consumed_by_request = {} - self._consumed_metric_messages = consumed_by_request - return consumed_by_request.setdefault(request_id, set()) - @property def is_running(self) -> bool: return self.engine.is_alive() @@ -449,9 +441,6 @@ def _log_summary_and_cleanup(self, request_id: str, reason: str = "stage_error") ) finally: self.request_states.pop(request_id, None) - consumed_by_request = getattr(self, "_consumed_metric_messages", None) - if consumed_by_request is not None: - consumed_by_request.pop(request_id, None) # Republish gauges so any stale value left by the per-stage # publish in _process_single_result (which runs while the request # is still in self.request_states) is corrected after the pop. @@ -532,7 +521,7 @@ def _handle_output_message( stage_meta = self.engine.get_stage_metadata(stage_id) output_type = getattr(msg.engine_outputs, "final_output_type", stage_meta.final_output_type) msg_id = id(msg) - consumed = self._consumed_metric_message_ids(req_id) + consumed = req_state.consumed_metric_message_ids if msg_id not in consumed: req_state.metrics.on_stage_metrics(stage_id, req_id, msg.metrics, output_type) submit_ts = msg.stage_submit_ts @@ -642,8 +631,9 @@ def _process_single_result( output_type = getattr(engine_outputs, "final_output_type", stage_meta.final_output_type) if finished and _m is not None: msg_id = id(result) - consumed = self._consumed_metric_message_ids(req_id) - if msg_id not in consumed: + req_state = self.request_states.get(req_id) + consumed = req_state.consumed_metric_message_ids if req_state is not None else None + if consumed is not None and msg_id not in consumed: metrics.accumulate_diffusion_metrics(stage_meta.stage_type, req_id, engine_outputs) metrics.on_stage_metrics(stage_id, req_id, _m, output_type) consumed.add(msg_id) From 5c69fc2a8a6bca8f87505d1e0cad112daa5941e4 Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Thu, 10 Sep 2026 17:44:32 -0700 Subject: [PATCH 64/93] [Refactor] Declare model-local KV held outside the paged manager (#6171) Signed-off-by: Yueqian Lin --- docs/design/feature/model_local_kv_caches.md | 132 ++++++ docs/design/index.md | 1 + .../models/test_model_local_kv.py | 404 ++++++++++++++++++ tests/tools/test_check_tts_adapter.py | 12 +- .../models/mimo_audio/mimo_audio_llm.py | 67 +++ .../ming_flash_omni/ming_flash_omni_talker.py | 11 + .../models/ming_flash_omni/talker_module.py | 58 ++- .../minicpmo_4_5/minicpmo_4_5_omni_llm.py | 52 +++ .../model_executor/models/model_local_kv.py | 347 +++++++++++++++ .../qwen3_tts/segmented_graph_wrapper.py | 86 ++++ .../modeling_qwen3_tts_tokenizer_v2.py | 69 +++ vllm_omni/worker/gpu_model_runner.py | 47 ++ 12 files changed, 1283 insertions(+), 3 deletions(-) create mode 100644 docs/design/feature/model_local_kv_caches.md create mode 100644 tests/model_executor/models/test_model_local_kv.py create mode 100644 vllm_omni/model_executor/models/model_local_kv.py diff --git a/docs/design/feature/model_local_kv_caches.md b/docs/design/feature/model_local_kv_caches.md new file mode 100644 index 00000000000..371b989ff13 --- /dev/null +++ b/docs/design/feature/model_local_kv_caches.md @@ -0,0 +1,132 @@ +# Model-local KV caches + +Some models keep their attention KV as HuggingFace `transformers` cache objects +instead of using the engine's paged KV manager. Each such model declares what it +holds by implementing `model_local_kv_specs()`, and +`OmniGPUModelRunner.load_model` sums and logs the result after load. + +## What is and is not already accounted for + +This distinction decides what the protocol is for, and an earlier revision of +this page got it wrong by claiming none of this memory was accounted for. + +CUDA-graph capture happens *during* `load_model` -- MiMo in its constructor, +Qwen3-TTS in `load_weights`. `determine_available_memory()` then samples +per-process GPU memory through NVML and subtracts it from the requested budget +(`vllm_omni/worker/base.py`). So the **graph-resident** portion is already +charged. Subtracting the declared total from the KV budget would double-count +it. + +What is *not* charged is the **per-request** portion, which is allocated after +profiling. That half is not small: Qwen3-TTS's codec decoder is 4.44 MiB per +request, so a server at `max_num_seqs=256` reaches roughly 1.1 GiB that no +per-stage number predicted. + +The consumer therefore reports and does not subtract. + +## What the four models declare + +| Model | Cache | Bounded by | Per row | Rows | +|---|---|---|---|---| +| Qwen3-TTS codec decoder (async-chunk only) | sliding `DynamicCache` | `sliding_window - 1 = 71` | 4.44 MiB | `max_num_seqs`, x2 for the working copy | +| Qwen3-TTS graph pool | retained captures | same | 4.44 MiB | fixed, from captured shapes | +| MiMo-Audio local transformer | `DynamicCache` | `group_size + max(delay_pattern) = 11` | 704 KiB | `max_num_seqs` | +| MiMo-Audio graph pool | captured | same | 704 KiB | fixed 261 (bucket sum) | +| MiniCPM-o Whisper encoder (streaming only) | `EncoderDecoderCache` | `embed_positions` rows = 1500 | 140.62 MiB | fixed 1, per session | +| ming_flash_omni talker | `StaticCache` | hardcoded `max_cache_len = 2048` | 24.00 MiB | fixed 1 | + +**No cache is bounded by `max_model_len`.** Each has its own mechanism: a +sliding window, a decode-loop trip count, an encoder frame limit, a constant. + +**ming's geometry does not exist in this repo.** `self._llm_config` is a +`Qwen2Config` from the checkpoint, so layers, kv-heads, head-dim and dtype are +only knowable after load. That is what makes this a runtime query rather than a +table. + +## Lifetime is not multiplicity + +`scope` records when memory appears and goes. It does **not** decide how wide +the allocation is, and an earlier revision that derived scaling from it was +wrong in two different directions at once: + +- ming's talker is per-call but serialized -- `forward` takes + `runtime_additional_information[0]` and runs the whole AR loop inline -- so + scaling it by `max_num_seqs` over-reported it by that factor. +- MiniCPM-o's encoder is per *session*, not per sequence, so `max_num_seqs` is + simply the wrong number for it. + +So width is declared explicitly, by naming the driver. `RowDriver` has two +values, because two are all the known caches need: + +- `FIXED` -- a count the model controls, with `rows_fixed` and a required + `rows_reason`. `rows_fixed=1` on a cache that looks per-request is a claim a + reviewer must be able to check. +- `MAX_NUM_SEQS` -- one row per in-flight sequence; only the engine knows the + value. + +A third value belongs here when a model actually has a cache the engine widens +by some other number. An earlier revision added one speculatively, for a single +declarer, and paid for it with a class hierarchy, an engine-capacity object and +a conditional-annotation field. + +`rows` counts row-equivalents at peak and deliberately does not distinguish +"one allocation N rows wide" from "N allocations of one row". Those cost the +same, and a revision that tried to carry both got the object topology wrong in +three of four declarers: MiniCPM-o keeps one cache object per session while +asserting batch size 1, and both graph pools keep one distinct object per +captured bucket. The layout belongs in `allocation_note`, where it cannot be +mistaken for arithmetic. + +## Declaring nothing is a valid declaration + +A declaration must be inert when its path is. Qwen3-TTS holds KV only in +async-chunk mode -- stateless decoding runs `_forward_exact`, which calls the +transformer without `use_cache` -- so the stateless path declares nothing. +Getting this wrong is easy in both directions, and both have happened here: one +revision hung the declaration off the CUDA-graph wrapper and reported zero under +`enforce_eager`, where the eager async-chunk path does allocate; the next +declared unconditionally and invented roughly a gigabyte for stateless +deployments. + +Where the model cannot see the deciding setting, it declares the honest unit +and says so in `allocation_note`. MiniCPM-o's encoder cache exists only on the +streaming path and is one object per session; the session cap is not visible +from the model, so it declares one session rather than scaling by a number it +would have to guess at. + +## Finding the declarers + +`collect_model_local_kv_specs()` walks `named_modules()`. Two wrinkles: + +- The runner's `self.model` may be a `CUDAGraphWrapper`, which is not an + `nn.Module` but forwards attribute access to its runnable. The walk therefore + works either way, but entering through both the wrapper and its runnable would + count a root-level declarer twice. The collector calls the supported + `unwrap()` accessor first. +- An owner that is not a submodule at all (Qwen3-TTS's graph wrapper, ming's + `MingAudioGenerator`) is reached by one explicit hop from the module that + holds it. + +## Adding a model + +Return one entry per distinct allocation. Prefer `spec_from_hf_config()` and +pass the same config object the cache is built from; it raises rather than +guessing when an attribute is missing. `capacity_source` is diagnostic text -- +never branch on it. + +## Known limits + +Row counts taken from live state after warmup miss caches `transformers` has +not yet materialized, which makes Qwen3-TTS's graph-pool figure a floor for the +xvec path. + +MiMo's graph-pool cache has no reachable Python object -- it is a local inside +the captured frame -- so its width comes from the captured bucket list rather +than from tensors. Its two entries are also mutually exclusive per call: a call +either replays a bucket or runs eagerly. The sum is the right ceiling for the +process and the wrong number for a single request. + +MiMo captures every bucket gated only on `torch.cuda.is_available()` +(`mimo_audio_llm.py:670`), so those 179.44 MiB stay resident even when the +deploy config sets `enforce_eager: true`. Gating it is a one-line change that +needs a model boot to verify, so it is noted rather than bundled here. diff --git a/docs/design/index.md b/docs/design/index.md index cff0b1eb3ce..6481aa49b2f 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -26,6 +26,7 @@ implementation contract; it is not, by itself, a general support claim. - [Async Omni Output Materialization](feature/omni_async_output_materialization.md) - [Runner-to-model Prefill/Decode Phase Contract](feature/preprocess_phase_contract.md) - [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md) +- [Model-local KV Caches](feature/model_local_kv_caches.md) - [Realtime AR-Diffusion Sessions](feature/realtime_ar_diffusion.md) ### Communication diff --git a/tests/model_executor/models/test_model_local_kv.py b/tests/model_executor/models/test_model_local_kv.py new file mode 100644 index 00000000000..5e2c9cdca6a --- /dev/null +++ b/tests/model_executor/models/test_model_local_kv.py @@ -0,0 +1,404 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""A declaration must match the cache that actually gets allocated. + +Two earlier versions of this file could not fail. The first asserted hand-typed +constants against each other. The second built real `transformers` caches but +compared them to specs written in the test, so no production declaration was +ever executed. What matters here is that a real declarer is called and its +numbers are checked against a real cache built the way the model builds it. +""" + +import ast +import pathlib +from types import SimpleNamespace + +import pytest +import torch +from transformers import Qwen2Config +from transformers.cache_utils import DynamicCache, StaticCache + +from vllm_omni.model_executor.models.model_local_kv import ( + HasModelLocalKV, + ModelLocalKVScope, + ModelLocalKVSpec, + RowDriver, + collect_model_local_kv_specs, + spec_from_hf_config, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +MODELS_DIR = pathlib.Path(__file__).resolve().parents[3] / "vllm_omni" / "model_executor" / "models" + + +def _total_bytes(model: object, max_num_seqs: int = 1) -> int: + """Sum declared peak bytes across a model's declarations.""" + return sum(spec.peak_bytes(max_num_seqs) for _, spec in collect_model_local_kv_specs(model)) + + +def _cache_bytes(cache) -> int: + total = 0 + for layer in cache.layers: + for tensor in (layer.keys, layer.values): + if tensor is not None: + total += tensor.numel() * tensor.element_size() + return total + + +def _qwen2_config(layers=4, kv_heads=2, heads=8, hidden=256) -> Qwen2Config: + return Qwen2Config( + num_hidden_layers=layers, + num_attention_heads=heads, + num_key_value_heads=kv_heads, + hidden_size=hidden, + intermediate_size=hidden * 2, + vocab_size=1000, + ) + + +# -------------------------------------------------------------------------- +# A real declarer, called, checked against a real cache. +# -------------------------------------------------------------------------- + + +def test_ming_declaration_matches_the_cache_it_describes(): + """Calls MingAudioGenerator.model_local_kv_specs and checks the bytes. + + Builds the declarer without its weights -- only `_llm_config` and `_model` + are read -- then builds the `StaticCache` that `_init_kv_cache` builds from + the same config and compares. A wrong layer count, kv-head count, dtype or + capacity in the production declaration fails here. + """ + from vllm_omni.model_executor.models.ming_flash_omni.talker_module import MingAudioGenerator + + config = _qwen2_config() + generator = object.__new__(MingAudioGenerator) + generator._llm_config = config + generator._model = torch.nn.Linear(1, 1).to(torch.float16) + + (spec,) = generator.model_local_kv_specs() + + cache = StaticCache( + config=config, + max_batch_size=1, + max_cache_len=MingAudioGenerator._STATIC_CACHE_LEN, + device="cpu", + dtype=torch.float16, + ) + head_dim = config.hidden_size // config.num_attention_heads + for layer_idx in range(config.num_hidden_layers): + keys = torch.zeros(1, config.num_key_value_heads, 1, head_dim, dtype=torch.float16) + cache.update(keys, keys.clone(), layer_idx) + + assert spec.peak_bytes(1) == _cache_bytes(cache) + + +def test_ming_does_not_scale_with_max_num_seqs(): + """forward() takes runtime_additional_information[0] and loops inline. + + One worker holds one of these no matter how many sequences the scheduler + admits. A previous revision derived the multiplier from scope and reported + `max_num_seqs` times the real figure. + """ + from vllm_omni.model_executor.models.ming_flash_omni.talker_module import MingAudioGenerator + + generator = object.__new__(MingAudioGenerator) + generator._llm_config = _qwen2_config() + generator._model = torch.nn.Linear(1, 1).to(torch.float16) + + (spec,) = generator.model_local_kv_specs() + assert spec.peak_bytes(64) == spec.peak_bytes(1) + + +def test_growing_cache_declaration_matches_a_real_dynamic_cache(): + """The declared positions are the extent the tensor reaches when full.""" + layers, kv_heads, head_dim, positions = 4, 2, 32, 11 + config = _qwen2_config(layers=layers, kv_heads=kv_heads, heads=kv_heads, hidden=kv_heads * head_dim) + spec = spec_from_hf_config( + config, + name="local_transformer", + dtype=torch.bfloat16, + physical_capacity_positions=positions, + capacity_source="group_size + max(delay_pattern)", + scope=ModelLocalKVScope.INVOCATION, + rows=RowDriver.FIXED, + rows_reason="test", + ) + + cache = DynamicCache() + for layer_idx in range(layers): + keys = torch.zeros(1, kv_heads, positions, head_dim, dtype=torch.bfloat16) + cache.update(keys, torch.zeros_like(keys), layer_idx) + + assert spec.peak_bytes() == _cache_bytes(cache) + + +def test_multiplicity_does_not_claim_an_object_layout(): + """One allocation B rows wide costs the same as B one-row allocations. + + The protocol reports bytes and deliberately does not distinguish the two. + An earlier revision carried a separate `allocations` count and got the + object topology wrong in three of four declarers, which is worse than not + claiming it: the layout belongs in `allocation_note`, where it cannot be + mistaken for arithmetic. + """ + config = _qwen2_config(layers=4, kv_heads=2, heads=2, hidden=64) + common = dict( + config=config, + name="c", + dtype=torch.bfloat16, + physical_capacity_positions=11, + capacity_source="test", + scope=ModelLocalKVScope.INVOCATION, + ) + batched = spec_from_hf_config(rows=RowDriver.MAX_NUM_SEQS, **common) + + assert batched.row_count(64) == 64 + assert batched.peak_bytes(64) == batched.bytes_per_row * 64 + assert not hasattr(batched, "allocations"), "the object-count axis was removed on purpose" + + +def test_a_fixed_driver_ignores_concurrency(): + """A serialized or capture-pinned cache does not widen with max_num_seqs.""" + spec = spec_from_hf_config( + _qwen2_config(), + name="graph_pool", + dtype=torch.bfloat16, + physical_capacity_positions=11, + capacity_source="test", + scope=ModelLocalKVScope.MODEL, + rows=RowDriver.FIXED, + rows_fixed=261, + rows_reason="sum of captured buckets", + ) + assert spec.row_count(1) == 261 + assert spec.row_count(64) == 261 + assert spec.peak_bytes(64) == spec.bytes_per_row * 261 + + +def test_qwen3_tts_declares_nothing_on_the_stateless_path(): + """Calls the real decoder declarer with no cache-holding mode configured. + + `_forward_exact` runs `pre_transformer` without `use_cache`, so stateless + decoding allocates no KV. A revision that declared a per-request cache + unconditionally invented over a gigabyte at max_num_seqs=256, and the + revision before it declared nothing at all under enforce_eager, where the + eager async-chunk path does allocate. Both are regressions this pins. + """ + from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz import ( + modeling_qwen3_tts_tokenizer_v2 as mod, + ) + + decoder = object.__new__(mod.Qwen3TTSTokenizerV2Decoder) + # nn.Module.__setattr__ needs the machinery __init__ would have built, and + # this instance deliberately has none of it. + setattr_ = object.__setattr__ + setattr_( + decoder, + "config", + SimpleNamespace( + sliding_window=72, + num_hidden_layers=8, + num_key_value_heads=16, + head_dim=64, + hidden_size=1024, + num_attention_heads=16, + ), + ) + setattr_(decoder, "parameters", lambda: iter([torch.zeros(1, dtype=torch.float32)])) + + setattr_(decoder, "_cudagraph_wrapper", None) + assert decoder.model_local_kv_specs() == [] + + setattr_(decoder, "_cudagraph_wrapper", SimpleNamespace(async_chunk=False, model_local_kv_specs=lambda: [])) + assert decoder.model_local_kv_specs() == [] + + setattr_(decoder, "_cudagraph_wrapper", SimpleNamespace(async_chunk=True, model_local_kv_specs=lambda: [])) + names = [s.name for s in decoder.model_local_kv_specs()] + assert names == ["codec_decoder_request", "codec_decoder_working_copy"], names + + +def test_fixed_rejects_unreviewable_values(): + """A bare or unexplained constant cannot be reviewed, so it is rejected.""" + common = dict( + name="x", + layers=1, + kv_heads=1, + head_dim=1, + dtype=torch.float16, + physical_capacity_positions=1, + capacity_source="", + scope=ModelLocalKVScope.REQUEST, + ) + with pytest.raises(ValueError, match="rows_reason"): + ModelLocalKVSpec(rows=RowDriver.FIXED, rows_fixed=1, rows_reason=" ", **common) + with pytest.raises(ValueError, match=">= 1"): + ModelLocalKVSpec(rows=RowDriver.FIXED, rows_fixed=0, rows_reason="test", **common) + + +def test_rows_is_required(): + """Defaulting the batch extent is how the previous guarantee died. + + The dataclass declared it required while the helper every declarer uses + defaulted it to 1, so nothing enforced it where it mattered. + """ + with pytest.raises(TypeError): + spec_from_hf_config( # type: ignore[call-arg] + _qwen2_config(), + name="x", + dtype=torch.float16, + physical_capacity_positions=1, + capacity_source="", + scope=ModelLocalKVScope.REQUEST, + ) + + +def test_spec_from_hf_config_raises_rather_than_guessing(): + class Empty: + pass + + with pytest.raises(ValueError, match="layer count"): + spec_from_hf_config( + Empty(), + name="x", + dtype=torch.float16, + physical_capacity_positions=1, + capacity_source="", + scope=ModelLocalKVScope.REQUEST, + rows=RowDriver.FIXED, + rows_reason="test", + ) + + +# -------------------------------------------------------------------------- +# Collection +# -------------------------------------------------------------------------- + + +def _spec(**overrides) -> ModelLocalKVSpec: + kwargs = dict( + name="c", + layers=2, + kv_heads=2, + head_dim=4, + dtype=torch.float16, + physical_capacity_positions=8, + capacity_source="test", + scope=ModelLocalKVScope.REQUEST, + rows=RowDriver.FIXED, + rows_reason="test", + ) + kwargs.update(overrides) + return ModelLocalKVSpec(**kwargs) + + +class _Owner(torch.nn.Module): + def model_local_kv_specs(self): + return [_spec()] + + +class _CUDAGraphWrapperLike: + """Mirrors vllm.compilation.cuda_graph.CUDAGraphWrapper. + + Not an nn.Module, and its __getattr__ forwards everything the runnable + has. That forwarding is why a previous revision's "unwrap fix" was + unnecessary, and why a walk that does not unwrap can enter the tree twice. + """ + + def __init__(self, runnable): + self.runnable = runnable + + def unwrap(self): + return self.runnable + + def __getattr__(self, key): + return getattr(self.runnable, key) + + +def test_collector_finds_owners_nested_in_the_module_tree(): + class Mid(torch.nn.Module): + def __init__(self): + super().__init__() + self.inner = _Owner() + + class Top(torch.nn.Module): + def __init__(self): + super().__init__() + self.mid = Mid() + + model = Top() + assert [path for path, _ in collect_model_local_kv_specs(model)] == ["mid.inner"] + # 2 layers * 2 kv * 4 head_dim * 8 positions * 2 bytes * 2 (K and V) = 512 + assert _total_bytes(model) == 512 + + +def test_a_root_declarer_behind_a_wrapper_is_counted_once(): + """The wrapper forwards attributes, so a naive walk sees it twice.""" + + class Root(_Owner): + pass + + wrapped = _CUDAGraphWrapperLike(Root()) + collected = collect_model_local_kv_specs(wrapped) + + assert len(collected) == 1, collected + assert _total_bytes(wrapped) == 512 + + +def test_collector_reaches_owners_through_a_wrapper(): + class Top(torch.nn.Module): + def __init__(self): + super().__init__() + self.owner = _Owner() + + wrapped = _CUDAGraphWrapperLike(Top()) + assert [path for path, _ in collect_model_local_kv_specs(wrapped)] == ["owner"] + + +def test_collector_survives_an_owner_that_raises(): + class Broken(torch.nn.Module): + def model_local_kv_specs(self): + raise RuntimeError("checkpoint field missing") + + class Top(torch.nn.Module): + def __init__(self): + super().__init__() + self.broken = Broken() + self.ok = _Owner() + + assert _total_bytes(Top()) == 512 + + +def test_undeclared_models_report_zero_not_an_error(): + assert _total_bytes(object()) == 0 + assert _total_bytes(torch.nn.Linear(2, 2)) == 0 + + +def test_protocol_is_structural(): + assert isinstance(_Owner(), HasModelLocalKV) + assert not isinstance(torch.nn.Linear(2, 2), HasModelLocalKV) + + +@pytest.mark.parametrize( + ("relative_path", "owner"), + [ + ("qwen3_tts/segmented_graph_wrapper.py", "CUDAGraphDecoderWrapper"), + ("qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py", "Qwen3TTSTokenizerV2Decoder"), + ("mimo_audio/mimo_audio_llm.py", "MiMoAudioLLMForConditionalGeneration"), + ("minicpmo_4_5/minicpmo_4_5_omni_llm.py", "MiniCPMWhisperEncoder"), + ("ming_flash_omni/talker_module.py", "MingAudioGenerator"), + ("ming_flash_omni/ming_flash_omni_talker.py", "MingFlashOmniTalkerForConditionalGeneration"), + ], +) +def test_known_cache_owners_still_declare(relative_path: str, owner: str): + """Parsed rather than imported: these modules need CUDA and real weights. + + Weak by construction -- it only proves the method exists. The ming tests + above are the ones that check a declaration's contents. + """ + tree = ast.parse((MODELS_DIR / relative_path).read_text()) + classes = [node for node in ast.walk(tree) if isinstance(node, ast.ClassDef) and node.name == owner] + assert classes, f"{owner} not found in {relative_path}" + methods = {child.name for cls in classes for child in cls.body if isinstance(child, ast.FunctionDef)} + assert "model_local_kv_specs" in methods, f"{owner} no longer declares its model-local KV" diff --git a/tests/tools/test_check_tts_adapter.py b/tests/tools/test_check_tts_adapter.py index 47c6ac89fed..3107900c0eb 100644 --- a/tests/tools/test_check_tts_adapter.py +++ b/tests/tools/test_check_tts_adapter.py @@ -63,7 +63,15 @@ def test_gate_passes_on_current_tree(): assert main([]) == 0 -def test_budgets_are_the_real_counts(): +def test_budgets_are_not_exceeded(): + """Budgets are a ceiling, not a target. + + This used to assert equality, which re-imposed through pytest exactly what + #6008 removed from the checker: a PR that deleted branches failed CI unless + it also hand-edited the constant. That is what broke `main` for 5h19m in + #5746. The checker reports an under-budget count as a passing notice + (check_tts_adapter.py:140); the test must agree. + """ root = Path(__file__).resolve().parents[2] serving = root / "vllm_omni" / "entrypoints" / "openai" / "serving_speech.py" - assert len(_model_type_branches(ast.parse(serving.read_text()))) == MAX_MODEL_TYPE_BRANCHES + assert len(_model_type_branches(ast.parse(serving.read_text()))) <= MAX_MODEL_TYPE_BRANCHES diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py index bb01e522c6c..ec4faadb828 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py @@ -58,6 +58,12 @@ from vllm.utils.tensor_schema import TensorSchema from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import MiMoAudioConfig +from vllm_omni.model_executor.models.model_local_kv import ( + ModelLocalKVScope, + ModelLocalKVSpec, + RowDriver, + spec_from_hf_config, +) from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) @@ -744,6 +750,67 @@ def _parse_and_validate_audio_input(self, **kwargs: object) -> MimoAudioInputs | def get_language_model(self) -> torch.nn.Module: return self.model + def model_local_kv_specs(self) -> list[ModelLocalKVSpec]: + """Declare the local transformer's decode cache and its graph copies. + + The cache is created per call inside ``base_local_forward`` and is + bounded by the decode loop, not by sequence length: the loop runs + ``group_size + max(delay_pattern)`` steps regardless of how long the + utterance is. + + The second entry is the one that matters. ``base_local_forward`` runs + inside ``torch.cuda.graph`` during capture, so each captured bucket + leaves a cache's worth of tensors pinned in the graph pool for the + process lifetime. There is no Python object left to inspect -- the + cache is a local in the captured frame -- so the width comes from the + buckets that captured successfully, which is why it reads + ``local_forward_cg_by_bs`` rather than MIMO_CUDAGRAPH_BATCH_SIZES: a + bucket whose capture raised is logged and skipped. + + Both entries are declared, and summing them over-states a single step: + a given call either replays a captured bucket or runs eagerly, never + both. The graph pool is resident regardless, so the sum is the right + ceiling for the process and the wrong number for one request. + """ + delay_iters = self.group_size + max(self.delay_pattern) + dtype = next(self.local_transformer.parameters()).dtype + bound = f"group_size({self.group_size}) + max(delay_pattern)({max(self.delay_pattern)})" + specs = [ + spec_from_hf_config( + self.local_config, + name="local_transformer", + dtype=dtype, + physical_capacity_positions=delay_iters, + capacity_source=bound, + scope=ModelLocalKVScope.INVOCATION, + # One cache, B rows wide -- not B caches. base_local_forward + # builds a single DynamicCache whose batch dimension is the + # number of requests in the group. + rows=RowDriver.MAX_NUM_SEQS, + allocation_note="one batched allocation per eager call; graph replay uses the pool entry instead", + ) + ] + captured_rows = sum(self.local_forward_cg_by_bs) + if captured_rows: + specs.append( + spec_from_hf_config( + self.local_config, + name="local_transformer_graph_pool", + dtype=dtype, + physical_capacity_positions=delay_iters, + capacity_source=bound, + scope=ModelLocalKVScope.MODEL, + rows=RowDriver.FIXED, + rows_fixed=captured_rows, + rows_reason=f"sum of captured buckets {sorted(self.local_forward_cg_by_bs)}", + allocation_note=( + "pinned in the CUDA graph pool; capture is gated only on cuda.is_available(), " + "so this stays resident under enforce_eager too" + ), + ) + ) + return specs + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: if kwargs.get("modality_preprocess") is None: mm_dummy_embeddings = [] diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py index d34cce46542..f01e206881b 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py @@ -158,6 +158,17 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): patch_size=self.patch_size, ) + def model_local_kv_specs(self): + """Forward to the generator, which owns the talker's StaticCache. + + ``MingAudioGenerator`` is a plain object rather than a submodule, so + the module-tree walk cannot reach it on its own. + """ + generator = getattr(self, "audio_generator", None) + if generator is None: + return [] + return generator.model_local_kv_specs() + @property def device(self) -> torch.device: return next(self.model.parameters()).device diff --git a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py index 8e3e556117f..293ec2fb3a8 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py @@ -41,6 +41,12 @@ integrate_cfm_steps, ) from vllm_omni.model_executor.models.common.ming.dit_blocks import CondEmbedder, DiTBlock, FinalLayer +from vllm_omni.model_executor.models.model_local_kv import ( + ModelLocalKVScope, + ModelLocalKVSpec, + RowDriver, + spec_from_hf_config, +) logger = init_logger(__name__) @@ -575,6 +581,13 @@ class MingAudioGenerator: for a single TTS request. The generator is stateless across requests. """ + _STATIC_CACHE_LEN = 2048 + """Positions the talker's StaticCache is preallocated for. + + Read by both ``_init_kv_cache`` and ``model_local_kv_specs`` so the + declaration cannot drift from the allocation. + """ + def __init__( self, config, @@ -609,6 +622,49 @@ def __init__( # trailing latent frames prepended on each decode call self._vae_decode_pad_frames = 32 + def model_local_kv_specs(self) -> list[ModelLocalKVSpec]: + """Declare the talker's preallocated decode cache. + + This is the case that a static table cannot hold. ``self._llm_config`` + is a ``Qwen2Config`` resolved from the checkpoint at load time, so + layers, kv-heads and head-dim do not exist anywhere in this repo, and + the dtype is whatever the weights loaded as. Reading them off the same + objects ``_init_kv_cache`` builds from is the only way to get a real + number. + + Unlike the other three this is not a ceiling that a short utterance + stays under. transformers v5 initializes ``StaticLayer`` lazily, but + the first write to a layer allocates all ``max_cache_len`` positions at + once rather than growing, so a prefill brings the full extent into + existence on step 0 of every generation. + + Width is fixed at one row, not ``max_num_seqs``. ``forward`` takes only + ``runtime_additional_information[0]`` and runs the whole AR loop inline + before returning, so one worker holds one cache no matter how many + sequences the scheduler admits. Declaring this per-sequence -- which an + earlier revision did, by deriving the multiplier from scope -- reported + ``max_num_seqs`` times the real figure. + + The caveat is that nothing enforces the serialization. ``StaticCache`` + is built outside ``_sampler_pool``'s lock, so overlapping calls to + ``generate_latents`` would coexist. If this model ever gains a + concurrent entry point, this row count is the line to revisit. + """ + return [ + spec_from_hf_config( + self._llm_config, + name="talker_llm_static_cache", + dtype=next(self._model.parameters()).dtype, + physical_capacity_positions=self._STATIC_CACHE_LEN, + capacity_source=f"hardcoded max_cache_len={self._STATIC_CACHE_LEN} in _init_kv_cache", + scope=ModelLocalKVScope.INVOCATION, + rows=RowDriver.FIXED, + rows_fixed=1, + rows_reason="forward() handles one request and runs the AR loop inline", + allocation_note="full extent on first write per layer, not grown; rebuilt per text segment", + ) + ] + @cached_property def _sampler_pool(self) -> CFMGraphExecutorPool | None: device = next(self._model.parameters()).device @@ -790,7 +846,7 @@ def _init_his_lat( def _init_kv_cache( self, use_static_cache: bool, device: torch.device, dtype: torch.dtype ) -> tuple[StaticCache | None, int]: - max_cache_len = 2048 + max_cache_len = self._STATIC_CACHE_LEN if not use_static_cache: return None, max_cache_len cache = StaticCache( diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index 315ae159d3c..d70bc30d0ab 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -136,6 +136,13 @@ def _encode_tokens(tokenizer: Any, prompt: str) -> list[int]: from vllm.utils.tensor_schema import TensorSchema, TensorShape +from vllm_omni.model_executor.models.model_local_kv import ( + ModelLocalKVScope, + ModelLocalKVSpec, + RowDriver, + spec_from_hf_config, +) + logger = init_logger(__name__) hf_logger = logging.get_logger(__name__) @@ -2621,6 +2628,51 @@ def __init__(self, config: WhisperConfig): [MiniCPMWhisperEncoderLayer(config, layer_idx=i) for i in range(config.encoder_layers)] ) + def model_local_kv_specs(self) -> list[ModelLocalKVSpec]: + """Declare the streaming encoder's self-attention cache. + + Bounded by learned position embeddings, not by ``max_model_len``: the + cache cannot outgrow ``embed_positions``, and past that point the + forward pass repeats the last position rather than extending. + + Whisper is encoder-only self-attention here, so kv-heads equal + attention heads. Layer count is read off the built ``self.layers`` + rather than the config, so a partially built encoder reports what it + actually has. + + Declared per session, not per sequence. Each streaming session state + owns its own ``audio_past_key_values`` and the encoder runs at batch + size one, so ``max_num_seqs`` is the wrong number here. + + One session is the unit rather than the configured cap, because the cap + is not visible from the model and the setting that carries it cannot + distinguish "duplex off" from "one session" -- so a cap-scaled figure + would report a cache that a non-duplex deployment never allocates. An + earlier revision grew a capacity driver and an engine-side field to + carry that number for this one declarer; the arithmetic is a + multiplication the reader can do when they need the ceiling. + """ + return [ + spec_from_hf_config( + self.config, + name="whisper_encoder_self_attn", + dtype=self.conv1.weight.dtype, + layers=len(self.layers), + kv_heads=self.config.encoder_attention_heads, + head_dim=self.config.d_model // self.config.encoder_attention_heads, + physical_capacity_positions=int(self.embed_positions.weight.shape[0]), + capacity_source="embed_positions rows (max_source_positions)", + scope=ModelLocalKVScope.SESSION, + rows=RowDriver.FIXED, + rows_fixed=1, + rows_reason="one EncoderDecoderCache per streaming session, encoder runs at batch size 1", + allocation_note=( + "only allocated on the streaming path, where use_cache is set; scales with the " + "configured duplex session cap" + ), + ) + ] + def forward( self, input_features, diff --git a/vllm_omni/model_executor/models/model_local_kv.py b/vllm_omni/model_executor/models/model_local_kv.py new file mode 100644 index 00000000000..7328e996719 --- /dev/null +++ b/vllm_omni/model_executor/models/model_local_kv.py @@ -0,0 +1,347 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Declaration protocol for attention KV kept outside the paged manager. + +Several models hold their own attention KV as HuggingFace ``transformers`` cache +objects rather than using the engine's paged KV manager. That memory is +allocated after the profiling run that decides how much KV the engine may claim, +so no per-stage footprint number accounts for it, and nothing catches a cache +that silently stops working. + +This module lets a model *declare* what it holds. It describes; it never +allocates. Allocation for the diffusion path is owned by the DiT KV manager +(RFC #5244 / PR #6094) and this protocol must not grow into a second allocator. + +Why a runtime query rather than a static table: geometry is frequently not +knowable until weights are loaded. ``ming_flash_omni``'s talker builds its cache +from a ``Qwen2Config`` that comes from the checkpoint, so layers, kv-heads, +head-dim and dtype simply do not exist in this repo. A table would have a hole +in it; a post-load query does not. + +Why capacity is a resolved number rather than a formula: the four known caches +are each bounded by a different mechanism -- a sliding window, a decode-loop +trip count, an encoder frame limit, a hardcoded constant -- and not one of them +is ``max_model_len``. Only the model knows its own bound. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from enum import Enum +from typing import Protocol, runtime_checkable + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +__all__ = [ + "HasModelLocalKV", + "ModelLocalKVScope", + "ModelLocalKVSpec", + "RowDriver", + "collect_model_local_kv_specs", + "spec_from_hf_config", +] + + +class ModelLocalKVScope(str, Enum): + """How long one allocation lives. + + Diagnostic only. Lifetime and size are independent: a per-call cache can be + the widest thing a model owns, and a process-lifetime one can be a single + row. ``ModelLocalKVSpec.rows`` carries the size; this carries only the + reader's mental model of when the memory appears and goes. + """ + + INVOCATION = "invocation" + """Dies when the call returns (e.g. a per-step working copy).""" + + REQUEST = "request" + """Retained across steps of one request, released when it finishes.""" + + SESSION = "session" + """Outlives a request; belongs to a duplex/streaming session.""" + + MODEL = "model" + """Lives as long as the model (e.g. captured into a CUDA-graph pool).""" + + +class RowDriver(str, Enum): + """What sets how many rows of a cache are live at peak. + + Two values, because two are all the known caches need. A third belongs + here only when a model actually has a cache the engine widens by some + other number -- adding one speculatively costs a branch everywhere and + buys nothing. + """ + + FIXED = "fixed" + """A count the model controls: a captured-bucket sum, a serialized call path.""" + + MAX_NUM_SEQS = "max_num_seqs" + """One row per in-flight sequence. Only the engine knows the value.""" + + +@dataclass(frozen=True) +class ModelLocalKVSpec: + """One declared KV allocation. + + A model returns one entry per *distinct lifetime*, not per cache object: if + the same geometry exists as both a retained per-request cache and a + short-lived working copy, that is two entries. + """ + + name: str + layers: int + kv_heads: int + head_dim: int + dtype: torch.dtype + + physical_capacity_positions: int + """Positions the tensor can physically hold. + + Not the logical sequence position. A sliding-window cache truncates on every + write, so a stream of thousands of tokens may only ever occupy + ``sliding_window - 1`` slots. + """ + + capacity_source: str + """Free-text note on where the bound comes from, for diagnostics only. + + Never branch on this. It exists so a reader can tell "71 because sliding + window" from "2048 because someone typed 2048", which is otherwise + invisible. + """ + + scope: ModelLocalKVScope + """How long one allocation lives. Diagnostic only. + + Deliberately not the multiplier. An earlier revision derived scaling from + scope, which made every non-``MODEL`` cache multiply by ``max_num_seqs``. + That over-reported ming by that factor (its calls are serialized) and used + the wrong driver entirely for MiniCPM-o (duplex sessions, not sequences). + Lifetime does not imply replication topology; ``rows`` says what does. + """ + + rows: RowDriver + """What sets how many rows of this shape are live at peak. + + A row is one batch entry's worth of the geometry above. Deliberately does + not distinguish "one allocation N rows wide" from "N allocations of one + row": those cost the same, and an earlier revision that tried to carry both + got the object topology wrong in three of four declarers. Bytes are what + this protocol reports; describe the layout in ``allocation_note``. + + Required, because defaulting it silently understates every cache whose + multiplicity is not 1. + """ + + rows_fixed: int = 1 + """The count, when ``rows`` is ``FIXED``. Ignored otherwise.""" + + rows_reason: str = "" + """Why ``rows_fixed`` is that number. Required when ``rows`` is ``FIXED``. + + A bare constant cannot be reviewed. ``rows_fixed=1`` on a cache that looks + per-request is exactly the claim a reader has to be able to check. + """ + + allocation_note: str | None = None + """Optional detail such as "rebuilt per text segment". Diagnostic only.""" + + def __post_init__(self) -> None: + if self.rows is RowDriver.FIXED: + if self.rows_fixed < 1: + raise ValueError(f"rows_fixed must be >= 1, got {self.rows_fixed}") + if not self.rows_reason.strip(): + raise ValueError("a FIXED row count needs rows_reason; an unexplained constant cannot be reviewed") + + @property + def bytes_per_row(self) -> int: + """Bytes for a single batch row of this cache.""" + return ( + 2 # K and V + * self.layers + * self.kv_heads + * self.head_dim + * self.physical_capacity_positions + * torch.empty((), dtype=self.dtype).element_size() + ) + + def row_count(self, max_num_seqs: int = 1) -> int: + """Rows live at peak, given the engine's concurrency.""" + if self.rows is RowDriver.FIXED: + return self.rows_fixed + return max(1, max_num_seqs) + + def peak_bytes(self, max_num_seqs: int = 1) -> int: + """Peak bytes once the engine's concurrency is known. + + The model supplies geometry and names the driver; only the engine knows + the number behind it. + """ + return self.bytes_per_row * self.row_count(max_num_seqs) + + +def _resolve(config: object, candidates: Sequence[str], what: str) -> int: + for attr in candidates: + value = getattr(config, attr, None) + if value is not None: + return int(value) + raise ValueError( + f"Cannot determine {what} from {type(config).__name__}: tried " + f"{', '.join(candidates)}. Pass it explicitly to spec_from_hf_config()." + ) + + +def spec_from_hf_config( + config: object, + *, + name: str, + physical_capacity_positions: int, + capacity_source: str, + scope: ModelLocalKVScope, + dtype: torch.dtype, + rows: RowDriver, + rows_fixed: int = 1, + rows_reason: str = "", + allocation_note: str | None = None, + layers: int | None = None, + kv_heads: int | None = None, + head_dim: int | None = None, +) -> ModelLocalKVSpec: + """Build a spec from the same config object the cache is built from. + + Every known model-local cache is constructed from an HF config, so the + geometry half of a declaration is the same three lookups each time. Pass + ``config`` and this derives them; pass ``layers``/``kv_heads``/``head_dim`` + explicitly to override when the config uses encoder-decoder naming or when + the built module is a more truthful source than the config. + + Raises rather than guessing when an attribute is absent: a silently wrong + geometry would under-report, which is the failure this protocol exists to + prevent. ``rows`` is keyword-required for the same reason -- an earlier + revision defaulted the batch extent to 1 here while documenting it as + required on the dataclass, so every declarer that used this helper got the + default and the guarantee was decorative. + """ + if layers is None: + layers = _resolve(config, ("num_hidden_layers", "encoder_layers", "num_layers"), "layer count") + if kv_heads is None: + kv_heads = _resolve( + config, + ("num_key_value_heads", "num_attention_heads", "encoder_attention_heads"), + "kv head count", + ) + if head_dim is None: + explicit = getattr(config, "head_dim", None) + if explicit is not None: + head_dim = int(explicit) + else: + hidden = _resolve(config, ("hidden_size", "d_model"), "hidden size") + heads = _resolve(config, ("num_attention_heads", "encoder_attention_heads"), "attention head count") + head_dim = hidden // heads + + return ModelLocalKVSpec( + name=name, + layers=layers, + kv_heads=kv_heads, + head_dim=head_dim, + dtype=dtype, + physical_capacity_positions=physical_capacity_positions, + capacity_source=capacity_source, + scope=scope, + rows=rows, + rows_fixed=rows_fixed, + rows_reason=rows_reason, + allocation_note=allocation_note, + ) + + +@runtime_checkable +class HasModelLocalKV(Protocol): + """Implemented by whichever object owns the cache. + + Implement it on the owner, not on the registered model. The owner is + usually several levels down (a codec decoder, a talker backbone), and + ``collect_model_local_kv_specs`` finds it by walking the module tree, so + no intermediate class has to forward anything. + """ + + def model_local_kv_specs(self) -> Sequence[ModelLocalKVSpec]: ... + + +def _unwrap_to_module(model: object) -> object: + """Unwrap runner-side wrappers down to the real module tree. + + ``vllm.compilation.cuda_graph.CUDAGraphWrapper`` is a plain callable, not + an ``nn.Module``, but its ``__getattr__`` forwards anything the runnable + has -- including ``named_modules``. So a check for ``named_modules`` does + not detect the wrapper, and the wrapper and its runnable both answer to the + same attributes. That matters twice over: it means walking the wrapper + happens to work, and it means a naive walk can enter the tree as both the + wrapper and the root module and count a root-level declarer twice. + + Prefer the supported ``unwrap()`` accessor, which returns the runnable + directly. + """ + seen: set[int] = set() + current = model + while current is not None and id(current) not in seen: + seen.add(id(current)) + unwrap = getattr(current, "unwrap", None) + if not callable(unwrap): + return current + try: + nxt = unwrap() + except Exception: + return current + if nxt is None or nxt is current: + return current + current = nxt + return model + + +def collect_model_local_kv_specs(model: object) -> list[tuple[str, ModelLocalKVSpec]]: + """Collect every declaration in a loaded model, with the owner's path. + + Walks ``named_modules()`` because the cache owner is an inner module in all + known cases. Requiring each registered model to forward the call would put + the burden on classes that do not own a cache and would silently report + zero the moment someone forgot -- exactly the failure mode this is meant to + surface. + + A raising declaration is logged and skipped rather than propagated: + reporting memory must not be able to break model load. + """ + root = _unwrap_to_module(model) + + owners: list[tuple[str, object]] = [] + seen_ids: set[int] = set() + if isinstance(root, HasModelLocalKV): + owners.append(("", root)) + seen_ids.add(id(root)) + + named_modules = getattr(root, "named_modules", None) + if callable(named_modules): + for path, module in named_modules(): + if id(module) in seen_ids or not isinstance(module, HasModelLocalKV): + continue + seen_ids.add(id(module)) + owners.append((path, module)) + + collected: list[tuple[str, ModelLocalKVSpec]] = [] + for path, owner in owners: + try: + specs = owner.model_local_kv_specs() + except Exception: + logger.warning( + "model_local_kv_specs() raised on %s; skipping its declaration", + type(owner).__name__, + exc_info=True, + ) + continue + collected.extend((path, spec) for spec in specs) + return collected diff --git a/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py b/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py index a1eb071ee1b..3c73f5d13f0 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py +++ b/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py @@ -17,6 +17,13 @@ from vllm.logger import init_logger from vllm.platforms import current_platform +from vllm_omni.model_executor.models.model_local_kv import ( + ModelLocalKVScope, + ModelLocalKVSpec, + RowDriver, + spec_from_hf_config, +) + logger = init_logger(__name__) @@ -166,6 +173,85 @@ def _maybe_log_stats(self) -> None: if self._stats_log_every > 0 and self._stats_total_requests % self._stats_log_every == 0: self.log_decode_stats() + @staticmethod + def _retained_kv_caches(states: dict) -> list[tuple[int, DynamicCache]]: + """Find the KV caches a capture kept alive, with their batch rows. + + Walks the stored ``cache`` payload instead of assuming which capture + path populates it: ``_decode_icl_first_chunk`` fills its dict on the + decoder side, so the only reliable count is the one taken from the + objects themselves after warmup. + """ + found: list[tuple[int, DynamicCache]] = [] + for state in states.values(): + payload = state.get("cache") + if not isinstance(payload, dict): + continue + for value in payload.values(): + if not isinstance(value, DynamicCache): + continue + layers = getattr(value, "layers", None) + keys = getattr(layers[0], "keys", None) if layers else None + if keys is None: + continue + found.append((int(keys.shape[0]), value)) + return found + + def model_local_kv_specs(self) -> list[ModelLocalKVSpec]: + """Declare only the graph-resident copies this wrapper retains. + + The per-decode cache is declared by the decoder itself, because it + exists whether or not graphs were captured. What is specific to this + wrapper is that every captured shape keeps its dummy ``DynamicCache`` + alive for replay (``combined_states``, ``icl_prefix_states`` and + ``xvec_prefix_states`` all store ``cache``), for the process lifetime. + + Rows come from the state dicts rather than the configured capture + lists: a capture that raised is logged and skipped, so the configured + shape count would over-report. The trade is that a capture whose + tensors transformers has not yet materialized contributes nothing, so + this is a floor rather than a contract. + """ + config = self.decoder.config + try: + dtype = next(self.decoder.parameters()).dtype + except StopIteration: # parameterless stub, only reachable in tests + return [] + + # Physical, not logical: the sliding window truncates on every write, + # so a stream of thousands of frames never occupies more than this. + physical = self.prefix_length - 1 + retained = [ + *self._retained_kv_caches(self.combined_states), + *self._retained_kv_caches(self.icl_prefix_states), + *self._retained_kv_caches(self.xvec_prefix_states), + ] + batched_captures = sum(rows for rows, _ in retained) + + if not batched_captures: + return [] + return [ + spec_from_hf_config( + config, + name="codec_decoder_graph_pool", + dtype=dtype, + physical_capacity_positions=physical, + capacity_source=f"sliding_window({self.prefix_length}) - 1", + scope=ModelLocalKVScope.MODEL, + rows=RowDriver.FIXED, + rows_fixed=batched_captures, + rows_reason=( + f"rows across {len(retained)} retained caches in " + f"{len(self.combined_states)} suffix / {len(self.icl_prefix_states)} icl-prefix / " + f"{len(self.xvec_prefix_states)} xvec-prefix captures" + ), + allocation_note=( + "counted from caches transformers has materialized; captures whose tensors are still " + "unallocated report nothing, so this is a floor for the xvec path" + ), + ) + ] + def log_decode_stats(self) -> None: if not getattr(self, "_stats_enabled", False) or self._stats_total_requests == 0: return diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py index ea11e64a745..9572d9ee211 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py @@ -854,6 +854,75 @@ def __init__(self, config: Qwen3TTSTokenizerV2DecoderConfig): self._cudagraph_enabled = False self._cudagraph_wrapper = None + def model_local_kv_specs(self): + """Declare this decoder's attention KV. + + Only the async-chunk path holds KV. In stateless mode the decoder runs + ``_forward_exact``, which calls ``pre_transformer`` without + ``use_cache`` and allocates no cache at all, so declaring a per-request + cache there invents memory that does not exist. The previous revision + did exactly that and reported over a gigabyte for stateless + deployments; the one before it hung the declaration off the CUDA-graph + wrapper and so reported zero under ``enforce_eager``, where the eager + async-chunk path does allocate. Both mistakes come from tying this to + the wrong condition. + + Two entries when async-chunk is on: the retained per-request cache, and + the working copy ``_decode_icl_first_chunk`` deep-copies from it, which + is live at the same time. + + The wrapper, when present, adds the graph-resident copies. It is a + plain object rather than a submodule, so the module-tree walk cannot + reach it; this is the hop it needs. + """ + from vllm_omni.model_executor.models.model_local_kv import ( + ModelLocalKVScope, + RowDriver, + spec_from_hf_config, + ) + + config = self.config + try: + dtype = next(self.parameters()).dtype + except StopIteration: # parameterless stub, only reachable in tests + return [] + + window = int(getattr(config, "sliding_window", 0) or 0) + if window <= 1: + return [] + + specs = [] + wrapper = getattr(self, "_cudagraph_wrapper", None) + # async_chunk is the wrapper's record of the mode; without a wrapper the + # decoder has not been told, and only the stateless path is reachable. + if wrapper is not None and getattr(wrapper, "async_chunk", False): + common = dict( + config=config, + dtype=dtype, + physical_capacity_positions=window - 1, + capacity_source=f"sliding_window({window}) - 1", + rows=RowDriver.MAX_NUM_SEQS, + ) + specs.append( + spec_from_hf_config( + name="codec_decoder_request", + scope=ModelLocalKVScope.REQUEST, + allocation_note="retained across chunks of one request", + **common, + ) + ) + specs.append( + spec_from_hf_config( + name="codec_decoder_working_copy", + scope=ModelLocalKVScope.INVOCATION, + allocation_note="deep-copied from the retained cache each chunk; live alongside it", + **common, + ) + ) + if wrapper is not None: + specs.extend(wrapper.model_local_kv_specs()) + return specs + def precompute_snake_caches(self): """Precompute exp(alpha) and 1/(exp(beta)+eps) for all SnakeBeta modules.""" count = 0 diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 80d55020f25..8dd9f5f58b5 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -36,6 +36,7 @@ from vllm_omni.core.prefix_cache import OmniTensorPrefixCache from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.model_executor.layers.rotary_embedding.mrope import OmniMRotaryEmbedding as MRotaryEmbedding +from vllm_omni.model_executor.models.model_local_kv import collect_model_local_kv_specs from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.platforms import current_omni_platform @@ -194,6 +195,52 @@ def load_model(self, *args, **kwargs) -> None: self._maybe_enable_output_token_ids_for_model_sampler() self._init_talker_mtp() self._prewarm_attention_capture_workspaces() + self._report_model_local_kv() + + def _report_model_local_kv(self) -> None: + """Log attention KV this model holds outside the paged manager. + + Reporting only. Two of these caches are captured into CUDA graphs + during ``load_model``, so NVML already charges them to the process + before ``determine_available_memory()`` samples it; subtracting the + declared total from the KV budget would double-count them. What is not + charged is the per-request half, which appears after profiling and + scales with ``max_num_seqs``. + + Wrapped whole: a memory report has no business breaking model load, + and this runs on platforms these declarations have never been + exercised on. + """ + try: + self._log_model_local_kv() + except Exception: + logger.warning("Model-local KV reporting failed; continuing", exc_info=True) + + def _log_model_local_kv(self) -> None: + specs = collect_model_local_kv_specs(getattr(self, "model", None)) + if not specs: + return + max_num_seqs = max(1, int(self.scheduler_config.max_num_seqs)) + total = sum(spec.peak_bytes(max_num_seqs) for _, spec in specs) + logger.info( + "Model-local KV (outside the paged manager): %.2f MiB across %d declaration(s) at max_num_seqs=%d", + total / (1 << 20), + len(specs), + max_num_seqs, + ) + for path, spec in specs: + logger.info( + " %s.%s: %.2f MiB (%d rows from %s x %.2f MiB/row, scope=%s, %d positions from %s)", + path or type(self.model).__name__, + spec.name, + spec.peak_bytes(max_num_seqs) / (1 << 20), + spec.row_count(max_num_seqs), + spec.rows.value if spec.rows.value != "fixed" else f"fixed({spec.rows_fixed}; {spec.rows_reason})", + spec.bytes_per_row / (1 << 20), + spec.scope.value, + spec.physical_capacity_positions, + spec.capacity_source, + ) def _maybe_enable_output_token_ids_for_model_sampler(self) -> None: if getattr(self.model, "logitsprocs_need_output_token_ids", False): From 498c41b8d60270086c9a226c08a55ff6547800a3 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Thu, 10 Sep 2026 20:50:12 -0400 Subject: [PATCH 65/93] [Realtime] Emit current (non-beta) OpenAI audio/transcript event names (#7339) Signed-off-by: Nick Cao Co-authored-by: Claude Sonnet 5 --- docs/design/fullduplex-personaplex.md | 4 +- docs/design/fullduplex.md | 14 +- docs/serving/full_duplex_api.md | 2 +- docs/serving/realtime_duplex_api.md | 40 +++--- examples/online_serving/barge_in_client.py | 2 +- .../minicpmo/realtime_duplex_demo.py | 10 +- .../minicpmo/realtime_web/app/static/app.js | 8 +- tests/benchmarks/patch/test_patch.py | 4 +- tests/benchmarks/test_omniinteract.py | 6 +- tests/clients/test_duplex_client.py | 24 ++-- tests/e2e/features/fullduplex/test_client.py | 12 +- .../fullduplex/test_omni_duplex_eval_cli.py | 2 +- .../minicpmo_realtime_duplex_scenarios.py | 38 +++--- .../nemotron_voicechat_realtime_duplex.py | 8 +- .../personaplex_realtime_duplex.py | 6 +- .../run_minicpmo_realtime_duplex_demo_pair.py | 5 +- ...run_minicpmo_realtime_duplex_server_vad.py | 13 +- ...minicpmo_realtime_duplex_soft_interrupt.py | 7 +- .../test_minicpmo_realtime_duplex_drivers.py | 125 +++++++++++------- tests/engine/test_duplex_import_boundary.py | 4 +- .../duplex/test_websocket_actor.py | 6 +- .../openai/test_duplex_session_attachment.py | 2 +- .../openai_api/test_duplex_handler.py | 40 +++--- ...st_minicpmo_realtime_duplex_simple_demo.py | 8 +- .../test_minicpmo_realtime_web_static.py | 4 +- .../duplex/test_unified_runtime.py | 2 +- .../duplex/omni_duplex_eval_clock.py | 6 +- vllm_omni/benchmarks/omniinteract.py | 10 +- vllm_omni/clients/duplex.py | 6 +- .../entrypoints/duplex/realtime_output.py | 8 +- vllm_omni/entrypoints/duplex/warmup.py | 2 +- vllm_omni/entrypoints/duplex/websocket.py | 2 - .../entrypoints/openai/serving_speech.py | 2 +- vllm_omni/experimental/fullduplex/client.py | 8 +- 34 files changed, 251 insertions(+), 189 deletions(-) diff --git a/docs/design/fullduplex-personaplex.md b/docs/design/fullduplex-personaplex.md index 249554bbb92..33d125c9716 100644 --- a/docs/design/fullduplex-personaplex.md +++ b/docs/design/fullduplex-personaplex.md @@ -20,7 +20,7 @@ The target is the engine-native path: -> PersonaPlex Talker -> streaming PersonaPlex Code2Wav -> PersonaPlex data-plane projector - -> response.audio.delta + response.audio_transcript.delta + -> response.output_audio.delta + response.output_audio_transcript.delta ``` The standalone `/api/chat` and `/v1/audio/duplex` server that accompanied the @@ -197,7 +197,7 @@ model-neutral native results: It owns per-request audio and text cursors so a cumulative output cannot replay old audio. The generic projector turns each native result into the Realtime -`response.audio.delta` / `response.audio_transcript.delta` pair (and +`response.output_audio.delta` / `response.output_audio_transcript.delta` pair (and `response.output_text.delta` for text) under one `response_id`; the full mapping is the name map in the [Realtime Duplex API](../serving/realtime_duplex_api.md) serving guide. diff --git a/docs/design/fullduplex.md b/docs/design/fullduplex.md index 87fb1f79f30..7d3d5ec9572 100644 --- a/docs/design/fullduplex.md +++ b/docs/design/fullduplex.md @@ -606,11 +606,11 @@ scoped to one `response_id` unless stated otherwise. | Public event | Trigger and ordering | Cardinality | | --- | --- | --- | | `response.created` | First event for a visible assistant response; precedes its output-item, content-part, speak, text, and audio events. | Exactly once for every visible response. | -| `response.speak` | The model selected speak; emitted after `response.created` and no later than the first `response.audio.delta`. It carries decision metadata, not transcript text. | At most once. | -| `response.audio.delta` | Carries one ordered audio chunk. Every emitted audio delta is followed by a `response.audio_transcript.delta` for the same chunk. | Zero or more. | -| `response.audio_transcript.delta` | Append-only transcript contribution paired with an audio delta; it may be empty for a text-less audio unit and must not repeat or overlap earlier text. | One per audio delta. | -| `response.audio.done` | Closes the audio stream after its final delta and before response terminal events. It is omitted for a response with no audio or transcript. | At most once. | -| `response.audio_transcript.done` | Contains the exact concatenation of all non-overlapping transcript deltas for the response. | At most once, and only for a non-empty transcript. | +| `response.speak` | The model selected speak; emitted after `response.created` and no later than the first `response.output_audio.delta`. It carries decision metadata, not transcript text. | At most once. | +| `response.output_audio.delta` | Carries one ordered audio chunk. Every emitted audio delta is followed by a `response.output_audio_transcript.delta` for the same chunk. | Zero or more. | +| `response.output_audio_transcript.delta` | Append-only transcript contribution paired with an audio delta; it may be empty for a text-less audio unit and must not repeat or overlap earlier text. | One per audio delta. | +| `response.output_audio.done` | Closes the audio stream after its final delta and before response terminal events. It is omitted for a response with no audio or transcript. | At most once. | +| `response.output_audio_transcript.done` | Contains the exact concatenation of all non-overlapping transcript deltas for the response. | At most once, and only for a non-empty transcript. | | `response.output_item.done` / `conversation.item.done` | Finalize the assistant item after content-part terminal events and before `response.done`. | At most once each. | | `response.done` | Terminal event for a created response; follows audio, transcript, content-part, and item terminal events. | Exactly once for every created response that reaches a terminal state. | | `rate_limits.updated` | Compatibility event emitted immediately after `response.done`; the current payload has an empty rate-limit list. | Once per emitted `response.done`. | @@ -625,8 +625,8 @@ response and never reuse that response ID for later model turns. Compatibility changes must update this table and its protocol contract tests in the same change. The golden transcript test requires joined -`response.audio_transcript.delta` values to equal -`response.audio_transcript.done.transcript`; response tests separately enforce +`response.output_audio_transcript.delta` values to equal +`response.output_audio_transcript.done.transcript`; response tests separately enforce at-most-once `response.speak` and terminal event cardinality. ### Session internal ledgers diff --git a/docs/serving/full_duplex_api.md b/docs/serving/full_duplex_api.md index f648d6fdb3a..4e8bf974890 100644 --- a/docs/serving/full_duplex_api.md +++ b/docs/serving/full_duplex_api.md @@ -86,7 +86,7 @@ A typical `/v1/realtime?duplex=1` session follows this lifecycle: 3. Send `input_audio_buffer.append` events while microphone audio arrives. 4. Send `input_audio_buffer.commit` at a user-turn boundary when required by the model policy. -5. Consume `response.created`, transcript deltas, `response.audio.delta`, and +5. Consume `response.created`, transcript deltas, `response.output_audio.delta`, and `response.done` or `response.listen` events. 6. Send `playback.ack` after audio has been played when the session advertises playback acknowledgement support. diff --git a/docs/serving/realtime_duplex_api.md b/docs/serving/realtime_duplex_api.md index d37d0229164..c855c449139 100644 --- a/docs/serving/realtime_duplex_api.md +++ b/docs/serving/realtime_duplex_api.md @@ -311,13 +311,19 @@ work unmodified. The overlap falls into the three tiers used by the event catalo | Direction | Messages | | --- | --- | | Client → server | `session.update`, `input_audio_buffer.append` / `.commit` / `.clear`, `conversation.item.create` / `.retrieve` / `.truncate` / `.delete`, `response.create`, `response.cancel`, `output_audio_buffer.clear` (OpenAI: WebRTC-only; here also a WebSocket event) | -| Server → client | `error` (same `{type, code, message, event_id, param}` envelope and the same classes `invalid_request_error` / `server_error` / `rate_limit_error`), `session.created`, `session.updated`, `input_audio_buffer.committed` / `.cleared` / `.speech_started` / `.speech_stopped`, `conversation.item.created` / `.added` / `.done` / `.deleted` / `.retrieved` / `.truncated`, `conversation.item.input_audio_transcription.completed`, `response.created`, `response.done`, `response.output_item.added` / `.done`, `response.content_part.added` / `.done`, `response.audio.delta` / `.done`, `response.audio_transcript.delta` / `.done`, `response.output_text.delta` / `.done`, `response.function_call_arguments.delta` / `.done`, `rate_limits.updated`, `output_audio_buffer.cleared` | +| Server → client | `error` (same `{type, code, message, event_id, param}` envelope and the same classes `invalid_request_error` / `server_error` / `rate_limit_error`), `session.created`, `session.updated`, `input_audio_buffer.committed` / `.cleared` / `.speech_started` / `.speech_stopped`, `conversation.item.created` / `.added` / `.done` / `.deleted` / `.retrieved` / `.truncated`, `conversation.item.input_audio_transcription.completed`, `response.created`, `response.done`, `response.output_item.added` / `.done`, `response.content_part.added` / `.done`, `response.output_audio.delta` / `.done`, `response.output_audio_transcript.delta` / `.done`, `response.output_text.delta` / `.done`, `response.function_call_arguments.delta` / `.done`, `rate_limits.updated`, `output_audio_buffer.cleared` | | Session fields | `model`, `modalities` / `output_modalities`, `instructions`, `voice`, `input_audio_format` / `output_audio_format` (beta spelling) **and** `audio.input.format` / `audio.output.format` objects (GA spelling), `turn_detection.server_vad {threshold, prefix_padding_ms, silence_duration_ms}`, `input_audio_transcription`, `input_audio_noise_reduction`, `tools`, `tool_choice`, `temperature`, `max_response_output_tokens`, `speed`, `tracing` | | Object shapes | `realtime.session`, `realtime.item`, `realtime.response`; item content parts `input_text` / `input_audio` / `output_text` / `output_audio`; `function_call` and `function_call_output` items; `previous_item_id` chaining; `output_index` / `content_index` addressing | | Sequencing | `response.created → output_item.added → content_part.added → deltas → *.done → content_part.done → output_item.done → response.done → rate_limits.updated`; `speech_started → speech_stopped → committed` | -The server deliberately emits and accepts **both** the beta and the GA -spellings at once (`response.audio.delta` alongside `output_modalities`, +Audio/transcript output uses only the current (non-beta) OpenAI event names +(`response.output_audio.delta`, not the deprecated `response.audio.delta`), +matching `openai.types.realtime` in the official `openai` Python SDK and what +clients built against it (e.g. `livekit-plugins-openai`'s non-Azure code path) +expect; a client still on the old beta SDK types would not recognize these +events. Session fields and `conversation.item.*` events, in contrast, are +deliberately dual: the server emits and accepts **both** the beta and the GA +spellings at once (`output_modalities` alongside `modalities`, `conversation.item.added` alongside `conversation.item.created`) so either generation of OpenAI client parses the stream. @@ -329,7 +335,7 @@ A stock client ignores the extra keys; the extensions are additive. | --- | --- | | `session.created` / `session.updated` | top-level `incarnation`, `attachment_generation`, `resume_token`; inside `session`: `state`, `turn_state`, `epoch`, `turn_id`, `active_request_id`, `active_response_id`, `active_response_turn_id`, `overlap_policy`, `overlap_*_ms/rms`, `playback_commit_policy`, `playback`, `capabilities`, `ref_audio`, `extra_body`, `idle_timeout_s`, `response_format` | | `response.created` / `response.done` / `response.listen` | `response_id` at top level; the raw duplex event under `response.metadata` (`duplex_event`); `status_details.reason` uses vLLM reasons (`barge_in`, `client_cancelled`, `new_response`, …) | -| `response.audio.delta` | `format`, `sample_rate_hz`, `metadata{session_id, epoch, model_speak, end_of_turn, audio_duration_ms, audio_text_marks, playback}` | +| `response.output_audio.delta` | `format`, `sample_rate_hz`, `metadata{session_id, epoch, model_speak, end_of_turn, audio_duration_ms, audio_text_marks, playback}` | | `response.speak` (inserted before the first delta) | not an OpenAI event, but rides the OpenAI response envelope (`response_id`, `item_id`, `output_index`, `content_index`) | | `input_audio_buffer.append` | `is_speech`, `video_frames` (+ `max_slice_nums`), `duration_ms`, `audio_end_ms`, per-event `format` / `sample_rate_hz` | | `input_audio_buffer.commit` | `final`, `response_create`, `is_speech:false` (silence declaration) | @@ -443,12 +449,12 @@ table lives in the normative contract section of [Full-Duplex Runtime (MiniCPM-o | `response.output_item.added` | 1 | Assistant item (or `function_call` item) attached to the response. | | `response.content_part.added` | 1 | Audio/text content part opened on the assistant item. | | `response.speak` | 3 | Model chose to speak (native lane); at most once per response, before the first audio delta. | -| `response.audio.delta` | 2 | One ordered audio chunk. OpenAI name; adds `format`, `sample_rate_hz`, `metadata{epoch, model_speak, end_of_turn, audio_text_marks, playback}`. | -| `response.audio_transcript.delta` | 1 | Transcript text paired one-to-one with each audio delta. | +| `response.output_audio.delta` | 2 | One ordered audio chunk. OpenAI name; adds `format`, `sample_rate_hz`, `metadata{epoch, model_speak, end_of_turn, audio_text_marks, playback}`. | +| `response.output_audio_transcript.delta` | 1 | Transcript text paired one-to-one with each audio delta. | | `response.output_text.delta` | 1 | Text delta for text-modality responses. | | `response.output_text.done` | 1 | Final text of a text-modality response. | -| `response.audio.done` | 1 | Audio stream closed for the response. | -| `response.audio_transcript.done` | 1 | Full transcript (concatenation of the deltas). | +| `response.output_audio.done` | 1 | Audio stream closed for the response. | +| `response.output_audio_transcript.done` | 1 | Full transcript (concatenation of the deltas). | | `response.content_part.done` | 1 | Content part finalized with its transcript/text. | | `response.output_item.done` | 1 | Assistant item finalized. | | `response.done` | 2 | Terminal event; `status` ∈ `completed` \| `cancelled` \| `failed`. OpenAI name; vLLM-Omni `status_details.reason` values and raw duplex event under `metadata`. | @@ -902,11 +908,11 @@ deferred commit during an active response (`event.response_create_deferred`) } ``` -`S→C response.audio.delta` (one ~1 s unit, 24 kHz PCM16) +`S→C response.output_audio.delta` (one ~1 s unit, 24 kHz PCM16) ```json { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp_01", "item_id": "item_resp_01", "output_index": 0, @@ -926,10 +932,10 @@ deferred commit during an active response (`event.response_create_deferred`) } ``` -`S→C response.audio_transcript.delta` (exactly one per audio delta) +`S→C response.output_audio_transcript.delta` (exactly one per audio delta) ```json -{"type": "response.audio_transcript.delta", "response_id": "resp_01", "item_id": "item_resp_01", "output_index": 0, "content_index": 0, "delta": "It is sunny"} +{"type": "response.output_audio_transcript.delta", "response_id": "resp_01", "item_id": "item_resp_01", "output_index": 0, "content_index": 0, "delta": "It is sunny"} ``` `S→C response.output_text.delta` / `response.output_text.done` (text modality) @@ -942,14 +948,14 @@ deferred commit during an active response (`event.response_create_deferred`) {"type": "response.output_text.done", "response_id": "resp_02", "item_id": "item_resp_02", "output_index": 0, "content_index": 0, "text": "Tomorrow looks cloudy."} ``` -`S→C response.audio.done` / `response.audio_transcript.done` +`S→C response.output_audio.done` / `response.output_audio_transcript.done` ```json -{"type": "response.audio.done", "response_id": "resp_01", "item_id": "item_resp_01", "output_index": 0, "content_index": 0} +{"type": "response.output_audio.done", "response_id": "resp_01", "item_id": "item_resp_01", "output_index": 0, "content_index": 0} ``` ```json -{"type": "response.audio_transcript.done", "response_id": "resp_01", "item_id": "item_resp_01", "output_index": 0, "content_index": 0, "transcript": "It is sunny, 24 degrees."} +{"type": "response.output_audio_transcript.done", "response_id": "resp_01", "item_id": "item_resp_01", "output_index": 0, "content_index": 0, "transcript": "It is sunny, 24 degrees."} ``` `S→C response.content_part.done` @@ -1191,8 +1197,8 @@ Realtime projector renames and fans out these events before they reach a | `session.create` (`open_session`, `session.config`) | first `session.update` | | `input.committed` | `conversation.item.added` / `conversation.item.created`, `input_audio_buffer.committed`, `conversation.item.input_audio_transcription.completed`, `conversation.item.done` | | `input.cancelled` | `input_audio_buffer.cleared` | -| `response.output_audio.delta` (with `audio_transcript`) | optional `response.speak`, then `response.audio.delta` and `response.audio_transcript.delta` | -| `response.output_audio.done` / `response.output_text.done` | `response.audio.done` / `response.output_text.done` | +| `response.output_audio.delta` (with `audio_transcript`) | optional `response.speak`, then `response.output_audio.delta` and `response.output_audio_transcript.delta` | +| `response.output_audio.done` / `response.output_text.done` | `response.output_audio.done` / `response.output_text.done` | | `response.text.delta` | `response.output_text.delta` | | `response.message` (chat-fallback raw chunk) | passed through | | `audio.cancelled` | optional `output_audio_buffer.cleared`, then the cancelled terminal events and `response.done` with `status: "cancelled"` | diff --git a/examples/online_serving/barge_in_client.py b/examples/online_serving/barge_in_client.py index 1ac627d355e..b144da9f40c 100644 --- a/examples/online_serving/barge_in_client.py +++ b/examples/online_serving/barge_in_client.py @@ -65,7 +65,7 @@ ) _TRANSCRIPT_DELTA_EVENTS = { - "response.audio_transcript.delta", + "response.output_audio_transcript.delta", "response.output_text.delta", "response.text.delta", } diff --git a/examples/online_serving/minicpmo/realtime_duplex_demo.py b/examples/online_serving/minicpmo/realtime_duplex_demo.py index 9cbc7f29bc9..9fc75621e6d 100644 --- a/examples/online_serving/minicpmo/realtime_duplex_demo.py +++ b/examples/online_serving/minicpmo/realtime_duplex_demo.py @@ -63,14 +63,14 @@ def __init__(self, output_dir: Path) -> None: def handle(self, event: dict[str, object]) -> None: event_type = event.get("type") if event_type in { - "response.audio_transcript.delta", + "response.output_audio_transcript.delta", "response.output_text.delta", }: delta = event.get("delta") if isinstance(delta, str) and delta: print(delta, end="", file=sys.stderr, flush=True) return - if event_type != "response.audio.delta": + if event_type != "response.output_audio.delta": return delta = event.get("delta") or event.get("audio") @@ -509,12 +509,12 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]: audio = collector.audio_bytes() first_text_at_s = collector.first_received_at( - "response.audio_transcript.delta", + "response.output_audio_transcript.delta", "response.output_text.delta", after_s=commit_sent_at_s, ) first_audio_at_s = collector.first_received_at( - "response.audio.delta", + "response.output_audio.delta", after_s=commit_sent_at_s, ) response_created_at_s = collector.first_received_at( @@ -536,7 +536,7 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]: for event in collector.events if event.get("type") in { - "response.audio_transcript.delta", + "response.output_audio_transcript.delta", "response.output_text.delta", } ] diff --git a/examples/online_serving/minicpmo/realtime_web/app/static/app.js b/examples/online_serving/minicpmo/realtime_web/app/static/app.js index 2943eec743e..1ea2ba09343 100644 --- a/examples/online_serving/minicpmo/realtime_web/app/static/app.js +++ b/examples/online_serving/minicpmo/realtime_web/app/static/app.js @@ -356,7 +356,7 @@ case 'response.speak': beginAssistant(responseId); break; - case 'response.audio.delta': + case 'response.output_audio.delta': currentResponseId = responseId || currentResponseId; assistantActive = true; setModel('Speaking'); @@ -364,13 +364,13 @@ .then((decoded) => feedPlayback(decoded, responseId)) .catch((error) => appendLog(`audio decode failed: ${error.message || error}`, true)); break; - case 'response.audio.done': + case 'response.output_audio.done': requestPlaybackDrain(responseId); break; - case 'response.audio_transcript.delta': + case 'response.output_audio_transcript.delta': addTranscript('assistant', event.delta || ''); break; - case 'response.audio_transcript.done': + case 'response.output_audio_transcript.done': finishTranscript('assistant', event.transcript || ''); break; case 'conversation.item.input_audio_transcription.delta': diff --git a/tests/benchmarks/patch/test_patch.py b/tests/benchmarks/patch/test_patch.py index accb3010077..6dd35e67750 100644 --- a/tests/benchmarks/patch/test_patch.py +++ b/tests/benchmarks/patch/test_patch.py @@ -168,7 +168,7 @@ async def send(self, event): ) self.events.add( { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": response_id, "delta": f"turn {self.response_count}", }, @@ -176,7 +176,7 @@ async def send(self, event): ) self.events.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": response_id, "delta": base64.b64encode(b"\x00\x00" * 2400).decode(), "sample_rate_hz": 24_000, diff --git a/tests/benchmarks/test_omniinteract.py b/tests/benchmarks/test_omniinteract.py index f9488b11cb8..b9dd10870b3 100644 --- a/tests/benchmarks/test_omniinteract.py +++ b/tests/benchmarks/test_omniinteract.py @@ -65,7 +65,7 @@ def _audio( value: int = 1, ) -> dict[str, object]: event: dict[str, object] = { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": response_id, "format": "pcm16", "delta": base64.b64encode(bytes((value, 0)) * samples).decode(), @@ -76,7 +76,7 @@ def _audio( def _text(response_id: str = "r1", value: str = "hello") -> dict[str, object]: - return {"type": "response.audio_transcript.delta", "response_id": response_id, "delta": value} + return {"type": "response.output_audio_transcript.delta", "response_id": response_id, "delta": value} def _listen(*, buffering: bool = False) -> dict[str, object]: @@ -427,7 +427,7 @@ def test_artifacts_publish_official_bundle_and_sparse_deferred_state(tmp_path: P assert result.success and result.eligible_for_official_eval assert summary["scene_type"] == "multi_turn" assert context and context.spans - assert all("delta" not in event for event in context.events if event["type"] == "response.audio.delta") + assert all("delta" not in event for event in context.events if event["type"] == "response.output_audio.delta") oi.publish_deferred_case_artifacts(tmp_path / "out", case, result) directory = oi._output_dir(tmp_path / "out", case) assert all((directory / name).is_file() for name in oi.SUCCESS_ARTIFACTS) diff --git a/tests/clients/test_duplex_client.py b/tests/clients/test_duplex_client.py index 1be3eb9432e..ce1795e0f85 100644 --- a/tests/clients/test_duplex_client.py +++ b/tests/clients/test_duplex_client.py @@ -279,7 +279,7 @@ async def test_response_handle_flow(): sock.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3}) sock.feed( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-1", "delta": _b64(chunk), "sample_rate_hz": 24_000, @@ -288,7 +288,7 @@ async def test_response_handle_flow(): ) sock.feed( { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": "resp-1", "delta": "hi there", "server_event_seq": 5, @@ -370,7 +370,7 @@ async def test_terminal_listen_keeps_spoken_decision(): sock.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3}) sock.feed( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-1", "delta": _b64(chunk), "sample_rate_hz": 24_000, @@ -418,7 +418,7 @@ async def test_slow_audio_consumer_drops_oldest_instead_of_stalling(): ) for payload in (b"c1", b"c2", b"c3"): # _feed must never block the reader, even against a full queue. - handle._feed(AudioDelta({"type": "response.audio.delta", "delta": _b64(payload)})) + handle._feed(AudioDelta({"type": "response.output_audio.delta", "delta": _b64(payload)})) handle._finish(None) chunks = [chunk async for chunk in handle.audio()] assert chunks == [b"c3"] # oldest chunks were dropped, the sentinel landed @@ -561,7 +561,7 @@ def test_event_collector_accumulates_audio(): collector = EventCollector() collector.add({"type": "response.created", "response": {"id": "r1"}}, received_at_s=1.0) collector.add( - {"type": "response.audio.delta", "response_id": "r1", "delta": _b64(b"ab"), "sample_rate_hz": 16_000}, + {"type": "response.output_audio.delta", "response_id": "r1", "delta": _b64(b"ab"), "sample_rate_hz": 16_000}, received_at_s=1.1, ) collector.add( @@ -632,7 +632,7 @@ def test_event_collector_partitions_audio_by_response(): collector.add({"type": "response.created", "response": {"id": "resp-a"}}) collector.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-a", "delta": base64.b64encode(b"audio-a").decode("ascii"), "sample_rate_hz": 16_000, @@ -643,7 +643,7 @@ def test_event_collector_partitions_audio_by_response(): assert collector.audio_bytes("resp-a") == b"audio-a" assert collector.output_sample_rate_hz == 16_000 assert collector.first_received_at("response.created") is not None - assert collector.last_received_at("response.audio.delta") is not None + assert collector.last_received_at("response.output_audio.delta") is not None def test_event_collector_reports_engine_token_and_audio_intervals(): @@ -664,7 +664,7 @@ def test_event_collector_reports_engine_token_and_audio_intervals(): for received_at_s, cumulative_audio_ms in ((10.2, 80), (10.25, 160), (10.36, 240)): collector.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-a", "delta": base64.b64encode(b"audio").decode("ascii"), "sample_rate_hz": 16_000, @@ -676,11 +676,11 @@ def test_event_collector_reports_engine_token_and_audio_intervals(): received_at_s=received_at_s, ) collector.add( - {"type": "response.audio_transcript.delta", "response_id": "resp-a", "delta": ""}, + {"type": "response.output_audio_transcript.delta", "response_id": "resp-a", "delta": ""}, received_at_s=10.1, ) collector.add( - {"type": "response.audio_transcript.delta", "response_id": "resp-a", "delta": "hello"}, + {"type": "response.output_audio_transcript.delta", "response_id": "resp-a", "delta": "hello"}, received_at_s=10.15, ) collector.add( @@ -752,7 +752,7 @@ def test_response_timing_ignores_unowned_session_level_metrics(): ) collector.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-a", "delta": base64.b64encode(b"audio").decode("ascii"), "metadata": { @@ -887,7 +887,7 @@ def test_reference_audio_data_url(tmp_path): def test_event_collector_response_text_joins_deltas_per_response(): collector = EventCollector() collector.add({"type": "response.created", "response_id": "r1"}) - collector.add({"type": "response.audio_transcript.delta", "response_id": "r1", "delta": "he"}) + collector.add({"type": "response.output_audio_transcript.delta", "response_id": "r1", "delta": "he"}) collector.add({"type": "response.output_text.delta", "response_id": "r2", "delta": "other"}) collector.add({"type": "response.text.delta", "response_id": "r1", "delta": "llo"}) assert collector.response_text("r1") == "hello" diff --git a/tests/e2e/features/fullduplex/test_client.py b/tests/e2e/features/fullduplex/test_client.py index 266df1fa04a..250391fe21d 100644 --- a/tests/e2e/features/fullduplex/test_client.py +++ b/tests/e2e/features/fullduplex/test_client.py @@ -241,7 +241,7 @@ def test_realtime_event_collector_partitions_audio_by_response(): collector.add({"type": "response.created", "response": {"id": "resp-a"}}) collector.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-a", "delta": base64.b64encode(b"audio-a").decode("ascii"), "sample_rate_hz": 16_000, @@ -252,7 +252,7 @@ def test_realtime_event_collector_partitions_audio_by_response(): assert collector.audio_bytes("resp-a") == b"audio-a" assert collector.output_sample_rate_hz == 16_000 assert collector.first_received_at("response.created") is not None - assert collector.last_received_at("response.audio.delta") is not None + assert collector.last_received_at("response.output_audio.delta") is not None def test_realtime_event_collector_reports_engine_token_and_audio_intervals(): @@ -273,7 +273,7 @@ def test_realtime_event_collector_reports_engine_token_and_audio_intervals(): for received_at_s, cumulative_audio_ms in ((10.2, 80), (10.25, 160), (10.36, 240)): collector.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-a", "delta": base64.b64encode(b"audio").decode("ascii"), "sample_rate_hz": 16_000, @@ -286,7 +286,7 @@ def test_realtime_event_collector_reports_engine_token_and_audio_intervals(): ) collector.add( { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": "resp-a", "delta": "", }, @@ -294,7 +294,7 @@ def test_realtime_event_collector_reports_engine_token_and_audio_intervals(): ) collector.add( { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": "resp-a", "delta": "hello", }, @@ -369,7 +369,7 @@ def test_response_timing_ignores_unowned_session_level_metrics(): ) collector.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-a", "delta": base64.b64encode(b"audio").decode("ascii"), "metadata": { diff --git a/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py b/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py index 192859de5e5..a4d7faf9b4b 100644 --- a/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py +++ b/tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py @@ -119,7 +119,7 @@ async def handler(websocket): await websocket.send( json.dumps( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "r1", "delta": base64.b64encode(b"\0\0" * 240).decode(), "sample_rate_hz": 24_000, diff --git a/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py b/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py index ebd59fc3e2c..f5bba11e431 100644 --- a/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py +++ b/tests/e2e/online_serving/helpers/minicpmo_realtime_duplex_scenarios.py @@ -213,7 +213,7 @@ def add(self, event: dict[str, object], *, received_at_s: float | None = None) - item_id = item.get("id") if isinstance(item_id, str) and item_id not in self.assistant_item_ids: self.assistant_item_ids.append(item_id) - elif event_type == "response.audio.delta": + elif event_type == "response.output_audio.delta": delta = event.get("delta") or event.get("audio") if isinstance(delta, str) and delta: try: @@ -386,8 +386,8 @@ def first_response_lifecycle_indices(self) -> dict[str, int]: "response.output_item.added", "response.content_part.added", "response.speak", - "response.audio.delta", - "response.audio.done", + "response.output_audio.delta", + "response.output_audio.done", "response.content_part.done", "response.output_item.done", "response.done", @@ -402,8 +402,8 @@ def first_response_lifecycle_indices(self) -> dict[str, int]: ) if index is None and event_type not in { "response.speak", - "response.audio.delta", - "response.audio.done", + "response.output_audio.delta", + "response.output_audio.done", }: return {} if index is not None: @@ -451,13 +451,13 @@ def event_order_ok( "response.output_item.added", "response.content_part.added", "response.speak", - "response.audio.delta", - "response.audio.done", + "response.output_audio.delta", + "response.output_audio.done", "response.content_part.done", "response.output_item.done", "response.done", ] - if "response.audio.delta" not in indices_by_type: + if "response.output_audio.delta" not in indices_by_type: listen_index = self.first_index( "response.listen", lambda event: ( @@ -515,7 +515,7 @@ def model_policy_event_order_ok( def model_speak_before_audio_ok(self) -> bool: speak_index = self.first_index("response.speak") - audio_index = self.first_index("response.audio.delta") + audio_index = self.first_index("response.output_audio.delta") return speak_index is not None and audio_index is not None and speak_index < audio_index def response_done(self, response_id: str | None) -> bool: @@ -532,7 +532,7 @@ def response_audio_delta_count(self, response_id: str | None) -> int: return sum( 1 for event in self.events - if event.get("type") == "response.audio.delta" and self._event_response_id(event) == response_id + if event.get("type") == "response.output_audio.delta" and self._event_response_id(event) == response_id ) def response_playback_sent_ms(self, response_id: str | None) -> int: @@ -565,14 +565,16 @@ def response_transcript_delta(self, response_id: str) -> str: return "".join( str(event.get("delta", "")) for event in self.events - if event.get("type") == "response.audio_transcript.delta" and self._event_response_id(event) == response_id + if event.get("type") == "response.output_audio_transcript.delta" + and self._event_response_id(event) == response_id ) def response_transcript_done(self, response_id: str) -> list[str]: return [ str(event.get("transcript", "")) for event in self.events - if event.get("type") == "response.audio_transcript.done" and self._event_response_id(event) == response_id + if event.get("type") == "response.output_audio_transcript.done" + and self._event_response_id(event) == response_id ] def completed_response_ids(self) -> list[str]: @@ -603,7 +605,7 @@ def stale_audio_delta_count(self) -> int: return 0 stale = 0 for index, event in enumerate(self.events): - if event.get("type") != "response.audio.delta": + if event.get("type") != "response.output_audio.delta": continue metadata = event.get("metadata") if not isinstance(metadata, dict): @@ -1126,7 +1128,7 @@ def eligible_response_id() -> str | None: state, lambda: state.response_audio_delta_count(response_id) > 0, timeout_s=timeout_s, - label=f"{transcript} response.audio.delta", + label=f"{transcript} response.output_audio.delta", ) await _wait_for( state, @@ -1341,7 +1343,7 @@ async def continue_until_first_speak() -> bool: state, lambda: state.response_audio_delta_count(first_response_id) > 0, timeout_s=timeout_s, - label=f"{transcripts[0]} response.audio.delta", + label=f"{transcripts[0]} response.output_audio.delta", ) if state.response_done(first_response_id): raise RuntimeError("first response completed before overlap input could start") @@ -1600,7 +1602,9 @@ async def run_demo(args: DemoArgs) -> dict[str, object]: transcript_hints_enabled=transcript_hints_enabled and not continuous_input, ) model_speak_event_ok = state.model_speak_before_audio_ok() - realtime_audio_lifecycle_ok = state.count("response.audio.delta") > 0 and state.count("response.audio.done") > 0 + realtime_audio_lifecycle_ok = ( + state.count("response.output_audio.delta") > 0 and state.count("response.output_audio.done") > 0 + ) completed_response_ids = state.completed_response_ids() observed_turn_response_ids = [response_id for response_id in turn_response_ids if isinstance(response_id, str)] expected_empty_response_ids = { @@ -1625,7 +1629,7 @@ async def run_demo(args: DemoArgs) -> dict[str, object]: state.count("response.created") == state.count("response.done") and len(state.response_ids) == len(completed_response_ids) and len(completed_response_ids) == len(set(completed_response_ids)) - and state.count("response.audio.done") <= state.count("response.done") + and state.count("response.output_audio.done") <= state.count("response.done") and _response_cardinality_ok( completed_response_ids, expected_turns=expected_turns, diff --git a/tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py b/tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py index 11c35cecc10..6aba7cfe718 100644 --- a/tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py +++ b/tests/e2e/online_serving/nemotron_voicechat_realtime_duplex.py @@ -215,7 +215,7 @@ async def run(args: argparse.Namespace) -> dict[str, object]: client.events.count("response.function_call_arguments.done") > 0 if args.expect_function_call else ( - client.events.count("response.audio.delta") >= args.minimum_audio_chunks + client.events.count("response.output_audio.delta") >= args.minimum_audio_chunks and ( args.allow_incomplete_response or client.events.count("response.done") > completed_responses_at_commit @@ -279,10 +279,10 @@ def tool_result_completed() -> bool: transcript = "".join( str(event.get("delta", "")) for event in later - if event.get("type") == "response.audio_transcript.delta" + if event.get("type") == "response.output_audio_transcript.delta" ).lower() return ( - any(event.get("type") == "response.audio.delta" for event in later) + any(event.get("type") == "response.output_audio.delta" for event in later) and any(event.get("type") == "response.done" for event in later) and (args.expected_post_tool_text is None or args.expected_post_tool_text.lower() in transcript) ) @@ -297,7 +297,7 @@ def tool_result_completed() -> bool: await asyncio.sleep(args.drain_s) audio = client.events.audio_bytes() - audio_events = _events(client, "response.audio.delta") + audio_events = _events(client, "response.output_audio.delta") rates = {event.get("sample_rate_hz") for event in audio_events} if not args.expect_function_call and audio and rates != {OUTPUT_SAMPLE_RATE_HZ}: raise AssertionError(f"unexpected output sample rates: {rates}") diff --git a/tests/e2e/online_serving/personaplex_realtime_duplex.py b/tests/e2e/online_serving/personaplex_realtime_duplex.py index 26ee22037f2..20da1598d27 100644 --- a/tests/e2e/online_serving/personaplex_realtime_duplex.py +++ b/tests/e2e/online_serving/personaplex_realtime_duplex.py @@ -244,7 +244,7 @@ def _audio_frame_stats( def _response_ids(client: RawRealtimeProbe) -> set[str]: return { str(event["response_id"]) - for event in _events(client, "response.audio.delta") + for event in _events(client, "response.output_audio.delta") if isinstance(event.get("response_id"), str) } @@ -268,7 +268,7 @@ def _session_result( raise AssertionError(f"output audio is non-finite or silent: samples={pcm.size}, rms={rms}") rates = { int(event["sample_rate_hz"]) - for event in _events(client, "response.audio.delta") + for event in _events(client, "response.output_audio.delta") if isinstance(event.get("sample_rate_hz"), int) } if rates != {SAMPLE_RATE_HZ}: @@ -289,7 +289,7 @@ def _session_result( ) runtime_metadata = [ metadata["vllm_omni"] - for event in _events(client, "response.audio.delta") + for event in _events(client, "response.output_audio.delta") if isinstance((metadata := event.get("metadata")), dict) and isinstance(metadata.get("vllm_omni"), dict) ] if not any( diff --git a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py index f5539968bc4..b5d04097bfe 100644 --- a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py +++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_demo_pair.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Run two MiniCPM-o Realtime duplex demo processes concurrently. This E2E driver covers the same process boundary users exercise manually: @@ -20,7 +23,7 @@ REPO_ROOT = Path(__file__).resolve().parents[3] DEMO_PATH = REPO_ROOT / "examples/online_serving/minicpmo/realtime_duplex_demo.py" -AUDIO_DELTA_EVENTS = {"response.audio.delta", "response.output_audio.delta"} +AUDIO_DELTA_EVENTS = {"response.output_audio.delta"} def _canonical_path(path: str) -> Path: diff --git a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py index 4047a0901bb..0fb1c57b325 100644 --- a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py +++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_server_vad.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import base64 from pathlib import Path @@ -41,7 +44,10 @@ async def until(predicate, label: str) -> None: next(e for e in client.events.events if e.get("type") == "response.created") ) assert target_id - await until(lambda: _events(client.events.events, target_id, "response.audio.delta"), "response.audio.delta") + await until( + lambda: _events(client.events.events, target_id, "response.output_audio.delta"), + "response.output_audio.delta", + ) cursor = len(client.events.events) await client.stream_pcm16(interrupt + bytes(16_000 * 2 * 800 // 1000), chunk_ms=args.chunk_ms, realtime=True) await until(lambda: _events(client.events.events, target_id, "response.done"), "cancelled response.done") @@ -64,10 +70,11 @@ async def until(predicate, label: str) -> None: terminal = done[0] if len(done) == 1 else None trailing = events[events.index(terminal) + 1 :] if terminal is not None else events stale = any( - e.get("type") == "response.audio.delta" and RealtimeEventCollector.response_id(e) == target_id for e in trailing + e.get("type") == "response.output_audio.delta" and RealtimeEventCollector.response_id(e) == target_id + for e in trailing ) followup_audio = any( - e.get("type") == "response.audio.delta" and RealtimeEventCollector.response_id(e) != target_id + e.get("type") == "response.output_audio.delta" and RealtimeEventCollector.response_id(e) != target_id for e in events[cursor:] ) ok = ( diff --git a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py index 3aff0cf8d0e..5e0cfff3a00 100644 --- a/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py +++ b/tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Validate MiniCPM-o Realtime duplex soft-interrupt delta streaming. This E2E driver runs the public ``realtime_duplex_demo.py`` against a live @@ -21,9 +24,9 @@ REPO_ROOT = Path(__file__).resolve().parents[3] DEMO_PATH = REPO_ROOT / "examples/online_serving/minicpmo/realtime_duplex_demo.py" -AUDIO_DELTA_EVENTS = {"response.audio.delta", "response.output_audio.delta"} +AUDIO_DELTA_EVENTS = {"response.output_audio.delta"} TRANSCRIPT_DELTA_EVENTS = { - "response.audio_transcript.delta", + "response.output_audio_transcript.delta", "response.output_text.delta", } diff --git a/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py index 22c8e6b34a6..d45c535db44 100644 --- a/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py +++ b/tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py @@ -168,9 +168,9 @@ def test_realtime_duplex_demo_pair_launches_demo_processes_concurrently(tmp_path "events = [", " {'type': 'session.created'},", " {'type': 'response.created', 'response': {'id': response_id}},", - " {'type': 'response.audio.delta', 'response_id': response_id, " + " {'type': 'response.output_audio.delta', 'response_id': response_id, " "'delta': 'AAAA', 'sample_rate_hz': 24000, '_client_received_at_s': 10.1},", - " {'type': 'response.audio.delta', 'response_id': response_id, " + " {'type': 'response.output_audio.delta', 'response_id': response_id, " "'delta': 'AAAA', 'sample_rate_hz': 24000, '_client_received_at_s': 10.6},", " {'type': 'response.done', 'response_id': response_id},", " {'type': 'session.closed'},", @@ -237,8 +237,8 @@ def test_realtime_duplex_demo_pair_rejects_false_green_audio_contract(tmp_path): response_id = "resp-a" events = [ {"type": "response.created", "response": {"id": response_id}}, - {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "sample_rate_hz": 24000}, - {"type": "response.audio.delta", "response_id": "resp-other", "delta": "AAAA", "sample_rate_hz": 24000}, + {"type": "response.output_audio.delta", "response_id": response_id, "delta": "AAAA", "sample_rate_hz": 24000}, + {"type": "response.output_audio.delta", "response_id": "resp-other", "delta": "AAAA", "sample_rate_hz": 24000}, {"type": "response.done", "response_id": response_id}, ] (output / "events.jsonl").write_text( @@ -459,25 +459,25 @@ def test_realtime_duplex_soft_interrupt_accepts_multi_delta_handoff_sequence(tmp {"type": "response.created", "response": {"id": first_response_id}, "_client_received_at_s": 2.0}, {"type": "response.speak", "response_id": first_response_id, "_client_received_at_s": 2.0}, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": first_response_id, "delta": "AAAA", "_client_received_at_s": 2.1, }, { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": first_response_id, "delta": "中国古代四大发明", "_client_received_at_s": 2.1, }, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": first_response_id, "delta": "AAAA", "_client_received_at_s": 2.6, }, { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": first_response_id, "delta": "是造纸术。", "_client_received_at_s": 2.6, @@ -486,25 +486,25 @@ def test_realtime_duplex_soft_interrupt_accepts_multi_delta_handoff_sequence(tmp {"type": "response.created", "response": {"id": second_response_id}, "_client_received_at_s": 4.0}, {"type": "response.speak", "response_id": second_response_id, "_client_received_at_s": 4.0}, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": second_response_id, "delta": "AAAA", "_client_received_at_s": 4.1, }, { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": second_response_id, "delta": "一加一等于", "_client_received_at_s": 4.1, }, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": second_response_id, "delta": "AAAA", "_client_received_at_s": 4.6, }, { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": second_response_id, "delta": "二。", "_client_received_at_s": 4.6, @@ -561,25 +561,25 @@ def test_realtime_duplex_soft_interrupt_accepts_followup_done_after_commit(tmp_p {"type": "response.created", "response": {"id": first_response_id}, "_client_received_at_s": 2.0}, {"type": "response.speak", "response_id": first_response_id, "_client_received_at_s": 2.0}, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": first_response_id, "delta": "AAAA", "_client_received_at_s": 2.1, }, { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": first_response_id, "delta": "中国古代四大发明", "_client_received_at_s": 2.1, }, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": first_response_id, "delta": "AAAA", "_client_received_at_s": 2.6, }, { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": first_response_id, "delta": "是造纸术。", "_client_received_at_s": 2.6, @@ -589,25 +589,25 @@ def test_realtime_duplex_soft_interrupt_accepts_followup_done_after_commit(tmp_p {"type": "response.created", "response": {"id": second_response_id}, "_client_received_at_s": 3.5}, {"type": "response.speak", "response_id": second_response_id, "_client_received_at_s": 3.5}, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": second_response_id, "delta": "AAAA", "_client_received_at_s": 3.6, }, { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": second_response_id, "delta": "一加一等于", "_client_received_at_s": 3.6, }, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": second_response_id, "delta": "AAAA", "_client_received_at_s": 3.7, }, { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": second_response_id, "delta": "二。", "_client_received_at_s": 3.7, @@ -653,8 +653,18 @@ def test_realtime_duplex_soft_interrupt_model_policy_accepts_single_response(tmp events = [ {"type": "response.listen", "_client_received_at_s": 1.0}, {"type": "response.created", "response": {"id": response_id}, "_client_received_at_s": 2.0}, - {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "_client_received_at_s": 2.1}, - {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "_client_received_at_s": 2.15}, + { + "type": "response.output_audio.delta", + "response_id": response_id, + "delta": "AAAA", + "_client_received_at_s": 2.1, + }, + { + "type": "response.output_audio.delta", + "response_id": response_id, + "delta": "AAAA", + "_client_received_at_s": 2.15, + }, {"type": "response.done", "response_id": response_id, "_client_received_at_s": 2.2}, {"type": "response.listen", "_client_received_at_s": 2.3}, {"type": "input_audio_buffer.committed", "_client_received_at_s": 3.0}, @@ -692,10 +702,25 @@ def test_realtime_duplex_soft_interrupt_model_policy_accepts_commit_during_speak events = [ {"type": "response.listen", "_client_received_at_s": 1.0}, {"type": "response.created", "response": {"id": response_id}, "_client_received_at_s": 2.0}, - {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "_client_received_at_s": 2.1}, - {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "_client_received_at_s": 2.2}, + { + "type": "response.output_audio.delta", + "response_id": response_id, + "delta": "AAAA", + "_client_received_at_s": 2.1, + }, + { + "type": "response.output_audio.delta", + "response_id": response_id, + "delta": "AAAA", + "_client_received_at_s": 2.2, + }, {"type": "input_audio_buffer.committed", "_client_received_at_s": 3.0}, - {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "_client_received_at_s": 4.0}, + { + "type": "response.output_audio.delta", + "response_id": response_id, + "delta": "AAAA", + "_client_received_at_s": 4.0, + }, {"type": "response.done", "response_id": response_id, "_client_received_at_s": 5.0}, ] (output / "events.jsonl").write_text( @@ -729,8 +754,18 @@ def test_realtime_duplex_soft_interrupt_response_required_rejects_single_respons events = [ {"type": "response.listen", "_client_received_at_s": 1.0}, {"type": "response.created", "response": {"id": response_id}, "_client_received_at_s": 2.0}, - {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "_client_received_at_s": 2.1}, - {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA", "_client_received_at_s": 2.15}, + { + "type": "response.output_audio.delta", + "response_id": response_id, + "delta": "AAAA", + "_client_received_at_s": 2.1, + }, + { + "type": "response.output_audio.delta", + "response_id": response_id, + "delta": "AAAA", + "_client_received_at_s": 2.15, + }, {"type": "response.done", "response_id": response_id, "_client_received_at_s": 2.2}, {"type": "response.listen", "_client_received_at_s": 2.3}, {"type": "input_audio_buffer.committed", "_client_received_at_s": 3.0}, @@ -767,19 +802,19 @@ def test_realtime_duplex_soft_interrupt_reports_text_expectation_without_gating( {"type": "response.listen", "_client_received_at_s": 1.0}, {"type": "response.created", "response": {"id": first_response_id}, "_client_received_at_s": 2.0}, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": first_response_id, "delta": "AAAA", "_client_received_at_s": 2.1, }, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": first_response_id, "delta": "AAAA", "_client_received_at_s": 2.2, }, { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": first_response_id, "delta": "一加一等于二", "_client_received_at_s": 2.2, @@ -788,19 +823,19 @@ def test_realtime_duplex_soft_interrupt_reports_text_expectation_without_gating( {"type": "response.listen", "_client_received_at_s": 2.4}, {"type": "response.created", "response": {"id": second_response_id}, "_client_received_at_s": 3.0}, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": second_response_id, "delta": "AAAA", "_client_received_at_s": 3.1, }, { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": second_response_id, "delta": "AAAA", "_client_received_at_s": 3.2, }, { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": second_response_id, "delta": "不知道", "_client_received_at_s": 3.2, @@ -834,7 +869,7 @@ def test_realtime_duplex_soft_interrupt_reports_text_expectation_without_gating( events = [ event for event in events - if event.get("response_id") != second_response_id or event["type"] != "response.audio_transcript.delta" + if event.get("response_id") != second_response_id or event["type"] != "response.output_audio_transcript.delta" ] (output / "events.jsonl").write_text( "".join(demo.json.dumps(event) + "\n" for event in events), @@ -1123,7 +1158,7 @@ def test_realtime_duplex_demo_partitions_timing_by_response_identity(): ) state.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": response_id, "delta": base64.b64encode(b"audio").decode("ascii"), "metadata": { @@ -1142,7 +1177,7 @@ def test_realtime_duplex_demo_partitions_timing_by_response_identity(): ) state.add( { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": response_id, "delta": "audio", }, @@ -1251,8 +1286,8 @@ def test_realtime_duplex_demo_distinguishes_late_and_missing_commit(): {"type": "response.output_item.added", "response_id": response_id}, {"type": "response.content_part.added", "response_id": response_id}, {"type": "response.speak", "response_id": response_id}, - {"type": "response.audio.delta", "response_id": response_id, "delta": "AAAA"}, - {"type": "response.audio.done", "response_id": response_id}, + {"type": "response.output_audio.delta", "response_id": response_id, "delta": "AAAA"}, + {"type": "response.output_audio.done", "response_id": response_id}, {"type": "response.content_part.done", "response_id": response_id}, {"type": "response.output_item.done", "response_id": response_id}, {"type": "response.done", "response_id": response_id}, @@ -1314,7 +1349,7 @@ def test_realtime_duplex_demo_distinct_inputs_compare_audio_content(): def _add_response_transcript(state, response_id, *, transcript, audio=True): state.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": response_id, "delta": "YQ==" if audio else "", } @@ -1322,14 +1357,14 @@ def _add_response_transcript(state, response_id, *, transcript, audio=True): if transcript: state.add( { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": response_id, "delta": transcript, } ) state.add( { - "type": "response.audio_transcript.done", + "type": "response.output_audio_transcript.done", "response_id": response_id, "transcript": transcript, } @@ -1404,7 +1439,7 @@ def test_realtime_duplex_demo_speak_gate_rejects_text_channel(): { "type": "response.speak", "response_id": "resp-1", - "text": "must be emitted through response.audio_transcript.delta", + "text": "must be emitted through response.output_audio_transcript.delta", } ) @@ -1450,7 +1485,7 @@ def test_realtime_duplex_demo_gate_rejects_delta_done_mismatch(): demo = _load_demo_module() state = demo.DemoState() _add_response_transcript(state, "resp-1", transcript="delta文本") - done = next(event for event in state.events if event.get("type") == "response.audio_transcript.done") + done = next(event for event in state.events if event.get("type") == "response.output_audio_transcript.done") done["transcript"] = "另一个done文本" result = demo._evaluate_transcript_integrity( @@ -2180,7 +2215,7 @@ async def fake_send_pcm16(*args, **kwargs): state.add({"type": "response.created", "response": {"id": "resp-first"}}) state.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-first", "delta": "YQ==", } @@ -2290,7 +2325,7 @@ def test_realtime_duplex_demo_writes_audio_per_response(tmp_path): state.add({"type": "response.created", "response": {"id": response_id}}) state.add( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": response_id, "delta": demo.base64.b64encode(payload).decode(), "sample_rate_hz": 24000, diff --git a/tests/engine/test_duplex_import_boundary.py b/tests/engine/test_duplex_import_boundary.py index f30266d282b..013ab86617d 100644 --- a/tests/engine/test_duplex_import_boundary.py +++ b/tests/engine/test_duplex_import_boundary.py @@ -122,8 +122,8 @@ def test_client_package_stays_clear_of_the_server_import_graph() -> None: collector = EventCollector() chunk = base64.b64encode(b"\\x00\\x00" * 2400).decode("ascii") collector.add({"type": "response.created", "response_id": "r1"}, received_at_s=1.0) -collector.add({"type": "response.audio.delta", "response_id": "r1", "delta": chunk}, received_at_s=1.1) -collector.add({"type": "response.audio.delta", "response_id": "r1", "delta": chunk}, received_at_s=1.2) +collector.add({"type": "response.output_audio.delta", "response_id": "r1", "delta": chunk}, received_at_s=1.1) +collector.add({"type": "response.output_audio.delta", "response_id": "r1", "delta": chunk}, received_at_s=1.2) summary = collector.timing_summary(after_s=0.0, response_id="r1") if "audio_output" not in summary or "request_metrics" not in summary: raise SystemExit("timing_summary produced no audio_output/request_metrics sections") diff --git a/tests/entrypoints/duplex/test_websocket_actor.py b/tests/entrypoints/duplex/test_websocket_actor.py index 9bdb73667bf..4e5d151b9b1 100644 --- a/tests/entrypoints/duplex/test_websocket_actor.py +++ b/tests/entrypoints/duplex/test_websocket_actor.py @@ -122,12 +122,12 @@ async def test_stale_fence_payload_is_dropped_before_websocket_send(): actor = DuplexWebSocketActor(websocket, current_epoch=lambda: 2) writer = asyncio.create_task(actor.writer_loop()) - await actor.send_json({"type": "response.audio.delta", "epoch": 1}) - await actor.send_json({"type": "response.audio.delta", "epoch": 2}) + await actor.send_json({"type": "response.output_audio.delta", "epoch": 1}) + await actor.send_json({"type": "response.output_audio.delta", "epoch": 2}) await actor.close_writer() await writer - assert websocket.sent == [{"type": "response.audio.delta", "epoch": 2}] + assert websocket.sent == [{"type": "response.output_audio.delta", "epoch": 2}] assert actor.stale_output_dropped == 1 diff --git a/tests/entrypoints/openai/test_duplex_session_attachment.py b/tests/entrypoints/openai/test_duplex_session_attachment.py index b4ea9c3287d..a73159064dd 100644 --- a/tests/entrypoints/openai/test_duplex_session_attachment.py +++ b/tests/entrypoints/openai/test_duplex_session_attachment.py @@ -78,7 +78,7 @@ def test_event_journal_sequences_acknowledges_and_replays_exact_payloads() -> No clock = _Clock() journal = DuplexEventJournal(max_bytes=4096, ttl_s=60.0, clock=clock) - first = journal.record({"type": "response.audio.delta", "delta": "AAAA"}) + first = journal.record({"type": "response.output_audio.delta", "delta": "AAAA"}) second = journal.record({"type": "response.done", "response_id": "resp-1"}) assert first.sequence == 1 diff --git a/tests/entrypoints/openai_api/test_duplex_handler.py b/tests/entrypoints/openai_api/test_duplex_handler.py index d3bbc5d01d2..0f969845792 100644 --- a/tests/entrypoints/openai_api/test_duplex_handler.py +++ b/tests/entrypoints/openai_api/test_duplex_handler.py @@ -2197,9 +2197,11 @@ def test_native_realtime_protocol_audio_delta_preserves_sample_rate_hz(): ) audio_events = [ - payload for payload in payloads if payload["type"] in {"response.output_audio.delta", "response.audio.delta"} + payload + for payload in payloads + if payload["type"] in {"response.output_audio.delta", "response.output_audio.delta"} ] - assert {payload["type"] for payload in audio_events} == {"response.audio.delta"} + assert {payload["type"] for payload in audio_events} == {"response.output_audio.delta"} assert {payload["format"] for payload in audio_events} == {"pcm16"} assert {payload["sample_rate_hz"] for payload in audio_events} == {24000} @@ -2227,14 +2229,14 @@ def test_native_realtime_protocol_ignores_removed_legacy_event_switches(): ) event_types = {payload["type"] for payload in payloads} - assert "response.audio.delta" in event_types - assert "response.audio.done" in event_types - assert "response.audio_transcript.delta" in event_types - assert "response.audio_transcript.done" in event_types - assert "response.output_audio.delta" not in event_types - assert "response.output_audio.done" not in event_types - assert "response.output_audio_transcript.delta" not in event_types - assert "response.output_audio_transcript.done" not in event_types + assert "response.output_audio.delta" in event_types + assert "response.output_audio.done" in event_types + assert "response.output_audio_transcript.delta" in event_types + assert "response.output_audio_transcript.done" in event_types + assert "response.audio.delta" not in event_types + assert "response.audio.done" not in event_types + assert "response.audio_transcript.delta" not in event_types + assert "response.audio_transcript.done" not in event_types assert "response.output_item.created" not in event_types assert "response.text.delta" not in event_types @@ -2256,9 +2258,9 @@ def test_native_realtime_protocol_emits_terminal_audio_transcript_events(): ) by_type = {payload["type"]: payload for payload in payloads} - assert "response.audio.done" in by_type - assert "response.audio_transcript.delta" in by_type - assert "response.audio_transcript.done" in by_type + assert "response.output_audio.done" in by_type + assert "response.output_audio_transcript.delta" in by_type + assert "response.output_audio_transcript.done" in by_type assert by_type["response.content_part.done"]["part"]["transcript"] == "hello" assert by_type["response.output_item.done"]["item"]["object"] == "realtime.item" assert by_type["response.done"]["response"]["output"][0]["object"] == "realtime.item" @@ -2284,8 +2286,10 @@ def test_native_realtime_protocol_terminal_transcript_equals_joined_deltas(): ) ) - transcript = "".join(str(event["delta"]) for event in events if event["type"] == "response.audio_transcript.delta") - done = [str(event["transcript"]) for event in events if event["type"] == "response.audio_transcript.done"] + transcript = "".join( + str(event["delta"]) for event in events if event["type"] == "response.output_audio_transcript.delta" + ) + done = [str(event["transcript"]) for event in events if event["type"] == "response.output_audio_transcript.done"] assert transcript == "It's Canberra. Next question." assert done == [transcript] @@ -2374,7 +2378,7 @@ def test_native_realtime_protocol_response_done_emits_audio_done_before_terminal payloads = protocol._from_duplex_event({"type": "response.done", "response_id": "resp-done"}) event_types = [payload["type"] for payload in payloads] - assert event_types.index("response.audio.done") < event_types.index("response.content_part.done") + assert event_types.index("response.output_audio.done") < event_types.index("response.content_part.done") assert event_types.index("response.content_part.done") < event_types.index("response.output_item.done") assert event_types.index("conversation.item.done") < event_types.index("response.done") assert event_types.index("response.output_item.done") < event_types.index("response.done") @@ -3181,7 +3185,7 @@ def _stage_output(samples: int, *, turn_end: bool = False): def commit_overlap_after_first_response(ws: TimedWebSocket, data: dict[str, Any]) -> None: nonlocal overlap_sent - if data.get("type") != "response.audio.delta" or overlap_sent: + if data.get("type") != "response.output_audio.delta" or overlap_sent: return overlap_sent = True chunk = { @@ -8679,7 +8683,7 @@ async def emit(output: object) -> None: protocol = NativeRealtimeSessionProtocol(TimedWebSocket()) # type: ignore[arg-type] realtime_events = [event for payload in sent for event in protocol._from_duplex_event(payload)] transcript_deltas = [ - m for m in realtime_events if m.get("type") == "response.audio_transcript.delta" and m.get("delta") + m for m in realtime_events if m.get("type") == "response.output_audio_transcript.delta" and m.get("delta") ] assert len(created) == 2 assert len(done) == 2 diff --git a/tests/examples/test_minicpmo_realtime_duplex_simple_demo.py b/tests/examples/test_minicpmo_realtime_duplex_simple_demo.py index d4a70e1b5ed..ba17ec65692 100644 --- a/tests/examples/test_minicpmo_realtime_duplex_simple_demo.py +++ b/tests/examples/test_minicpmo_realtime_duplex_simple_demo.py @@ -60,7 +60,7 @@ def test_post_commit_decision_accepts_drained_speak_after_commit_ack(): events = [ {"type": "input_audio_buffer.committed"}, {"type": "response.created", "response": {"id": "resp-a"}}, - {"type": "response.audio.delta", "response_id": "resp-a", "delta": "AAAA"}, + {"type": "response.output_audio.delta", "response_id": "resp-a", "delta": "AAAA"}, {"type": "response.done", "response_id": "resp-a"}, ] @@ -267,7 +267,7 @@ def test_streaming_output_writer_persists_audio_deltas_as_they_arrive(tmp_path, writer.handle( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-a", "delta": base64.b64encode(first_pcm).decode("ascii"), "sample_rate_hz": 24_000, @@ -275,14 +275,14 @@ def test_streaming_output_writer_persists_audio_deltas_as_they_arrive(tmp_path, ) writer.handle( { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": "resp-a", "delta": "你好", } ) writer.handle( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "resp-a", "delta": base64.b64encode(second_pcm).decode("ascii"), "sample_rate_hz": 24_000, diff --git a/tests/examples/test_minicpmo_realtime_web_static.py b/tests/examples/test_minicpmo_realtime_web_static.py index 5fdab87d16d..926ae36ef70 100644 --- a/tests/examples/test_minicpmo_realtime_web_static.py +++ b/tests/examples/test_minicpmo_realtime_web_static.py @@ -42,8 +42,8 @@ def test_client_uses_proxy_relative_realtime_url_and_model_policy_session(): assert "input_audio_buffer.commit" not in source assert "playback.ack" in source assert "event.event || event" in source - assert "response.audio.delta" in source - assert "response.audio_transcript.delta" in source + assert "response.output_audio.delta" in source + assert "response.output_audio_transcript.delta" in source assert "conversation.item.input_audio_transcription" in source assert "force_barge_in" not in source assert "server_vad" not in source diff --git a/tests/model_executor/models/personaplex/duplex/test_unified_runtime.py b/tests/model_executor/models/personaplex/duplex/test_unified_runtime.py index 8d4c1b5893f..1f03e894b7e 100644 --- a/tests/model_executor/models/personaplex/duplex/test_unified_runtime.py +++ b/tests/model_executor/models/personaplex/duplex/test_unified_runtime.py @@ -61,7 +61,7 @@ def _audio_client( frames[:voiced_frames] = 1000 raw = frames.tobytes() event = { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": "response-1", "sample_rate_hz": e2e_driver.SAMPLE_RATE_HZ, "metadata": { diff --git a/vllm_omni/benchmarks/duplex/omni_duplex_eval_clock.py b/vllm_omni/benchmarks/duplex/omni_duplex_eval_clock.py index 03997fc8e30..494f9a856d2 100644 --- a/vllm_omni/benchmarks/duplex/omni_duplex_eval_clock.py +++ b/vllm_omni/benchmarks/duplex/omni_duplex_eval_clock.py @@ -49,7 +49,11 @@ def extract_timed_sentences(events: Iterable[dict[str, Any]], *, clock: str = "m last: float | None = None for event in events: event_type = event.get("type") - if event_type not in {"response.output_text.delta", "response.audio_transcript.delta", "response.text.delta"}: + if event_type not in { + "response.output_text.delta", + "response.output_audio_transcript.delta", + "response.text.delta", + }: continue delta = event.get("delta") if not isinstance(delta, str) or not delta.strip(): diff --git a/vllm_omni/benchmarks/omniinteract.py b/vllm_omni/benchmarks/omniinteract.py index b821d02356d..6b861d3399a 100644 --- a/vllm_omni/benchmarks/omniinteract.py +++ b/vllm_omni/benchmarks/omniinteract.py @@ -233,7 +233,7 @@ def ingest(self, events: EventCollector) -> None: if response_id: self.completed.add(response_id) continue - if event.get("type") != "response.audio.delta": + if event.get("type") != "response.output_audio.delta": continue encoded = event.get("delta") or event.get("audio") if not response_id: @@ -241,13 +241,13 @@ def ingest(self, events: EventCollector) -> None: if not isinstance(encoded, str): raise ValueError("response audio payload is missing") if event.get("format") is None: - self._warn_once("response.audio.delta omitted format; assumed pcm16") + self._warn_once("response.output_audio.delta omitted format; assumed pcm16") elif event.get("format") != "pcm16": raise ValueError("OmniInteract output must be pcm16") rate = event.get("sample_rate_hz") if rate is None: rate = events.output_sample_rate_hz or OUTPUT_SAMPLE_RATE - self._warn_once(f"response.audio.delta omitted sample_rate_hz; assumed {rate}") + self._warn_once(f"response.output_audio.delta omitted sample_rate_hz; assumed {rate}") if rate != OUTPUT_SAMPLE_RATE: raise ValueError(f"OmniInteract output must use {OUTPUT_SAMPLE_RATE} Hz audio") try: @@ -683,7 +683,7 @@ def _collect_output( timing = response_times.setdefault(response_id, [start_s, end_s]) timing[0], timing[1] = min(timing[0], start_s), max(timing[1], end_s) - text_event_types = {"response.audio_transcript.delta", "response.output_text.delta", "response.text.delta"} + text_event_types = {"response.output_audio_transcript.delta", "response.output_text.delta", "response.text.delta"} text_ids = {collector.response_id(event) for event in collector.events if event.get("type") in text_event_types} if None in text_ids or not text_ids <= created: raise ValueError("response transcript has no matching response.created") @@ -726,7 +726,7 @@ def _collect_output( **{key: value for key, value in event.items() if key not in {"delta", "audio"}}, "audio_bytes": audio_sizes[index], } - if event.get("type") == "response.audio.delta" + if event.get("type") == "response.output_audio.delta" else dict(event) for index, event in enumerate(collector.events) ] diff --git a/vllm_omni/clients/duplex.py b/vllm_omni/clients/duplex.py index eb8a6ded0d7..15a191cc15f 100644 --- a/vllm_omni/clients/duplex.py +++ b/vllm_omni/clients/duplex.py @@ -460,9 +460,7 @@ def to_exception(self) -> DuplexProtocolError: "session.expired": SessionExpired, "response.created": ResponseCreated, "response.done": ResponseDone, - "response.audio.delta": AudioDelta, "response.output_audio.delta": AudioDelta, - "response.audio_transcript.delta": TranscriptDelta, "response.output_audio_transcript.delta": TranscriptDelta, "response.text.delta": TextDelta, "response.output_text.delta": TextDelta, @@ -1359,7 +1357,7 @@ def response_text(self, response_id: str) -> str: if self.response_id(event) == response_id and event.get("type") in { - "response.audio_transcript.delta", + "response.output_audio_transcript.delta", "response.output_text.delta", "response.text.delta", } @@ -1406,7 +1404,7 @@ def timing_summary( if ( event.get("type") in { - "response.audio_transcript.delta", + "response.output_audio_transcript.delta", "response.output_text.delta", "response.text.delta", } diff --git a/vllm_omni/entrypoints/duplex/realtime_output.py b/vllm_omni/entrypoints/duplex/realtime_output.py index 82ce9761e76..fcbdf87d7da 100644 --- a/vllm_omni/entrypoints/duplex/realtime_output.py +++ b/vllm_omni/entrypoints/duplex/realtime_output.py @@ -168,7 +168,7 @@ def _from_duplex_event(self, event: dict[str, Any]) -> list[dict[str, object]]: text = "" payloads.append( { - "type": "response.audio_transcript.delta", + "type": "response.output_audio_transcript.delta", "response_id": response_id, "item_id": self._response_item_id(response_id), "output_index": 0, @@ -839,7 +839,7 @@ def _realtime_audio_delta_events( ) payloads.append( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "response_id": response_id, "item_id": item_id, "output_index": 0, @@ -872,7 +872,7 @@ def _realtime_audio_done_events( state.audio_done_emitted = True payloads.append( { - "type": "response.audio.done", + "type": "response.output_audio.done", "response_id": response_id, "item_id": item_id, "output_index": 0, @@ -882,7 +882,7 @@ def _realtime_audio_done_events( if transcript: payloads.append( { - "type": "response.audio_transcript.done", + "type": "response.output_audio_transcript.done", "response_id": response_id, "item_id": item_id, "output_index": 0, diff --git a/vllm_omni/entrypoints/duplex/warmup.py b/vllm_omni/entrypoints/duplex/warmup.py index c380775ef73..5051a1d318c 100644 --- a/vllm_omni/entrypoints/duplex/warmup.py +++ b/vllm_omni/entrypoints/duplex/warmup.py @@ -108,7 +108,7 @@ async def _recv_until(predicate, timeout_s: float) -> bool: sent += 1 await asyncio.sleep(0.08) # Wait for the pipeline's first audio output so every stage ran. - saw_audio = await _recv_until(lambda e: e.get("type") == "response.audio.delta", 30) + saw_audio = await _recv_until(lambda e: e.get("type") == "response.output_audio.delta", 30) # Close the session EXPLICITLY: a bare websocket close parks the # session in its disconnect grace window, where it would keep # occupying the max_sessions=1 slot and block the first client. diff --git a/vllm_omni/entrypoints/duplex/websocket.py b/vllm_omni/entrypoints/duplex/websocket.py index df50fd0c379..d51a7ed57bc 100644 --- a/vllm_omni/entrypoints/duplex/websocket.py +++ b/vllm_omni/entrypoints/duplex/websocket.py @@ -55,9 +55,7 @@ def normalize_duplex_input_event(event: dict[str, object]) -> dict[str, object]: "response.output_item.added", "response.content_part.added", "response.output_audio.delta", - "response.audio.delta", "response.output_audio.done", - "response.audio.done", "response.output_text.delta", "response.output_text.done", "response.text.delta", diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index a9256d95332..2f4e472f824 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -1600,7 +1600,7 @@ async def _generate_audio_sse_events( Field naming follows the OpenAI ``speech.audio.delta`` schema, which carries the base64 chunk in ``audio`` (not ``delta`` — that is the - Realtime API ``response.audio.delta`` convention, a different event). + Realtime API ``response.output_audio.delta`` convention, a different event). See https://platform.openai.com/docs/api-reference/audio-streaming. The terminal ``speech.audio.done`` event carries a ``usage`` object diff --git a/vllm_omni/experimental/fullduplex/client.py b/vllm_omni/experimental/fullduplex/client.py index 252fc92df2b..45cdde48262 100644 --- a/vllm_omni/experimental/fullduplex/client.py +++ b/vllm_omni/experimental/fullduplex/client.py @@ -234,7 +234,7 @@ def add(self, event: dict[str, object], *, received_at_s: float | None = None) - event_type = stored_event.get("type") if event_type == "response.created" and response_id and response_id not in self.response_ids: self.response_ids.append(response_id) - if event_type == "response.audio.delta": + if event_type == "response.output_audio.delta": delta = stored_event.get("delta") or stored_event.get("audio") if isinstance(delta, str) and response_id: try: @@ -263,7 +263,7 @@ def response_text(self, response_id: str) -> str: if self.response_id(event) == response_id and event.get("type") in { - "response.audio_transcript.delta", + "response.output_audio_transcript.delta", "response.output_text.delta", "response.text.delta", } @@ -322,7 +322,7 @@ def timing_summary( if ( event.get("type") in { - "response.audio_transcript.delta", + "response.output_audio_transcript.delta", "response.output_text.delta", "response.text.delta", } @@ -337,7 +337,7 @@ def timing_summary( if isinstance(stage0, dict): stage0_metrics = stage0 - if event.get("type") != "response.audio.delta" or ( + if event.get("type") != "response.output_audio.delta" or ( response_id is not None and event_response_id != response_id ): continue From 44fed170321fd8a8e0f117c9229dfcb13023f8c3 Mon Sep 17 00:00:00 2001 From: Jim Ban <77719403+BANANASJIM@users.noreply.github.com> Date: Thu, 10 Sep 2026 18:03:26 -0700 Subject: [PATCH 66/93] [Bugfix][Core] Clean up failed HWR atomic metadata writes (#6956) Signed-off-by: BANANASJIM Co-authored-by: Hongsheng Liu --- .../test_filesystem_store.py | 149 ++++++++++++++++++ .../host_weight_runtime/filesystem/store.py | 35 ++-- 2 files changed, 168 insertions(+), 16 deletions(-) diff --git a/tests/host_weight_runtime/test_filesystem_store.py b/tests/host_weight_runtime/test_filesystem_store.py index 7fe716c5ddd..d149d34bd0e 100644 --- a/tests/host_weight_runtime/test_filesystem_store.py +++ b/tests/host_weight_runtime/test_filesystem_store.py @@ -20,6 +20,7 @@ from dataclasses import replace from multiprocessing.process import BaseProcess from pathlib import Path +from typing import IO, Any import pytest import torch @@ -408,6 +409,154 @@ def _make_store( return FilesystemHostWeightStore(domain, capacity or CapacityPolicy(), integrity or IntegrityPolicy()) +class _FaultyTemporaryFile: + """NamedTemporaryFile proxy that injects a write or close failure.""" + + def __init__( + self, + handle: IO[bytes], + *, + write_error: BaseException | None = None, + close_error: BaseException | None = None, + ) -> None: + self._handle = handle + self._write_error = write_error + self._close_error = close_error + + def __getattr__(self, name: str) -> object: + return getattr(self._handle, name) + + def __enter__(self) -> _FaultyTemporaryFile: + return self + + def __exit__(self, *exc_info: object) -> None: + self._handle.close() + if self._close_error is not None: + raise self._close_error + + def write(self, data: bytes) -> int: + if self._write_error is not None: + raise self._write_error + return self._handle.write(data) + + +def _inject_temporary_file_failure( + monkeypatch: pytest.MonkeyPatch, + *, + write_error: BaseException | None = None, + close_error: BaseException | None = None, +) -> None: + original = tempfile.NamedTemporaryFile + + def faulty(*args: Any, **kwargs: Any) -> _FaultyTemporaryFile: + handle: IO[bytes] = original(*args, **kwargs) + return _FaultyTemporaryFile(handle, write_error=write_error, close_error=close_error) + + monkeypatch.setattr(tempfile, "NamedTemporaryFile", faulty) + + +def test_atomic_json_replace_failure_removes_temporary_file( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + path = tmp_path / "domain.json" + path.write_bytes(b"existing metadata") + replace_error = OSError(errno.EIO, "injected replace failure") + + def fail_replace(source: os.PathLike[str], destination: os.PathLike[str]) -> None: + raise replace_error + + monkeypatch.setattr(os, "replace", fail_replace) + + with pytest.raises(OSError) as error: + filesystem_store_module._write_atomic_json(path, {"schema_version": 1}) + + assert error.value is replace_error + assert path.read_bytes() == b"existing metadata" + assert not list(tmp_path.glob(".domain.json.*.tmp")) + + +def test_atomic_json_cleanup_failure_preserves_replace_error( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + path = tmp_path / "domain.json" + replace_error = OSError(errno.EIO, "injected replace failure") + original_unlink = Path.unlink + cleanup_attempted = False + + def fail_replace(source: os.PathLike[str], destination: os.PathLike[str]) -> None: + raise replace_error + + def fail_temporary_unlink(target: Path, missing_ok: bool = False) -> None: + nonlocal cleanup_attempted + if target.parent == tmp_path and target.name.startswith(".domain.json."): + cleanup_attempted = True + raise PermissionError(errno.EACCES, "injected cleanup failure") + original_unlink(target, missing_ok=missing_ok) + + monkeypatch.setattr(os, "replace", fail_replace) + monkeypatch.setattr(Path, "unlink", fail_temporary_unlink) + + with pytest.raises(OSError) as error: + filesystem_store_module._write_atomic_json(path, {"schema_version": 1}) + + assert cleanup_attempted + assert error.value is replace_error + assert len(list(tmp_path.glob(".domain.json.*.tmp"))) == 1 + + +@pytest.mark.parametrize("close_fails", [False, True]) +def test_atomic_json_write_failure_preserves_error_when_cleanup_fails( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + close_fails: bool, +) -> None: + path = tmp_path / "domain.json" + write_error = OSError(errno.EIO, "injected write failure") + original_unlink = Path.unlink + cleanup_attempted = False + + def fail_temporary_unlink(target: Path, missing_ok: bool = False) -> None: + nonlocal cleanup_attempted + if target.parent == tmp_path and target.name.startswith(".domain.json."): + cleanup_attempted = True + raise PermissionError(errno.EACCES, "injected cleanup failure") + original_unlink(target, missing_ok=missing_ok) + + _inject_temporary_file_failure( + monkeypatch, + write_error=write_error, + close_error=OSError(errno.EIO, "injected close failure") if close_fails else None, + ) + monkeypatch.setattr(Path, "unlink", fail_temporary_unlink) + + with pytest.raises(OSError) as error: + filesystem_store_module._write_atomic_json(path, {"schema_version": 1}) + + assert cleanup_attempted + assert error.value is write_error + assert not path.exists() + assert len(list(tmp_path.glob(".domain.json.*.tmp"))) == 1 + + +def test_atomic_json_close_failure_removes_temporary_file( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + path = tmp_path / "domain.json" + close_error = OSError(errno.EIO, "injected close failure") + + _inject_temporary_file_failure(monkeypatch, close_error=close_error) + + with pytest.raises(OSError) as error: + filesystem_store_module._write_atomic_json(path, {"schema_version": 1}) + + assert error.value is close_error + assert not path.exists() + assert not list(tmp_path.glob(".domain.json.*.tmp")) + + def _publish_test_artifact( store: FilesystemHostWeightStore, identity: WeightArtifactIdentity | None = None, diff --git a/vllm_omni/host_weight_runtime/filesystem/store.py b/vllm_omni/host_weight_runtime/filesystem/store.py index ec6b49f3c79..1542adefc00 100644 --- a/vllm_omni/host_weight_runtime/filesystem/store.py +++ b/vllm_omni/host_weight_runtime/filesystem/store.py @@ -12,6 +12,7 @@ import os import shutil import stat +import tempfile import time import uuid from collections import defaultdict @@ -298,25 +299,27 @@ def _move_artifact_locked(source: Path, destination: Path) -> None: def _write_atomic_json(path: Path, value: object, *, mode: int = 0o444) -> None: data = canonical_json(value) - temporary = path.parent / f".{path.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp" - fd = os.open(temporary, os.O_CREAT | os.O_EXCL | os.O_WRONLY | os.O_CLOEXEC, 0o600) + handle = tempfile.NamedTemporaryFile(dir=path.parent, prefix=f".{path.name}.", suffix=".tmp", delete=False) + temporary = Path(handle.name) + write_error: BaseException | None = None try: - offset = 0 - while offset < len(data): - count = os.write(fd, data[offset:]) - if count <= 0: - raise OSError(errno.EIO, f"short write for {temporary}") - offset += count - os.fsync(fd) - os.fchmod(fd, mode) - os.fsync(fd) + with handle: + try: + handle.write(data) + handle.flush() + os.fsync(handle.fileno()) + os.fchmod(handle.fileno(), mode) + os.fsync(handle.fileno()) + except BaseException as error: + write_error = error + raise + os.replace(temporary, path) except BaseException: - os.close(fd) - temporary.unlink(missing_ok=True) + with contextlib.suppress(OSError): + temporary.unlink(missing_ok=True) + if write_error is not None: + raise write_error raise - else: - os.close(fd) - os.replace(temporary, path) _fsync_directory(path.parent) From 3d952d1334f78ac6a4adc595974264457a8eb810 Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Fri, 11 Sep 2026 09:26:35 +0800 Subject: [PATCH 67/93] [Bugfix] Keep MiniMax-H3 reference audio budgets separate (#7281) Signed-off-by: david6666666 <530634352@qq.com> --- .../minimax_h3/test_minimax_h3_contract.py | 50 +++++++++++++++++++ .../models/minimax_h3/pipeline_minimax_h3.py | 5 +- 2 files changed, 53 insertions(+), 2 deletions(-) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index a94f842c5fe..5470cd9c477 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -2455,6 +2455,56 @@ def test_g4_standalone_audio_duration_and_total_duration_contract(): ) +def test_ref2va_video_soundtrack_does_not_consume_standalone_audio_budget(): + from PIL import Image + + from vllm_omni.diffusion.cache.cachedit.runtime import RequestScopedCacheDiTRuntime + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3.quality_policy import MiniMaxH3QualityPolicy + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "ref2va" + pipeline.supported_tasks = frozenset({"ref2va"}) + pipeline.device = torch.device("cpu") + pipeline.default_video_shift = 12.0 + pipeline.default_audio_shift = 3.0 + pipeline._resolve_sigma_positions = Mock(return_value=(None, 50)) + pipeline._quality_policy = MiniMaxH3QualityPolicy(None) + pipeline._cache_dit_runtime = Mock(spec=RequestScopedCacheDiTRuntime) + pipeline.text_encoder = object() + pipeline.encode_prompt = Mock(return_value=(torch.ones(1, 2), torch.ones(1, dtype=torch.long))) + pipeline._prepare_reference_videos = Mock(return_value=[{"input_has_audio": True}]) + pipeline._encode_visual_conditions = Mock(return_value=(torch.ones(1, 96), [(1, 16, 16), (17, 48, 84)])) + # A 2.35-second video soundtrack and a 15-second standalone reference + # each satisfy their modality's separate 15-second duration budget. + embedded = torch.full((188, 32), 1.0) + standalone = torch.full((1200, 32), 2.0) + pipeline._encode_reference_audio_conditions = Mock(return_value=(embedded, [94], standalone, [600])) + sampling = OmniDiffusionSamplingParams( + width=1344, + height=768, + fps=24, + num_inference_steps=50, + extra_args={"task": "ref2va", "duration": 15.0}, + ) + + context = pipeline._prepare_request_inputs( + prompt="A person carrying books down a street", + multi_modal_data={ + "image": Image.new("RGB", (256, 256)), + "video": ["reference.mp4"], + "audio": (torch.zeros(1, 15 * 16000), 16000), + }, + sampling=sampling, + ) + + assert context["audio_condition_lengths"] == [94, 600] + assert [block["kind"] for block in context["ref_blocks"]] == ["image", "video_audio", "audio"] + torch.testing.assert_close(context["audio_condition"], torch.cat([embedded, standalone])) + + def test_ref2va_audio_duration_validation_precedes_rank_branch(monkeypatch): from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as pipeline_module diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index ff1f8a3e5e7..fb9f02b28fa 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -2458,8 +2458,9 @@ def _prepare_request_inputs( if audio_lengths: if any(length < 80 or length > 600 for length in audio_lengths): raise OmniClientError("MiniMax H3 audio references must each be between 2 and 15 seconds") - if sum(audio_lengths) > 600: - raise OmniClientError("MiniMax H3 audio references must be at most 15 seconds in total") + # Video soundtracks and standalone audio have separate + # 15-second budgets, validated before encoding. Do not merge + # those budgets when concatenating their conditioning rows. if len(audio_lengths) == 1: ref_audio_t = audio_lengths[0] From bad50980c9e2679d4fa6352b0fbb121af9dfbb40 Mon Sep 17 00:00:00 2001 From: Yan Cao <31481315+yancaocn@users.noreply.github.com> Date: Fri, 11 Sep 2026 11:43:34 +0800 Subject: [PATCH 68/93] [Bugfix] Fix Helios USP: per-component split for correct sequence parallelism (#6930) Signed-off-by: yancaocn Co-authored-by: yancaocn --- tests/diffusion/models/helios/__init__.py | 0 .../models/helios/test_helios_usp_2rank.py | 258 ++++++++++++++++++ .../models/helios/test_helios_usp_shard.py | 225 +++++++++++++++ .../models/helios/helios_transformer.py | 90 +++++- 4 files changed, 561 insertions(+), 12 deletions(-) create mode 100644 tests/diffusion/models/helios/__init__.py create mode 100644 tests/diffusion/models/helios/test_helios_usp_2rank.py create mode 100644 tests/diffusion/models/helios/test_helios_usp_shard.py diff --git a/tests/diffusion/models/helios/__init__.py b/tests/diffusion/models/helios/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/diffusion/models/helios/test_helios_usp_2rank.py b/tests/diffusion/models/helios/test_helios_usp_2rank.py new file mode 100644 index 00000000000..7b16d1d4f5e --- /dev/null +++ b/tests/diffusion/models/helios/test_helios_usp_2rank.py @@ -0,0 +1,258 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Two-process USP smoke test for the Helios transformer. + +Pins the multi-rank mechanism that ``test_helios_usp_shard.py`` (unit tests +of ``_sp_split_seq`` alone) cannot see: + + * the manual ``_sp_shard_depth`` bump in ``forward()`` must make attn1 run + the Ulysses all-to-all (not ``NoParallelAttention``); + * attn2 cross-attention must stay rank-local (``skip_sequence_parallel= + True``): if replicated text K/V were fed through Ulysses they would be + ws-fold duplicated and outputs would diverge from the single-rank run; + * end-to-end correctness: ws=2 output equals the ws=1 baseline for the same + weights/inputs (divisible shapes, so the padding path is not involved). + +Runs ``HeliosTransformer3DModel.forward`` with tiny random weights via +``torch.multiprocessing.spawn``: first a baseline (world_size=1), then an +SP run (world_size=2, ulysses_degree=2); rank-0 outputs are compared. + +Requires >= 2 devices; skipped otherwise. +""" + +import os +import tempfile + +import pytest +import torch + +from tests.helpers.mark import hardware_marks +from vllm_omni.diffusion.config import set_current_diffusion_config +from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.parallel_state import ( + destroy_distributed_env, + init_distributed_environment, + initialize_model_parallel, +) +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelConfig, + get_sp_plan_from_model, +) +from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context +from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel +from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel +from vllm_omni.platforms import current_omni_platform + +_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2) +pytestmark = [pytest.mark.full_model, pytest.mark.diffusion, *_L4_TWO_GPU] + +RANDOM_SEED = 42 +DTYPE = torch.float32 # keep both runs on the SDPA fallback for tight tolerance + +# Tiny but architecturally faithful config: +# inner_dim = 4 * 32 = 128, rope_dim sums to head_dim (32), +# current chunk 2x4x6 latent -> (1,2,2) patch -> 2*2*3 = 12 tokens, +# short history 1x4x6 latent -> (1,2,2) patch -> 6 tokens. +# All component lengths are even, so ws=2 needs no padding replicas. +TINY_MODEL_KWARGS = dict( + patch_size=(1, 2, 2), + num_attention_heads=4, + attention_head_dim=32, + in_channels=4, + out_channels=4, + text_dim=64, + freq_dim=32, + ffn_dim=128, + num_layers=2, + rope_dim=(12, 10, 10), + rope_theta=10000.0, + guidance_cross_attn=True, + zero_history_timestep=True, + has_multi_term_memory_patch=True, + is_amplify_history=False, +) + +BATCH = 1 +CUR_FRAMES, LAT_H, LAT_W = 2, 4, 6 # -> 12 current tokens +HIST_FRAMES = 1 # short history -> 6 tokens +TEXT_LEN = 5 + + +def _update_env(envs_dict: dict[str, str]) -> None: + for k, v in envs_dict.items(): + os.environ[k] = v + + +def _make_inputs(device: torch.device) -> dict[str, torch.Tensor]: + torch.manual_seed(RANDOM_SEED + 1) + return { + "hidden_states": torch.randn(BATCH, 4, CUR_FRAMES, LAT_H, LAT_W, device=device, dtype=DTYPE), + "timestep": torch.full((BATCH,), 500.0, device=device), + "encoder_hidden_states": torch.randn( + BATCH, TEXT_LEN, TINY_MODEL_KWARGS["text_dim"], device=device, dtype=DTYPE + ), + "indices_latents_history_short": torch.tensor([[0]], device=device), + "latents_history_short": torch.randn(BATCH, 4, HIST_FRAMES, LAT_H, LAT_W, device=device, dtype=DTYPE), + } + + +def _spy_strategies(model: HeliosTransformer3DModel) -> dict[str, list[str]]: + """Record the parallel strategy each Attention resolves to during forward.""" + seen: dict[str, list[str]] = {"attn1": [], "attn2": []} + for name, attn in (("attn1", model.blocks[0].attn1.attn), ("attn2", model.blocks[0].attn2.attn)): + original = attn._get_active_parallel_strategy + + def spy(original=original, name=name): + strategy = original() + seen[name].append(getattr(strategy, "name", type(strategy).__name__)) + return strategy + + attn._get_active_parallel_strategy = spy + return seen + + +def helios_forward_on_model( + local_rank: int, + world_size: int, + ulysses_degree: int, + output_file: str, + model_state_file: str, + is_baseline: bool, +) -> None: + device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") + current_omni_platform.set_device(device) + torch.set_default_device(device) + torch.set_default_dtype(DTYPE) + + _update_env( + { + "RANK": str(local_rank), + "LOCAL_RANK": str(local_rank), + "WORLD_SIZE": str(world_size), + "MASTER_ADDR": "localhost", + "MASTER_PORT": "12361", + } + ) + init_distributed_environment() + + parallel_config = DiffusionParallelConfig( + pipeline_parallel_size=1, + data_parallel_size=1, + tensor_parallel_size=1, + sequence_parallel_size=ulysses_degree, + ulysses_degree=ulysses_degree, + ring_degree=1, + allgather_degree=1, + cfg_parallel_size=1, + ) + od_config = OmniDiffusionConfig.from_kwargs( + model="test_model", + dtype=DTYPE, + parallel_config=parallel_config, + diffusion_attention_backend="TORCH_SDPA", + ) + initialize_model_parallel( + data_parallel_size=1, + cfg_parallel_size=1, + sequence_parallel_size=ulysses_degree, + ulysses_degree=ulysses_degree, + ring_degree=1, + allgather_degree=1, + tensor_parallel_size=1, + pipeline_parallel_size=1, + ) + + with set_forward_context(omni_diffusion_config=od_config), set_current_diffusion_config(od_config): + torch.manual_seed(RANDOM_SEED) + model = HeliosTransformer3DModel(**TINY_MODEL_KWARGS).to(device).to(DTYPE) + + if is_baseline and local_rank == 0: + with open(model_state_file, "wb") as f: + torch.save({k: v.cpu() for k, v in model.state_dict().items()}, f) + + if world_size > 1: + torch.distributed.barrier() + + with open(model_state_file, "rb") as f: + model_state = torch.load(f) + model.load_state_dict({k: v.to(device).to(DTYPE) for k, v in model_state.items()}) + + if not is_baseline: + # Mirror registry._apply_sequence_parallel_if_enabled: install the + # proj_out gather hook from the model's own _sp_plan and flip + # sp_plan_hooks_applied so sp_active follows _sp_shard_depth + # (i.e. the manual bump in Helios forward is what enables Ulysses). + plan = get_sp_plan_from_model(model) + assert plan is not None, "HeliosTransformer3DModel must define _sp_plan" + apply_sequence_parallel(model, SequenceParallelConfig(ulysses_degree=ulysses_degree), plan) + get_forward_context().sp_plan_hooks_applied = True + + seen = _spy_strategies(model) + inputs = _make_inputs(device) + output = model(**inputs).sample + + assert output.shape == (BATCH, 4, CUR_FRAMES, LAT_H, LAT_W), f"unexpected output shape {output.shape}" + + if is_baseline: + assert seen["attn1"] and all(n == "none" for n in seen["attn1"]), ( + f"baseline attn1 must use NoParallelAttention ('none'), saw {seen['attn1']}" + ) + else: + # attn1: SP communication must be active inside the block loop + # (regression here = the manual _sp_shard_depth bump was lost). + assert seen["attn1"] and all(n == "ulysses" for n in seen["attn1"]), ( + f"attn1 must run Ulysses under USP (manual _sp_shard_depth bump), saw {seen['attn1']}" + ) + # attn2: replicated text K/V must stay rank-local + # (regression here = the ws-fold duplicated-keys bug). + assert seen["attn2"] and all(n == "none" for n in seen["attn2"]), ( + f"attn2 must skip sequence parallel (replicated text K/V), saw {seen['attn2']}" + ) + + if local_rank == 0: + with open(output_file, "wb") as f: + torch.save(output.detach().cpu().float(), f) + + destroy_distributed_env() + + +def test_helios_usp2_matches_single_rank(): + """ws=2 (Ulysses) must reproduce the ws=1 output for identical weights/inputs.""" + available_devices = current_omni_platform.get_device_count() + if available_devices < 2: + pytest.skip(f"Test requires 2 devices but only {available_devices} available") + + with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f: + baseline_output_file = f.name + with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f: + sp_output_file = f.name + with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f: + model_state_file = f.name + + try: + torch.multiprocessing.spawn( + helios_forward_on_model, + args=(1, 1, baseline_output_file, model_state_file, True), + nprocs=1, + ) + torch.multiprocessing.spawn( + helios_forward_on_model, + args=(2, 2, sp_output_file, model_state_file, False), + nprocs=2, + ) + + with open(baseline_output_file, "rb") as f: + baseline = torch.load(f) + with open(sp_output_file, "rb") as f: + sp = torch.load(f) + + assert baseline.shape == sp.shape + max_abs_diff = (baseline - sp).abs().max().item() + assert torch.allclose(baseline, sp, atol=1e-3, rtol=1e-3), ( + f"USP=2 output diverges from single-rank baseline: max_abs_diff={max_abs_diff:.6e}" + ) + finally: + for path in [baseline_output_file, sp_output_file, model_state_file]: + if os.path.exists(path): + os.remove(path) diff --git a/tests/diffusion/models/helios/test_helios_usp_shard.py b/tests/diffusion/models/helios/test_helios_usp_shard.py new file mode 100644 index 00000000000..88301a85159 --- /dev/null +++ b/tests/diffusion/models/helios/test_helios_usp_shard.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for Helios per-component sequence splitting. + +These tests verify the manual sharding logic that replaces the broken +rope/blocks.0 hooks in the Helios ``_sp_plan``. They exercise the *real* +``HeliosTransformer3DModel._sp_split_seq`` (the parallel-state getters it +imports lazily inside the method body are monkeypatched), so the model-side +helper cannot drift silently from what is tested. For the multi-rank +mechanism around the helper (manual ``_sp_shard_depth`` bump, attn1/attn2 +parallel strategy, ws=2 vs ws=1 output equivalence) see +``test_helios_usp_2rank.py``. + +Coverage: + 1. USP disabled (ws == 1) returns the tensor unchanged. + 2. Divisible split gives each rank the expected slice. + 3. Non-divisible split pads with replicated tokens (no crash). + 4. Per-component split: history + current tokens are distributed so + that *both* ranks receive some current tokens (no mosaic regression). + 5. Edge cases: 1-D tensor, empty seq, 4-D tensor. +""" + +from collections.abc import Iterator +from contextlib import contextmanager + +import pytest +import torch + +import vllm_omni.diffusion.distributed.parallel_state as sp_parallel_state +from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +@contextmanager +def _fake_sp_state(ws: int, rank: int) -> Iterator[None]: + """Patch the parallel-state getters used by the real helper. + + ``_sp_split_seq`` imports these names from the module *inside* its + function body, so replacing the module attributes takes effect. + """ + prev_ws = sp_parallel_state.get_sequence_parallel_world_size + prev_rank = sp_parallel_state.get_sequence_parallel_rank + sp_parallel_state.get_sequence_parallel_world_size = lambda: ws + sp_parallel_state.get_sequence_parallel_rank = lambda: rank + try: + yield + finally: + sp_parallel_state.get_sequence_parallel_world_size = prev_ws + sp_parallel_state.get_sequence_parallel_rank = prev_rank + + +def _sp_split_seq(x: torch.Tensor, ws: int, rank: int) -> torch.Tensor: + """Call the real model-side helper without constructing a model.""" + with _fake_sp_state(ws, rank): + return HeliosTransformer3DModel._sp_split_seq(None, x) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestSpSplitSeqDisabled: + """When USP is off (ws == 1) the helper must be a no-op.""" + + def test_ws1_returns_unchanged(self): + x = torch.randn(2, 100, 64) + out = _sp_split_seq(x, ws=1, rank=0) + assert torch.equal(out, x) + assert out is x + + +class TestSpSplitSeqDivisible: + """Divisible sequences: each rank gets exactly seq_len / ws tokens.""" + + def test_each_rank_gets_correct_slice(self): + for ws in [2, 4, 8]: + seq_len = ws * 10 + x = torch.arange(seq_len * 4).float().reshape(1, seq_len, 4) + for rank in range(ws): + out = _sp_split_seq(x, ws=ws, rank=rank) + n = seq_len // ws + expected = x[:, rank * n : (rank + 1) * n, :] + assert out.shape == (1, n, 4) + assert torch.equal(out, expected) + + def test_output_is_contiguous(self): + x = torch.randn(1, 20, 8) + out = _sp_split_seq(x, ws=2, rank=0) + assert out.is_contiguous() + + +class TestSpSplitSeqNonDivisible: + """Non-divisible sequences are padded with replicated tokens.""" + + def test_non_divisible_pads_instead_of_crash(self): + """400x272 → 9*25*17 = 3825 tokens, USP=2 → not divisible. + Should pad to 3826 and split, not crash.""" + x = torch.randn(1, 3825, 8) + for rank in range(2): + out = _sp_split_seq(x, ws=2, rank=rank) + # 3825 → pad 1 → 3826 → /2 = 1913 + assert out.shape[1] == 1913 + + def test_non_divisible_pads_with_replicated_last_token(self): + """Verify the padding is the replicated last token, not zeros.""" + x = torch.randn(1, 7, 4) # 7 % 2 = 1, pad 1 + last_token = x[:, -1:, :].clone() + out = _sp_split_seq(x, ws=2, rank=0) + # rank 0 gets tokens [0:4] (first 4 of the padded 8) + assert out.shape == (1, 4, 4) + # rank 1 gets tokens [4:8] — token 7 (the padded one) = last token + out1 = _sp_split_seq(x, ws=2, rank=1) + assert out1.shape == (1, 4, 4) + assert torch.equal(out1[:, -1:, :], last_token) + + def test_non_divisible_7_4(self): + """7 tokens, ws=4 → pad 1 → 8 → /4 = 2 per rank.""" + x = torch.randn(1, 7, 8) + for rank in range(4): + out = _sp_split_seq(x, ws=4, rank=rank) + assert out.shape[1] == 2 + + def test_non_divisible_15_4(self): + """15 tokens, ws=4 → pad 1 → 16 → /4 = 4 per rank.""" + x = torch.randn(1, 15, 8) + for rank in range(4): + out = _sp_split_seq(x, ws=4, rank=rank) + assert out.shape[1] == 4 + + +class TestSpSplitSeqPerComponent: + """Simulate Helios forward: split current + history components so that + *both* ranks receive some current tokens (no mosaic regression).""" + + def test_both_ranks_get_current_tokens(self): + """history=2400, current=540, ws=2. + After per-component split, each rank must have > 0 current tokens.""" + current = torch.randn(1, 540, 64) + history = torch.randn(1, 2400, 64) + + for rank in range(2): + cur_local = _sp_split_seq(current, ws=2, rank=rank) + hist_local = _sp_split_seq(history, ws=2, rank=rank) + + assert cur_local.shape[1] > 0, f"Rank {rank} received 0 current tokens" + assert cur_local.shape[1] == 270 # 540 / 2 + assert hist_local.shape[1] == 1200 # 2400 / 2 + + original_context_length = cur_local.shape[1] + assert original_context_length > 0 # no mosaic regression + + def test_gather_restores_full_current(self): + """After splitting current into 2 ranks, simulating proj_out + gather must restore the full current sequence length so that + the original unpatchify reshape succeeds.""" + current = torch.randn(1, 540, 64) + local_chunks = [_sp_split_seq(current, ws=2, rank=rank) for rank in range(2)] + gathered = torch.cat(local_chunks, dim=1) + assert gathered.shape[1] == 540 # full current restored + + def test_gather_with_padding_truncates_correctly(self): + """Non-divisible current (3825) → pad to 3826 → split → + gather → 3826. Slice to 3825 for unpatchify.""" + current = torch.randn(1, 3825, 64) + local_chunks = [_sp_split_seq(current, ws=2, rank=rank) for rank in range(2)] + gathered = torch.cat(local_chunks, dim=1) + assert gathered.shape[1] == 3826 # padded size + # Simulate the slice before unpatchify + expected_seq = 3825 + if gathered.shape[1] > expected_seq: + gathered = gathered[:, :expected_seq, :] + assert gathered.shape[1] == 3825 # correct after slice + + def test_rotary_emb_aligned_with_hidden_states(self): + """rotary_emb and hidden_states must be split at the same + positions so that tokens align inside the attention layer.""" + hidden = torch.randn(1, 540, 5120) + rotary = torch.randn(1, 540, 256) + for rank in range(2): + h_local = _sp_split_seq(hidden, ws=2, rank=rank) + r_local = _sp_split_seq(rotary, ws=2, rank=rank) + assert h_local.shape[1] == r_local.shape[1] + + def test_whole_split_causes_mosaic_but_per_component_does_not(self): + """Demonstrates why per-component split is needed: + with whole split (simulated), rank 0 gets 0 current tokens; + with per-component split, both ranks get current tokens.""" + history = torch.randn(1, 2400, 64) + current = torch.randn(1, 540, 64) + full = torch.cat([history, current], dim=1) # 2940 + + # Whole split (the OLD broken behavior) + whole_n = full.shape[1] // 2 # 1470 + rank0_whole = full[:, :whole_n, :] + rank0_current = rank0_whole.shape[1] - 2400 # negative → 0 current + assert rank0_current <= 0, "Whole split should give rank 0 no current" + + # Per-component split (the NEW correct behavior) + for rank in range(2): + cur_local = _sp_split_seq(current, ws=2, rank=rank) + assert cur_local.shape[1] == 270 # both ranks have current + + +class TestSpSplitSeqEdgeCases: + """Edge cases that should not crash or silently misbehave.""" + + def test_1d_tensor_skipped(self): + """1-D tensors must be returned unchanged (dim < 2 guard).""" + x = torch.randn(100) + out = _sp_split_seq(x, ws=2, rank=0) + assert torch.equal(out, x) + + def test_empty_seq_skipped(self): + """Empty sequence dimension must be returned unchanged.""" + x = torch.randn(1, 0, 64) + out = _sp_split_seq(x, ws=2, rank=0) + assert out.shape == x.shape + + def test_4d_tensor_split_on_dim1(self): + """4-D tensor [B, S, H, D] should split along dim=1.""" + x = torch.randn(1, 20, 8, 32) + out = _sp_split_seq(x, ws=2, rank=1) + assert out.shape == (1, 10, 8, 32) diff --git a/vllm_omni/diffusion/models/helios/helios_transformer.py b/vllm_omni/diffusion/models/helios/helios_transformer.py index 4e3cf5078f6..6c6ec223dc3 100644 --- a/vllm_omni/diffusion/models/helios/helios_transformer.py +++ b/vllm_omni/diffusion/models/helios/helios_transformer.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 # Adapted from Helios (https://github.com/BestWishYsh/Helios) import math @@ -27,10 +27,7 @@ from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig -from vllm_omni.diffusion.distributed.sp_plan import ( - SequenceParallelInput, - SequenceParallelOutput, -) +from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelOutput if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( @@ -448,6 +445,12 @@ def __init__( num_kv_heads=self.num_heads, softmax_scale=1.0 / (head_dim**0.5), causal=False, + # Text K/V is replicated on every SP rank, so Ulysses must not be + # used here: its all-to-all would treat the per-rank replicas as + # sequence shards and ws-fold duplicate the keys. Sharded Q over + # full replicated K/V is correct locally and needs no + # communication (same as Wan2.2 cross-attn). + skip_sequence_parallel=True, ) def project_kv(self, encoder_hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: @@ -624,16 +627,50 @@ def _is_transformer_block(name: str, module) -> bool: _hsdp_shard_conditions = [_is_transformer_block] _sp_plan = { - "rope": { - 0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True), - 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True), - }, - "blocks.0": { - "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True), - }, + # "rope" and "blocks.0" are intentionally omitted. + # The rope hook splits along dim=1 of the 5D output [B, D, T, H, W], + # which is the frequency dimension, NOT the sequence dimension. + # The blocks.0 hook splits the concatenated hidden_states as a whole, + # which can put all history tokens in one rank (original_context_length=0). + # Instead, hidden_states and rotary_emb are split per-component in + # forward() after flatten+transpose, ensuring each rank gets half + # of each component (history and current). "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), } + def _sp_split_seq(self, x: torch.Tensor) -> torch.Tensor: + """Split tensor along sequence dim (dim=1) for Ulysses SP. + + If seq_len is not divisible by world_size, replicate the last + token(s) to pad to a divisible length. This avoids crashing + on non-divisible history components (which use different patch + sizes and thus produce different token counts). + + Residual caveat: the replicated padding tokens attend unmasked + in attn1 on the last rank (at most ws-1 replicas per component). + The impact is bounded, but this is the first place to look if + odd resolutions ever show quality drift. + """ + from vllm_omni.diffusion.distributed.parallel_state import ( + get_sequence_parallel_rank, + get_sequence_parallel_world_size, + ) + + ws = get_sequence_parallel_world_size() + if ws > 1 and x.dim() >= 2 and x.shape[1] > 0: + seq_len = x.shape[1] + remainder = seq_len % ws + if remainder != 0: + pad = ws - remainder + # Replicate the last token to fill the padding + last = x[:, -1:, ...].expand(-1, pad, *([-1] * (x.dim() - 2))) + x = torch.cat([x, last], dim=1) + seq_len = x.shape[1] + r = get_sequence_parallel_rank() + n = seq_len // ws + x = x[:, r * n : (r + 1) * n, ...].contiguous() + return x + def __init__( self, patch_size: tuple[int, ...] = (1, 2, 2), @@ -860,6 +897,9 @@ def forward( device=hidden_states.device, ) rotary_emb = rotary_emb.flatten(2).transpose(1, 2) + # USP: per-component split (each rank gets half of current frames) + hidden_states = self._sp_split_seq(hidden_states) + rotary_emb = self._sp_split_seq(rotary_emb) original_context_length = hidden_states.shape[1] # 2. Process short history latents @@ -876,6 +916,9 @@ def forward( device=latents_history_short.device, ) rotary_emb_history_short = rotary_emb_history_short.flatten(2).transpose(1, 2) + # USP: per-component split + latents_history_short = self._sp_split_seq(latents_history_short) + rotary_emb_history_short = self._sp_split_seq(rotary_emb_history_short) hidden_states = torch.cat([latents_history_short, hidden_states], dim=1) rotary_emb = torch.cat([rotary_emb_history_short, rotary_emb], dim=1) @@ -896,6 +939,9 @@ def forward( rotary_emb_history_mid = pad_for_3d_conv(rotary_emb_history_mid, (2, 2, 2)) rotary_emb_history_mid = center_down_sample_3d(rotary_emb_history_mid, (2, 2, 2)) rotary_emb_history_mid = rotary_emb_history_mid.flatten(2).transpose(1, 2) + # USP: per-component split + latents_history_mid = self._sp_split_seq(latents_history_mid) + rotary_emb_history_mid = self._sp_split_seq(rotary_emb_history_mid) hidden_states = torch.cat([latents_history_mid, hidden_states], dim=1) rotary_emb = torch.cat([rotary_emb_history_mid, rotary_emb], dim=1) @@ -916,6 +962,9 @@ def forward( rotary_emb_history_long = pad_for_3d_conv(rotary_emb_history_long, (4, 4, 4)) rotary_emb_history_long = center_down_sample_3d(rotary_emb_history_long, (4, 4, 4)) rotary_emb_history_long = rotary_emb_history_long.flatten(2).transpose(1, 2) + # USP: per-component split + latents_history_long = self._sp_split_seq(latents_history_long) + rotary_emb_history_long = self._sp_split_seq(rotary_emb_history_long) hidden_states = torch.cat([latents_history_long, hidden_states], dim=1) rotary_emb = torch.cat([rotary_emb_history_long, rotary_emb], dim=1) @@ -958,6 +1007,18 @@ def forward( timestep_proj = timestep_proj.permute(0, 2, 1, 3) # 6. Transformer blocks + # Manually increment _sp_shard_depth so that attention layers + # know SP is active and perform Ulysses All-to-All communication. + # Normally this is done by split_output=True hooks (e.g. rope), + # but those are removed because they split the wrong dimension. + from vllm_omni.diffusion.forward_context import ( + get_forward_context, + is_forward_context_available, + ) + + if is_forward_context_available(): + get_forward_context()._sp_shard_depth += 1 + hidden_states = hidden_states.contiguous() encoder_hidden_states = encoder_hidden_states.contiguous() rotary_emb = rotary_emb.contiguous() @@ -979,6 +1040,11 @@ def forward( hidden_states = self.proj_out(hidden_states) # 8. Unpatchify + # proj_out gather may include padded tokens from non-divisible + # sequences. Slice to the expected size before reshape. + expected_seq = post_patch_num_frames * post_patch_height * post_patch_width + if hidden_states.shape[1] > expected_seq: + hidden_states = hidden_states[:, :expected_seq, :] hidden_states = hidden_states.reshape( batch_size, post_patch_num_frames, post_patch_height, post_patch_width, p_t, p_h, p_w, -1 ) From 96360ccb00d0c4c415eaec4369c827e7396e910a Mon Sep 17 00:00:00 2001 From: Samit <285365963@qq.com> Date: Fri, 11 Sep 2026 16:28:42 +0800 Subject: [PATCH 69/93] [Perf][Diffusion] Optimize HSDP startup via Rank-0 shared weight loading and accelerated LoRA delta computation (#7005) Signed-off-by: samithuang <285365963@qq.com> --- docs/mkdocs/hooks/search_index.py | 2 +- tests/config/test_omni_config.py | 1 + tests/diffusion/lora/test_loader.py | 12 + .../model_loader/test_diffusers_loader.py | 449 +++++++++++++++++- tests/engine/test_stage_engine_args.py | 1 + vllm_omni/config/omni_config.py | 3 + vllm_omni/config/stage_config.py | 1 + vllm_omni/diffusion/data.py | 1 + vllm_omni/diffusion/lora/loader.py | 21 +- .../model_loader/diffusers_loader.py | 151 +++++- .../model_loader/host_weight_plan.py | 3 + vllm_omni/engine/arg_utils.py | 1 + vllm_omni/entrypoints/cli/serve.py | 7 + 13 files changed, 637 insertions(+), 16 deletions(-) diff --git a/docs/mkdocs/hooks/search_index.py b/docs/mkdocs/hooks/search_index.py index 1b9e268d977..8cf23e8698f 100644 --- a/docs/mkdocs/hooks/search_index.py +++ b/docs/mkdocs/hooks/search_index.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 """Reduce browser search-index memory for generated API reference pages.""" from material.plugins.search.plugin import SearchIndex diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index da8ffb4525e..fdb2cd4b970 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -714,6 +714,7 @@ def test_sub_config_fields_match_structured_scopes(): "max_cudagraph_capture_size", "enable_flashinfer_autotune", "enable_multithread_weight_load", + "enable_broadcast_weight_load", "num_weight_load_threads", "disable_autocast", # Per-stage checkpoint resolution for repos whose stages live in diff --git a/tests/diffusion/lora/test_loader.py b/tests/diffusion/lora/test_loader.py index deb5ec80d93..4fd8769372c 100644 --- a/tests/diffusion/lora/test_loader.py +++ b/tests/diffusion/lora/test_loader.py @@ -509,3 +509,15 @@ def test_lora_loader_mixin_lora_is_fused_property(mocker: MockerFixture): assert pipeline2.lora_is_fused is False pipeline2.lora_is_fused = True assert pipeline2.lora_is_fused is True + + +def test_prepare_lora_delta_with_compute_device(): + base_key = "blocks.0.attn.to_q" + state_dict = { + f"{base_key}.lora_A.weight": torch.ones(RANK, HEAD_DIM), + f"{base_key}.lora_B.weight": torch.ones(HEAD_DIM, RANK), + } + delta, used_keys = _prepare_lora_delta(state_dict, base_key, compute_device=torch.device("cpu")) + assert delta is not None + assert delta.shape == (HEAD_DIM, HEAD_DIM) + assert len(used_keys) == 2 diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 565949ff616..c7a87ac343a 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -6,9 +6,11 @@ """ import json +import os +import tempfile from pathlib import Path from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest import torch @@ -260,17 +262,16 @@ def make_loader() -> tuple[DiffusersPipelineLoader, _HWRPipeline]: assert len(tuple((store_root / "source-digests-v1" / "entries").glob("*.json"))) == 1 -def test_maybe_fuse_distilled_lora_skips_when_hwr_warm_snapshot_present(): +def test_maybe_fuse_distilled_lora_skips_when_lora_path_unset(): cfg = SimpleNamespace( lora_backend="distill", - lora_path="/path/to/lora.safetensors", + lora_path=None, lora_scale=1.0, dtype=torch.bfloat16, quantization_config=None, parallel_config=SimpleNamespace(use_hsdp=False), ) loader = DiffusersPipelineLoader(LoadConfig(), cfg) - loader._hwr_state = {"warm_snapshot": {"transformer.weight": (1234, "hash")}} model = nn.Module() model.load_lora_weights = MagicMock() @@ -278,7 +279,7 @@ def test_maybe_fuse_distilled_lora_skips_when_hwr_warm_snapshot_present(): loader._maybe_fuse_distilled_lora(model) model.load_lora_weights.assert_not_called() - assert getattr(model, "lora_is_fused", False) is True + assert getattr(model, "lora_is_fused", False) is False def test_maybe_fuse_distilled_lora_fuses_when_not_warm_snapshot(): @@ -1053,6 +1054,34 @@ def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch): assert loader.take_host_weight_plan() is None +def test_dlo_mmap_plan_with_distilled_lora_falls_back_to_ordinary_loader(monkeypatch): + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace(use_hsdp=False, tensor_parallel_size=1), + quantization_config=None, + enable_distributed_layerwise_offload=True, + dlo_use_allgather=False, + lora_backend=LoRABackend.DISTILL, + lora_path="/fake/lora.safetensors", + model="unused", + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + model.load_lora_weights = MagicMock() + calls: list[str] = [] + + loader._init_from_load_format = lambda *_args, **_kwargs: model # type: ignore[method-assign] + loader.load_weights = lambda _model: calls.append("load") # type: ignore[method-assign] + loader._process_weights_after_loading = lambda *_args: calls.append("process") # type: ignore[method-assign] + loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign] + + assert loader.load_model(load_device="cpu") is model + assert calls == ["load", "process"] + assert model.load_lora_weights.call_count == 1 + assert loader.take_host_weight_plan() is None + + def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch): from vllm.model_executor.layers.quantization.online.fp8 import ( Fp8PerTensorOnlineLinearMethod, @@ -1487,3 +1516,413 @@ def test_load_model(prefetch_helios_model, mock_tp_group): ) model = loader.load_model(load_device="cpu") assert isinstance(model, HeliosPipeline) + + +def test_hsdp_broadcast_weight_load_rank0(mocker): + """Ensure that on rank 0 with enable_broadcast_weight_load=True, weights are loaded, fused, and broadcast.""" + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + from vllm_omni.diffusion.offloader.module_collector import PipelineModules + + mocker.patch("torch.distributed.is_initialized", return_value=True) + mocker.patch("torch.distributed.get_world_size", return_value=4) + mocker.patch("torch.distributed.get_rank", return_value=0) + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace( + use_hsdp=True, + hsdp_replicate_size=1, + hsdp_shard_size=4, + ), + enable_broadcast_weight_load=True, + lora_backend=LoRABackend.DISTILL, + lora_path="/fake/lora.safetensors", + quantization_config=None, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + events: list[str] = [] + + loader._init_from_load_format = mocker.Mock(return_value=model) # type: ignore[method-assign] + loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load")) # type: ignore[method-assign] + loader._maybe_fuse_distilled_lora = mocker.Mock( # type: ignore[method-assign] + side_effect=lambda _model: events.append("fuse") + ) + loader._broadcast_model_weights = mocker.Mock( # type: ignore[method-assign] + side_effect=lambda _model, **_kwargs: events.append("broadcast") + ) + loader._process_weights_after_loading = mocker.Mock( # type: ignore[method-assign] + side_effect=lambda _model, _device: events.append("process") + ) + mocker.patch.object( + loader_mod.ModuleDiscovery, + "discover", + return_value=PipelineModules( + dits=[model.transformer], + dit_names=["transformer"], + vaes=[], + encoders=[], + encoder_names=[], + resident_modules=[], + resident_names=[], + ), + ) + mocker.patch.object( + loader_mod, + "apply_hsdp_to_model", + side_effect=lambda *_args, **_kwargs: events.append("shard"), + ) + + loader._load_model_with_hsdp(torch.device("cpu")) + + assert events == ["load", "fuse", "broadcast", "process", "shard"] + + +def test_hsdp_broadcast_weight_load_rank_nonzero(mocker): + """Ensure that on rank > 0 with enable_broadcast_weight_load=True, loading is skipped and broadcast is received.""" + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + from vllm_omni.diffusion.offloader.module_collector import PipelineModules + + mocker.patch("torch.distributed.is_initialized", return_value=True) + mocker.patch("torch.distributed.get_world_size", return_value=4) + mocker.patch("torch.distributed.get_rank", return_value=1) + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace( + use_hsdp=True, + hsdp_replicate_size=1, + hsdp_shard_size=4, + ), + enable_broadcast_weight_load=True, + lora_backend=LoRABackend.DISTILL, + lora_path="/fake/lora.safetensors", + quantization_config=None, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + model.load_lora_weights = mocker.Mock() # type: ignore[assignment] + events: list[str] = [] + + loader._init_from_load_format = mocker.Mock(return_value=model) # type: ignore[method-assign] + loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load")) # type: ignore[method-assign] + loader._maybe_fuse_distilled_lora = mocker.Mock( # type: ignore[method-assign] + side_effect=lambda _model: events.append("fuse") + ) + loader._broadcast_model_weights = mocker.Mock( # type: ignore[method-assign] + side_effect=lambda _model, **_kwargs: events.append("broadcast") + ) + loader._process_weights_after_loading = mocker.Mock( # type: ignore[method-assign] + side_effect=lambda _model, _device: events.append("process") + ) + mocker.patch.object( + loader_mod.ModuleDiscovery, + "discover", + return_value=PipelineModules( + dits=[model.transformer], + dit_names=["transformer"], + vaes=[], + encoders=[], + encoder_names=[], + resident_modules=[], + resident_names=[], + ), + ) + mocker.patch.object( + loader_mod, + "apply_hsdp_to_model", + side_effect=lambda *_args, **_kwargs: events.append("shard"), + ) + + loader._load_model_with_hsdp(torch.device("cpu")) + + # Rank 1 should NOT call load_weights or _maybe_fuse_distilled_lora from disk + assert "load" not in events + assert "fuse" not in events + assert events == ["broadcast", "process", "shard"] + assert getattr(model, "lora_is_fused", False) is True + + +def test_broadcast_model_weights_invokes_dist_broadcast(mocker): + broadcast_calls = [] + mocker.patch("torch.distributed.is_initialized", return_value=True) + mocker.patch("torch.distributed.get_world_size", return_value=2) + mocker.patch("torch.distributed.get_rank", return_value=0) + mocker.patch("torch.distributed.get_backend", return_value="gloo") + mocker.patch("torch.distributed.broadcast", side_effect=lambda tensor, src: broadcast_calls.append((tensor, src))) + mocker.patch("torch.distributed.barrier") + + model = nn.Module() + model.linear = nn.Linear(2, 2) + model.register_buffer("buf", torch.tensor([1.0, 2.0])) + model.register_buffer("int_buf", torch.tensor([1, 2, 3], dtype=torch.long)) + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace(use_hsdp=True), + enable_broadcast_weight_load=True, + quantization_config=None, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + loader._broadcast_model_weights(model, target_device=torch.device("cpu"), src_rank=0) + + # float32 tensors (linear.weight, linear.bias, buf) are coalesced into 1 bucket, + # and int64 tensor (int_buf) is in a separate bucket -> 2 coalesced broadcast calls + assert len(broadcast_calls) == 2 + assert all(src == 0 for _, src in broadcast_calls) + + +def test_broadcast_model_weights_receiver_copies_data(mocker): + """Ensure receiver rank copies broadcasted flat tensor data into its model parameters/buffers.""" + mocker.patch("torch.distributed.is_initialized", return_value=True) + mocker.patch("torch.distributed.get_world_size", return_value=2) + mocker.patch("torch.distributed.get_rank", return_value=1) + mocker.patch("torch.distributed.get_backend", return_value="gloo") + mocker.patch("torch.distributed.barrier") + + # Simulate broadcast by filling the flat tensor with source values + def mock_broadcast(tensor, src): + if tensor.dtype == torch.float32: + tensor.fill_(3.14) + elif tensor.dtype == torch.long: + tensor.fill_(42) + + mocker.patch("torch.distributed.broadcast", side_effect=mock_broadcast) + + model = nn.Module() + model.linear = nn.Linear(2, 2) + model.register_buffer("buf", torch.zeros(2)) + model.register_buffer("int_buf", torch.zeros(3, dtype=torch.long)) + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace(use_hsdp=True), + enable_broadcast_weight_load=True, + quantization_config=None, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + loader._broadcast_model_weights(model, target_device=torch.device("cpu"), src_rank=0) + + assert torch.allclose(model.linear.weight, torch.tensor(3.14)) + assert torch.allclose(model.linear.bias, torch.tensor(3.14)) + assert torch.allclose(model.buf, torch.tensor(3.14)) + assert torch.equal(model.int_buf, torch.tensor([42, 42, 42], dtype=torch.long)) + + +def test_broadcast_model_weights_mixed_cpu_and_accelerator_devices(mocker): + """Ensure that when some pipeline modules are on CPU (transformer) and others on accelerator (VAE/encoder), + broadcasting normalizes tensors to CPU before concatenation and succeeds without device mismatch. + """ + dev = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu") + mocker.patch("torch.distributed.is_initialized", return_value=True) + mocker.patch("torch.distributed.get_world_size", return_value=2) + mocker.patch("torch.distributed.get_rank", return_value=0) + mocker.patch("torch.distributed.get_backend", return_value="nccl" if dev.type == "cuda" else "gloo") + mocker.patch("torch.distributed.broadcast") + mocker.patch("torch.distributed.barrier") + + model = nn.Module() + # Transformer parameters on CPU + model.transformer = nn.Linear(4, 4) + # VAE / Encoder parameters on accelerator + model.vae = nn.Linear(4, 4).to(dev) + model.text_encoder = nn.Linear(4, 4).to(dev) + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace(use_hsdp=True), + enable_broadcast_weight_load=True, + quantization_config=None, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + # Should not raise RuntimeError: Tensors must have same device + loader._broadcast_model_weights(model, target_device=dev, src_rank=0) + + +def test_broadcast_model_weights_mixed_devices_receiver(mocker): + """Ensure receiver rank with mixed CPU/accelerator tensors correctly receives broadcast data.""" + dev = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu") + mocker.patch("torch.distributed.is_initialized", return_value=True) + mocker.patch("torch.distributed.get_world_size", return_value=2) + mocker.patch("torch.distributed.get_rank", return_value=1) + mocker.patch("torch.distributed.get_backend", return_value="nccl" if dev.type == "cuda" else "gloo") + mocker.patch("torch.distributed.barrier") + + def mock_broadcast(tensor, src): + tensor.fill_(7.0) + + mocker.patch("torch.distributed.broadcast", side_effect=mock_broadcast) + + model = nn.Module() + model.transformer = nn.Linear(2, 2) # CPU + model.vae = nn.Linear(2, 2).to(dev) # Accelerator + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace(use_hsdp=True), + enable_broadcast_weight_load=True, + quantization_config=None, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + loader._broadcast_model_weights(model, target_device=dev, src_rank=0) + + assert torch.allclose(model.transformer.weight, torch.tensor(7.0)) + assert torch.allclose(model.vae.weight, torch.tensor(7.0, device=dev)) + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_hsdp_broadcast_weight_load_falls_back_when_online_quantization_enabled(mocker, rank): + """Ensure that when online quantization is enabled, HSDP weight loading falls back to + ordinary per-rank loading on all ranks rather than rank-0 broadcast. + """ + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + from vllm_omni.diffusion.offloader.module_collector import PipelineModules + + mocker.patch("torch.distributed.is_initialized", return_value=True) + mocker.patch("torch.distributed.get_world_size", return_value=4) + mocker.patch("torch.distributed.get_rank", return_value=rank) + + mock_quant_config = SimpleNamespace( + is_checkpoint_quantized=False, + ) + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace( + use_hsdp=True, + hsdp_replicate_size=1, + hsdp_shard_size=4, + ), + enable_broadcast_weight_load=True, + lora_backend=None, + lora_path=None, + quantization_config=mock_quant_config, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + events: list[str] = [] + + loader._init_from_load_format = mocker.Mock(return_value=model) # type: ignore[method-assign] + loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load")) # type: ignore[method-assign] + loader._broadcast_model_weights = mocker.Mock( # type: ignore[method-assign] + side_effect=lambda _model, **_kwargs: events.append("broadcast") + ) + loader._process_weights_after_loading = mocker.Mock( # type: ignore[method-assign] + side_effect=lambda _model, _device: events.append("process") + ) + mocker.patch.object( + loader_mod.ModuleDiscovery, + "discover", + return_value=PipelineModules( + dits=[model.transformer], + dit_names=["transformer"], + vaes=[], + encoders=[], + encoder_names=[], + resident_modules=[], + resident_names=[], + ), + ) + mocker.patch.object( + loader_mod, + "apply_hsdp_to_model", + side_effect=lambda *_args, **_kwargs: events.append("shard"), + ) + + loader._load_model_with_hsdp(torch.device("cpu")) + + # Both rank 0 and non-zero ranks MUST execute ordinary load_weights, not broadcast + assert "load" in events + assert "broadcast" not in events + assert events == ["load", "process", "shard"] + + +def _mp_worker_online_quant(rank: int, world_size: int, rendezvous: str, temp_dir: str) -> None: + torch.set_num_threads(1) + import torch.distributed as dist + + from vllm_omni.diffusion.offloader.module_collector import PipelineModules + + dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=world_size) + try: + + class _MockOnlineQuantModel(nn.Module): + def __init__(self): + super().__init__() + self.transformer = nn.Linear(4, 4) + self.load_called = False + + def load_weights(self, weights): + self.load_called = True + return {"transformer.weight", "transformer.bias"} + + mock_quant_config = SimpleNamespace(is_checkpoint_quantized=False) + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace( + use_hsdp=True, + hsdp_replicate_size=1, + hsdp_shard_size=world_size, + ), + enable_broadcast_weight_load=True, + lora_backend=None, + lora_path=None, + quantization_config=mock_quant_config, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + model = _MockOnlineQuantModel() + loader._init_from_load_format = lambda *args, **kwargs: model # type: ignore[method-assign] + loader._process_weights_after_loading = lambda *args, **kwargs: None # type: ignore[method-assign] + + with ( + patch.object( + loader_module.ModuleDiscovery, + "discover", + return_value=PipelineModules( + dits=[model.transformer], + dit_names=["transformer"], + vaes=[], + encoders=[], + encoder_names=[], + resident_modules=[], + resident_names=[], + ), + ), + patch.object(loader_module, "apply_hsdp_to_model", return_value=None), + ): + loader._load_model_with_hsdp(torch.device("cpu")) + + if model.load_called: + with open(os.path.join(temp_dir, f"rank_{rank}_success.flag"), "w") as f: + f.write("ok") + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif( + not torch.distributed.is_available() or not torch.distributed.is_gloo_available(), + reason="requires torch.distributed gloo", +) +def test_hsdp_broadcast_weight_load_online_quant_multiprocess(): + """Multi-process regression test verifying that with online quantization and broadcast enabled, + all ranks fall back to independent load_weights in parallel processes without hanging. + """ + import torch.multiprocessing as mp + + with tempfile.TemporaryDirectory() as temp_dir: + rendezvous = f"file://{os.path.join(temp_dir, 'gloo-rendezvous')}" + world_size = 2 + with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}): + mp.spawn( + _mp_worker_online_quant, + args=(world_size, rendezvous, temp_dir), + nprocs=world_size, + join=True, + ) + 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")) diff --git a/tests/engine/test_stage_engine_args.py b/tests/engine/test_stage_engine_args.py index 250aff88464..f2053828447 100644 --- a/tests/engine/test_stage_engine_args.py +++ b/tests/engine/test_stage_engine_args.py @@ -82,6 +82,7 @@ def _effective_backend_values(config_cls: type, engine_args: dict) -> dict[str, "custom_voice_dir", "devices", "disable_autocast", + "enable_broadcast_weight_load", "enable_multithread_weight_load", "env", "has_sampling_extra_args", diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 2b016be0d16..f0752077bb4 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -165,6 +165,7 @@ class _ModelEngineOverrides(TypedDict, total=False): max_cudagraph_capture_size: int enable_flashinfer_autotune: bool enable_multithread_weight_load: bool + enable_broadcast_weight_load: bool num_weight_load_threads: int disable_autocast: bool @@ -435,6 +436,7 @@ class OmniStageModelConfig(_TrackExplicitConfigFields): max_cudagraph_capture_size: int | None = Field(default=None, ge=0) enable_flashinfer_autotune: bool | None = None enable_multithread_weight_load: bool = True + enable_broadcast_weight_load: bool = False num_weight_load_threads: int = Field(default=4, ge=1) disable_autocast: bool = False # Per-stage checkpoint/tokenizer subdirectories under the model root @@ -1028,6 +1030,7 @@ def enrich_config(self) -> None: "enable_sleep_mode", "enforce_eager", "enable_multithread_weight_load", + "enable_broadcast_weight_load", "num_weight_load_threads", "disable_autocast", } diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 0e962b4cc30..063a95b0e16 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -472,6 +472,7 @@ class StageDeployConfig: # Runtime optimizations used by diffusion loading/execution. enable_multithread_weight_load: bool | None = None + enable_broadcast_weight_load: bool | None = None num_weight_load_threads: int | None = None diffusion_offload_config: dict[str, Any] | None = None # Compatibility aliases for existing callers and model-specific stage diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index d27cc36e191..7a7de3bb1d6 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -920,6 +920,7 @@ class OmniDiffusionConfig: # Parallel weight loading (for faster diffusion model startup) enable_multithread_weight_load: bool = True + enable_broadcast_weight_load: bool = False num_weight_load_threads: int = 4 # Enable sleep mode diff --git a/vllm_omni/diffusion/lora/loader.py b/vllm_omni/diffusion/lora/loader.py index a13363c7a77..eb3f35188ea 100644 --- a/vllm_omni/diffusion/lora/loader.py +++ b/vllm_omni/diffusion/lora/loader.py @@ -43,6 +43,7 @@ def _prepare_lora_delta( lora_a_suffix: str = "lora_A.weight", lora_b_suffix: str = "lora_B.weight", lora_bias_suffix: str = "bias", + compute_device: torch.device | None = None, ): used_keys: set[str] = set() # stacked_params_mapping param_to_weight_names @@ -70,6 +71,8 @@ def _prepare_lora_delta( if lora_bias_key not in lora_state_dict: return None, used_keys delta = lora_state_dict[lora_bias_key] + if compute_device is not None and compute_device.type != "cpu": + delta = delta.to(compute_device) stacked_deltas.append(delta) used_keys.add(lora_bias_key) else: @@ -80,6 +83,9 @@ def _prepare_lora_delta( return None, used_keys a = lora_state_dict[lora_a_key] b = lora_state_dict[lora_b_key] + if compute_device is not None and compute_device.type != "cpu": + a = a.to(compute_device) + b = b.to(compute_device) delta = torch.matmul(b, a) stacked_deltas.append(delta) used_keys.add(lora_a_key) @@ -94,7 +100,10 @@ def _prepare_lora_delta( if lora_bias_key not in lora_state_dict: return None, used_keys used_keys.add(lora_bias_key) - return lora_state_dict[lora_bias_key], used_keys + delta = lora_state_dict[lora_bias_key] + if compute_device is not None and compute_device.type != "cpu": + delta = delta.to(compute_device) + return delta, used_keys lora_a_key = f"{base_key}.{lora_a_suffix}" lora_b_key = f"{base_key}.{lora_b_suffix}" @@ -102,9 +111,13 @@ def _prepare_lora_delta( return None, used_keys a = lora_state_dict[lora_a_key] b = lora_state_dict[lora_b_key] + if compute_device is not None and compute_device.type != "cpu": + a = a.to(compute_device) + b = b.to(compute_device) + delta = torch.matmul(b, a) used_keys.add(lora_a_key) used_keys.add(lora_b_key) - return torch.matmul(b, a), used_keys + return delta, used_keys def _load_lora_state_dict( @@ -267,6 +280,7 @@ def load_lora_into_module( else: continue + compute_dev = params.device if params.device.type != "cpu" else None delta, used_keys = _prepare_lora_delta( state_dict, base_key, @@ -275,6 +289,7 @@ def load_lora_into_module( lora_a_suffix, lora_b_suffix, lora_bias_suffix, + compute_device=compute_dev, ) if delta is None: continue @@ -319,6 +334,7 @@ def unload_module_lora( else: continue + compute_dev = param.device if param.device.type != "cpu" else None delta, used_keys = _prepare_lora_delta( state_dict, base_key, @@ -327,6 +343,7 @@ def unload_module_lora( lora_a_suffix, lora_b_suffix, lora_bias_suffix, + compute_device=compute_dev, ) if delta is None: continue diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 27774aa3724..78247e7b968 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -509,11 +509,6 @@ def _maybe_fuse_distilled_lora(self, model: nn.Module) -> None: if getattr(model, "lora_is_fused", False): return - # A warm HWR restore already brings back the fused final-layout weights. - if self._hwr_state is not None and self._hwr_state.get("warm_snapshot") is not None: - setattr(model, "lora_is_fused", True) - return - lora_path = getattr(self.od_config, "lora_path", None) if not lora_path: return @@ -531,6 +526,111 @@ def _maybe_fuse_distilled_lora(self, model: nn.Module) -> None: else: logger.warning("Pipeline %s does not support loading distilled LoRA weights.", model.__class__.__name__) + def _broadcast_model_weights( + self, + model: nn.Module, + target_device: torch.device, + src_rank: int = 0, + bucket_cap_bytes: int = 256 * 1024 * 1024, + ) -> None: + """Broadcast model parameters and buffers from src_rank to all other ranks using coalesced bucketing.""" + if not torch.distributed.is_initialized() or torch.distributed.get_world_size() <= 1: + return + + rank = torch.distributed.get_rank() + world_size = torch.distributed.get_world_size() + backend = torch.distributed.get_backend() + is_cpu_backend = backend == "gloo" + + logger.info( + "Worker %d: %s full model weights via %s across %d ranks", + rank, + "Broadcasting" if rank == src_rank else "Receiving", + backend, + world_size, + ) + t0 = time.perf_counter() + + all_tensors = [p.data for _, p in model.named_parameters() if p.numel() > 0] + [ + b.data for _, b in model.named_buffers() if b.numel() > 0 + ] + if not all_tensors: + return + + count = len(all_tensors) + total_bytes = sum(t.numel() * t.element_size() for t in all_tensors) + + # Group tensors into buckets by dtype and maximum byte size + buckets: list[list[torch.Tensor]] = [] + curr_bucket: list[torch.Tensor] = [] + curr_bytes = 0 + + for tensor in all_tensors: + t_bytes = tensor.numel() * tensor.element_size() + if curr_bucket and (curr_bucket[0].dtype != tensor.dtype or curr_bytes + t_bytes > bucket_cap_bytes): + buckets.append(curr_bucket) + curr_bucket = [] + curr_bytes = 0 + curr_bucket.append(tensor) + curr_bytes += t_bytes + + if curr_bucket: + buckets.append(curr_bucket) + + for bucket in buckets: + tot_numel = sum(t.numel() for t in bucket) + dtype = bucket[0].dtype + + if is_cpu_backend: + if rank == src_rank: + flat_cpu = torch.cat([t.detach().to("cpu").reshape(-1) for t in bucket]) + else: + flat_cpu = torch.empty(tot_numel, dtype=dtype) + torch.distributed.broadcast(flat_cpu, src=src_rank) + if rank != src_rank: + offset = 0 + for t in bucket: + n = t.numel() + t.copy_(flat_cpu[offset : offset + n].view_as(t)) + offset += n + del flat_cpu + else: + if rank == src_rank: + flat_cpu = torch.cat([t.detach().to("cpu").reshape(-1) for t in bucket]) + dev_flat = flat_cpu.to(target_device, non_blocking=False) + del flat_cpu + torch.distributed.broadcast(dev_flat, src=src_rank) + del dev_flat + else: + dev_flat = torch.empty(tot_numel, dtype=dtype, device=target_device) + torch.distributed.broadcast(dev_flat, src=src_rank) + flat_cpu = dev_flat.to("cpu", non_blocking=False) + del dev_flat + offset = 0 + for t in bucket: + n = t.numel() + t.copy_(flat_cpu[offset : offset + n].view_as(t), non_blocking=False) + offset += n + del flat_cpu + + if not is_cpu_backend and target_device.type != "cpu": + from vllm_omni.platforms import current_omni_platform + + current_omni_platform.synchronize() + torch.distributed.barrier() + + elapsed = time.perf_counter() - t0 + gb = total_bytes / 1e9 + logger.info( + "Worker %d: Shared weight broadcast complete (%d tensors, %d buckets, %.2f GB) in %.2fs (%.2f GB/s)", + rank, + count, + len(buckets), + gb, + elapsed, + gb / max(elapsed, 1e-6), + ) + def load_model( self, load_device: str, @@ -654,6 +754,10 @@ def load_model( if hwr_active and hwr_state is not None: self.host_weight_plan = cast(HostWeightPlan | None, hwr_state.get("plan")) if dit_distributed_offload and not hwr_active and not self._force_canonical_load: + lora_backend = getattr(self.od_config, "lora_backend", None) + has_distilled_lora = lora_backend in (LoRABackend.DISTILL, "distill") and bool( + getattr(self.od_config, "lora_path", None) + ) plan_result = build_checkpoint_mmap_plan( model, dit_modules=tuple(zip(modules.dit_names, modules.dits)), @@ -662,6 +766,7 @@ def load_model( tensor_parallel_size=tensor_parallel_size, use_hsdp=use_hsdp, online_quantization=any(self._has_online_quant(dit) for dit in modules.dits), + has_distilled_lora=has_distilled_lora, ) self.host_weight_plan = plan_result.plan @@ -689,7 +794,6 @@ def load_model( sources=ordinary_sources, planned_weights=host_weight_plan.bindings, ) - self._maybe_fuse_distilled_lora(model) else: if dit_distributed_offload and plan_result is not None: logger.info( @@ -1097,8 +1201,39 @@ def _load_model_with_hsdp( if load_format == "diffusers": raise ValueError("HSDP is not supported with the diffusers adapter load format") model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True) - self.load_weights(model) - self._maybe_fuse_distilled_lora(model) + world_size = 1 + rank = 0 + if torch.distributed.is_initialized(): + world_size = torch.distributed.get_world_size() + rank = torch.distributed.get_rank() + + has_online_quant = self._has_online_quant(model) or ( + self.quant_config is not None and not getattr(self.quant_config, "is_checkpoint_quantized", False) + ) + enable_broadcast = bool(getattr(self.od_config, "enable_broadcast_weight_load", False)) and world_size > 1 + + if enable_broadcast and has_online_quant: + logger.info( + "Worker %d: Online quantization detected; falling back to ordinary per-rank weight loading for HSDP", + rank, + ) + enable_broadcast = False + + if enable_broadcast: + if rank == 0: + self.load_weights(model) + self._maybe_fuse_distilled_lora(model) + self._broadcast_model_weights(model, target_device=target_device, src_rank=0) + if ( + rank != 0 + and getattr(self.od_config, "lora_backend", None) in (LoRABackend.DISTILL, "distill") + and getattr(self.od_config, "lora_path", None) + and hasattr(model, "load_lora_weights") + ): + setattr(model, "lora_is_fused", True) + else: + self.load_weights(model) + self._maybe_fuse_distilled_lora(model) # Quantization methods must finish while parameters are ordinary local # tensors. Some post-load transforms use operations (for example, diff --git a/vllm_omni/diffusion/model_loader/host_weight_plan.py b/vllm_omni/diffusion/model_loader/host_weight_plan.py index e6301b0ece6..f0ea8024d75 100644 --- a/vllm_omni/diffusion/model_loader/host_weight_plan.py +++ b/vllm_omni/diffusion/model_loader/host_weight_plan.py @@ -322,6 +322,7 @@ def build_checkpoint_mmap_plan( tensor_parallel_size: int, use_hsdp: bool, online_quantization: bool, + has_distilled_lora: bool = False, ) -> HostWeightPlanResult: """Build a complete direct-checkpoint plan or return a fallback reason.""" if tensor_parallel_size != 1: @@ -330,6 +331,8 @@ def build_checkpoint_mmap_plan( 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): diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index ee555275189..20319de866e 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -572,6 +572,7 @@ class OrchestratorArgs: vae_use_slicing: bool = False vae_use_tiling: bool = False enable_multithread_weight_load: bool = True + enable_broadcast_weight_load: bool = False num_weight_load_threads: int = 4 diffusion_offload_config: dict[str, Any] | None = None # Compatibility aliases for existing callers and model-specific stage diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 598cbe89328..f9297aa0c20 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -726,6 +726,13 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu default=True, help="Disable multi-threaded safetensors loading (default: enabled with 4 threads).", ) + omni_config_group.add_argument( + "--enable-broadcast-weight-load", + action="store_true", + dest="enable_broadcast_weight_load", + default=False, + help="Enable Rank-0 shared weight broadcast across workers for HSDP (default: disabled).", + ) omni_config_group.add_argument( "--num-weight-load-threads", type=int, From a7ab91fb85b1bb513d34f0a19c3289c55a38b013 Mon Sep 17 00:00:00 2001 From: SuyanLi <126558907+suyanli220@users.noreply.github.com> Date: Fri, 11 Sep 2026 16:47:25 +0800 Subject: [PATCH 70/93] [Example] Migrate HunyuanImage-3.0 to model_extras + shared task examples (#5559) Signed-off-by: suyanli220 Signed-off-by: suyan.li Co-authored-by: Claude Opus 4.8 Co-authored-by: suyan.li Co-authored-by: Hongsheng Liu --- .../offline_inference/hunyuan_image3.md | 35 -- .../hunyuan_image3/README.md | 154 -------- .../hunyuan_image3/end2end.py | 321 ----------------- .../image_to_image/image_edit.py | 126 +++++++ .../text_to_image/text_to_image.py | 126 +++++++ recipes/Tencent/HunyuanImage-3.0-Instruct.md | 52 ++- .../hunyuan_image3/test_prompt_utils.py | 129 ++++++- .../offline_inference/test_hunyuanimage3.py | 122 ++++++- .../openai_api/test_image_server.py | 69 ++++ ...test_serving_chat_multistage_generation.py | 81 +++++ tests/model_extras/test_model_extras.py | 241 +++++++++++++ .../test_shared_script_ar_integration.py | 340 ++++++++++++++++++ .../models/hunyuan_image3/prompt_utils.py | 136 +++++++ vllm_omni/entrypoints/openai/serving_chat.py | 11 +- vllm_omni/model_extras/__init__.py | 4 + vllm_omni/model_extras/hunyuan_image3.py | 171 ++++++++- vllm_omni/model_extras/registry.py | 53 +++ 17 files changed, 1627 insertions(+), 544 deletions(-) delete mode 100644 docs/user_guide/examples/offline_inference/hunyuan_image3.md delete mode 100644 examples/offline_inference/hunyuan_image3/README.md delete mode 100644 examples/offline_inference/hunyuan_image3/end2end.py create mode 100644 tests/model_extras/test_shared_script_ar_integration.py diff --git a/docs/user_guide/examples/offline_inference/hunyuan_image3.md b/docs/user_guide/examples/offline_inference/hunyuan_image3.md deleted file mode 100644 index 636aa0c11d5..00000000000 --- a/docs/user_guide/examples/offline_inference/hunyuan_image3.md +++ /dev/null @@ -1,35 +0,0 @@ -# HunyuanImage-3.0 Image-to-Text Inference - -Source . - - -This example demonstrates how to run HunyuanImage-3.0 Image-to-Text with the vLLM-Omni. - -## Local CLI Usage - -Download the example image: - -```bash -wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg -``` - -Run example: - -```bash -python image_to_text.py \ - --image cherry_blossom.jpg \ - --prompt "<|startoftext|>You are an assistant that understands images and outputs text.Describe the content of the picture." -``` - -Key arguments: - -- `--model`: Model used. Default is: tencent/HunyuanImage-3.0-Instruct (Optional). -- `--image`: Path to input image (required). -- `--prompt`: Text description used to guide image understanding (required). - -## Example materials - -??? abstract "image_to_text.py" - ``````py - --8<-- "examples/offline_inference/hunyuan_image3/image_to_text.py" - `````` diff --git a/examples/offline_inference/hunyuan_image3/README.md b/examples/offline_inference/hunyuan_image3/README.md deleted file mode 100644 index b4d5af2c7b9..00000000000 --- a/examples/offline_inference/hunyuan_image3/README.md +++ /dev/null @@ -1,154 +0,0 @@ -# HunyuanImage-3.0-Instruct - -This example runs HunyuanImage-3.0-Instruct offline with the unified deploy -YAMLs under `vllm_omni/deploy/`. - -## Deploy Configs - -| File | Topology | Default use | -| :--- | :--- | :--- | -| `vllm_omni/deploy/hunyuan_image_3_moe.yaml` | AR + DiT | Default for `text2img` and `img2img`. | -| `vllm_omni/deploy/hunyuan_image3_ar.yaml` | AR only | Default for `img2text` and `text2text`. | -| `vllm_omni/deploy/hunyuan_image3_dit.yaml` | DiT only | Standalone diffusion stage. Pass it explicitly with `--deploy-config`. | - -The example chooses a deploy config automatically when `--deploy-config` is -omitted: - -| `--modality` | `mode` passed to Omni | Default deploy | -| :--- | :--- | :--- | -| `text2img` | `text-to-image` | `hunyuan_image_3_moe.yaml` | -| `img2img` | `image-editing` | `hunyuan_image_3_moe.yaml` | -| `img2text` | `image-to-text` | `hunyuan_image3_ar.yaml` | -| `text2text` | `text-to-text` | `hunyuan_image3_ar.yaml` | - -`--modality` is an offline example convenience flag. It maps to the internal -`mode` argument passed to `Omni(...)` by this script. HunyuanImage3 uses -separate deploy YAMLs for AR + DiT, AR-only, and DiT-only topologies, so the -stage topology is selected by the deploy file rather than by YAML mode -overrides. - -Online serving does not expose a `--modality` flag or accept `mode` as an API -request field. Choose the deploy topology when starting the server with -`--deploy-config`, then use the OpenAI-compatible endpoint and request shape for -the scenario. The `modalities` request field is used by the chat completions -path; the image endpoints infer the image task from the endpoint and payload. - -| Online scenario | Server deploy | Request | -| :--- | :--- | :--- | -| Text to image | `--deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml` | `POST /v1/images/generations`, or `POST /v1/chat/completions` with `"modalities": ["image"]`. | -| Image editing | `--deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml` | `POST /v1/images/edits`. | -| Image/text to text | `--deploy-config vllm_omni/deploy/hunyuan_image3_ar.yaml` | `POST /v1/chat/completions` for text output, for example with `"modalities": ["text"]`. | -| DiT-only image generation | `--deploy-config vllm_omni/deploy/hunyuan_image3_dit.yaml` | `POST /v1/images/generations`. | - -## Run Examples - -Text to image, using the default AR + DiT deploy: - -```bash -python examples/offline_inference/hunyuan_image3/end2end.py \ - --model tencent/HunyuanImage-3.0-Instruct \ - --modality text2img \ - --prompts "A cute cat sitting on a windowsill watching the sunset" -``` - -Image editing, using the default AR + DiT deploy: - -```bash -python examples/offline_inference/hunyuan_image3/end2end.py \ - --model tencent/HunyuanImage-3.0-Instruct \ - --modality img2img \ - --image-path /path/to/image.png \ - --prompts "Make the petals neon pink" -``` - -Image to text, using the AR-only deploy: - -```bash -python examples/offline_inference/hunyuan_image3/end2end.py \ - --model tencent/HunyuanImage-3.0-Instruct \ - --modality img2text \ - --image-path /path/to/image.jpg \ - --prompts "Describe the content of the picture." -``` - -Text to text, using the AR-only deploy: - -```bash -python examples/offline_inference/hunyuan_image3/end2end.py \ - --model tencent/HunyuanImage-3.0-Instruct \ - --modality text2text \ - --prompts "What is the capital of France?" -``` - -Standalone DiT, using the DiT-only deploy explicitly: - -```bash -python examples/offline_inference/hunyuan_image3/end2end.py \ - --model tencent/HunyuanImage-3.0-Instruct \ - --modality text2img \ - --deploy-config vllm_omni/deploy/hunyuan_image3_dit.yaml \ - --prompts "A cinematic portrait of an astronaut in a greenhouse" -``` - -Override the default full AR + DiT deploy explicitly: - -```bash -python examples/offline_inference/hunyuan_image3/end2end.py \ - --model tencent/HunyuanImage-3.0-Instruct \ - --modality text2img \ - --deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml \ - --prompts "A cute cat" -``` - -## Additional Config - -You can pass diffusion worker `additional_config` from the offline example as a JSON object. -This maps to the upstream vLLM `VllmConfig.additional_config` platform extension field: -https://docs.vllm.ai/en/stable/api/vllm/config/#vllm.config.VllmConfig.additional_config - -```bash -python end2end.py --modality text2img \ - --prompts "A cute cat" \ - --additional-config '{"torchair_graph_config":{"enabled":true}}' -``` - - -## Key Arguments - -| Argument | Description | -| :--- | :--- | -| `--deploy-config` | Optional path for unified deploy YAML overrides. | -| `--additional-config` | JSON object forwarded to diffusion worker `additional_config`. | -| `--modality` | Offline-only convenience flag. One of `text2img`, `img2img`, `img2text`, `text2text`. It selects prompt formatting, internal `mode`, and default deploy config for this script. Online serving uses `--deploy-config` plus the endpoint and, for chat completions, request `modalities` instead. | -| `--steps` | Number of diffusion inference steps for image generation. | -| `--guidance-scale` | Classifier-free guidance scale for image generation. | -| `--height`, `--width` | Output image size for `text2img`. | -| `--bot-task` | Override prompt mode. `none`, `think`, `recaption`, `think_recaption`, or `vanilla`. | -| `--sys-type` | Override the system prompt type, for example `en_unified` or `en_vanilla`. | -| `--vae-use-tiling` | Enable VAE tiling for memory reduction. | - -## Notes - -- `hunyuan_image3_ar.yaml` is a 4-card AR-only text/comprehension deploy. -- `hunyuan_image3_dit.yaml` is a single-stage DiT deploy with `stage_id: 0`. -- Platform-specific overrides are included in the unified deploy YAMLs. - -## Prompt Format - -HunyuanImage-3.0-Instruct uses an instruct chat template: - -```text -<|startoftext|>{system_prompt} - -User: {?}{user_prompt} - -Assistant: {trigger_tag?} -``` - -- ``: Placeholder for each input image (single token; expanded by the multimodal pipeline). -- Trigger tags: `` for CoT and `` for recaptioning, placed after `Assistant: `. -- System prompt: Auto-selected from `task` and `bot_task`. -- `bot_task='vanilla'` with `task='t2i'` uses the bare pretrain template. - -The shared `vllm_omni.diffusion.models.hunyuan_image3.prompt_utils.build_prompt_tokens()` -helper handles segment-by-segment tokenization and matches HF `apply_chat_template`. diff --git a/examples/offline_inference/hunyuan_image3/end2end.py b/examples/offline_inference/hunyuan_image3/end2end.py deleted file mode 100644 index fbe4289befb..00000000000 --- a/examples/offline_inference/hunyuan_image3/end2end.py +++ /dev/null @@ -1,321 +0,0 @@ -""" -HunyuanImage-3.0-Instruct unified end-to-end inference script. -""" - -import argparse -import json -import os -from pathlib import Path - -from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( - MAX_IMAGES_PER_REQUEST, - build_prompt_tokens, - resolve_stop_token_ids, - resolve_sys_type, -) -from vllm_omni.entrypoints.omni import Omni -from vllm_omni.inputs.data import OmniPromptType - -_REPO_ROOT = Path(__file__).resolve().parents[3] -_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image_3_moe.yaml") -_DEFAULT_AR_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3_ar.yaml") - -_MODALITY_TASK_MAP: dict[str, tuple[str, str | None]] = { - "text2img": ("t2i", "think"), - "img2img": ("it2i", "think"), - "img2text": ("i2t", None), - "text2text": ("t2t", None), -} - -_MODALITY_DEFAULT_DEPLOY_CONFIG = { - "text2img": _DEFAULT_DEPLOY_CONFIG, - "img2img": _DEFAULT_DEPLOY_CONFIG, - "img2text": _DEFAULT_AR_DEPLOY_CONFIG, - "text2text": _DEFAULT_AR_DEPLOY_CONFIG, -} - -_MODALITY_MODE = { - "text2img": "text-to-image", - "img2img": "image-editing", - "img2text": "image-to-text", - "text2text": "text-to-text", -} - - -def parse_args(): - parser = argparse.ArgumentParser(description="HunyuanImage-3.0-Instruct end-to-end inference.") - parser.add_argument("--model", default="tencent/HunyuanImage-3.0-Instruct", help="Model name or local path.") - parser.add_argument( - "--modality", - default="text2img", - choices=list(_MODALITY_TASK_MAP), - ) - parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.") - parser.add_argument( - "--image-path", - type=str, - default=None, - help="Input image path(s) for img2img/img2text. Comma-separated for multi-image (up to 3).", - ) - parser.add_argument("--output", type=str, default=".", help="Output directory to save results.") - parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.") - parser.add_argument("--guidance-scale", type=float, default=5.0, help="Classifier-free guidance scale.") - parser.add_argument("--seed", type=int, default=42, help="Random seed.") - parser.add_argument("--height", type=int, default=None, help="Output image height.") - parser.add_argument("--width", type=int, default=None, help="Output image width.") - parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.") - parser.add_argument( - "--bot-task", - type=str, - default=None, - choices=["none", "think", "recaption", "think_recaption", "vanilla"], - help="Override prompt mode. Default: auto from --modality.", - ) - parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.") - parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.") - parser.add_argument( - "--stream", - action="store_true", - default=False, - help="Enable streaming CoT display; print AR text token-by-token in real time.", - ) - parser.add_argument("--log-stats", action="store_true", default=False) - parser.add_argument("--init-timeout", type=int, default=300, help="Initialization timeout in seconds.") - parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.") - parser.add_argument( - "--diffusion-kv-cache-dtype", - type=str, - default=None, - help="Diffusion attention KV cache dtype, for example 'fp8'. Separate from vLLM --kv-cache-dtype.", - ) - parser.add_argument( - "--diffusion-kv-cache-skip-steps", - type=str, - default=None, - help="Denoising step selector to keep diffusion KV cache in native dtype, for example '0,1,4-6'.", - ) - parser.add_argument( - "--diffusion-kv-cache-skip-layers", - type=str, - default=None, - help="Transformer layer selector to keep diffusion KV cache in native dtype, for example '0-2,10'.", - ) - parser.add_argument( - "--additional-config", - type=str, - default=None, - help=( - "JSON object forwarded to Omni/additional_config, for example " - '\'{"torchair_graph_config":{"enabled":true}}\'. Different platforms may support different ' - "configs. Make sure the configs are valid for the platform you are using. " - "Contents must be hashable." - ), - ) - - return parser.parse_args() - - -def parse_additional_config(raw_value: str | None) -> dict | None: - """Parse a JSON string into an additional_config mapping.""" - if raw_value is None: - return None - - try: - additional_config = json.loads(raw_value) - except json.JSONDecodeError as exc: - raise ValueError(f"Invalid --additional-config JSON: {exc}") from exc - - if additional_config is None: - return None - if not isinstance(additional_config, dict): - raise ValueError(f"--additional-config must decode to a JSON object, got {type(additional_config).__name__}") - return additional_config - - -def main(): - args = parse_args() - os.makedirs(args.output, exist_ok=True) - additional_config = parse_additional_config(args.additional_config) - - task, default_bot_task = _MODALITY_TASK_MAP[args.modality] - if args.bot_task is None: - bot_task: str | None = default_bot_task - elif args.bot_task == "none": - bot_task = None - else: - bot_task = args.bot_task - - deploy_config = args.deploy_config or _MODALITY_DEFAULT_DEPLOY_CONFIG[args.modality] - - omni_kwargs = { - "model": args.model, - "vae_use_tiling": args.vae_use_tiling, - "log_stats": args.log_stats, - "init_timeout": args.init_timeout, - "enforce_eager": args.enforce_eager, - "mode": _MODALITY_MODE[args.modality], - "diffusion_kv_cache_dtype": args.diffusion_kv_cache_dtype, - "diffusion_kv_cache_skip_steps": args.diffusion_kv_cache_skip_steps, - "diffusion_kv_cache_skip_layers": args.diffusion_kv_cache_skip_layers, - } - - if additional_config is not None: - omni_kwargs["additional_config"] = additional_config - omni_kwargs["deploy_config"] = deploy_config - - omni = Omni(**omni_kwargs) - - prompts = args.prompts or ["A cute cat"] - input_images: list = [] - if args.modality in ("img2img", "img2text"): - if not args.image_path: - raise ValueError(f"--image-path required for {args.modality}, got: {args.image_path}") - from PIL import Image - - image_paths = [p.strip() for p in args.image_path.split(",") if p.strip()] - if len(image_paths) > MAX_IMAGES_PER_REQUEST: - raise ValueError( - f"--image-path accepts at most {MAX_IMAGES_PER_REQUEST} images for " - f"HunyuanImage-3.0 IT2I, got {len(image_paths)}: {args.image_path}" - ) - for image_path in image_paths: - if not os.path.exists(image_path): - raise ValueError(f"Image path does not exist: {image_path}") - input_images.append(Image.open(image_path).convert("RGB")) - if not input_images: - raise ValueError(f"--image-path produced no usable paths: {args.image_path!r}") - - from transformers import AutoTokenizer - - tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) - mm_image_payload = (input_images[0] if len(input_images) == 1 else input_images) if input_images else None - - formatted_prompts: list[OmniPromptType] = [] - for prompt in prompts: - build_kwargs: dict = {"task": task, "bot_task": bot_task, "sys_type": args.sys_type} - if input_images: - build_kwargs["num_images"] = len(input_images) - result = build_prompt_tokens(prompt, tokenizer, **build_kwargs) - token_ids = result.token_ids - effective_sys_type = args.sys_type or resolve_sys_type(bot_task) - - prompt_dict: dict = { - "prompt_token_ids": token_ids, - "prompt": prompt, - "use_system_prompt": effective_sys_type, - } - if args.modality == "text2img": - prompt_dict["modalities"] = ["image"] - elif args.modality == "img2img": - prompt_dict["modalities"] = ["image"] - prompt_dict["multi_modal_data"] = {"image": mm_image_payload} - prompt_dict["height"] = input_images[0].height - prompt_dict["width"] = input_images[0].width - elif args.modality == "img2text": - prompt_dict["modalities"] = ["text"] - prompt_dict["multi_modal_data"] = {"image": mm_image_payload} - else: - prompt_dict["modalities"] = ["text"] - formatted_prompts.append(prompt_dict) - - from vllm.sampling_params import RequestOutputKind - - params_list = list(omni.default_sampling_params_list) - - from vllm_omni.inputs.data import OmniDiffusionSamplingParams - - if (args.height is None) != (args.width is None): - raise ValueError("--height and --width must both be specified or both omitted.") - user_specified_size = args.height is not None and args.width is not None - if args.modality in ("img2text", "text2text"): - ar_image_size = "auto" - elif user_specified_size: - ar_image_size = f"{args.width}x{args.height}" - else: - ar_image_size = None - ar_stop_token_ids = resolve_stop_token_ids( - task=task, bot_task=bot_task, tokenizer=tokenizer, image_size=ar_image_size - ) - print( - f"[AR Config] task={task}, bot_task={bot_task}, image_size={ar_image_size}, stop_token_ids={ar_stop_token_ids}" - ) - for sp in params_list: - if isinstance(sp, OmniDiffusionSamplingParams): - sp.num_inference_steps = args.steps - sp.guidance_scale = args.guidance_scale - sp.guidance_scale_provided = True - if args.seed is not None: - sp.seed = args.seed - if args.modality == "text2img": - sp.height = args.height - sp.width = args.width - elif hasattr(sp, "stop_token_ids"): - sp.stop_token_ids = ar_stop_token_ids - # When --stream is set, request DELTA output from the AR stage - # so we can display CoT text token-by-token in real time. - if args.stream and hasattr(sp, "output_kind"): - sp.output_kind = RequestOutputKind.DELTA - - print(f"\n{'=' * 60}") - print("HunyuanImage-3.0 Generation Configuration:") - print(f" Model: {args.model}") - print(f" Modality: {args.modality}") - print(f" Prompt task: {task}") - print(f" Bot task: {bot_task}") - print(f" Deploy config: {deploy_config}") - print(f" Num stages: {omni.num_stages}") - if args.modality in ("text2img", "img2img"): - print(f" Inference steps: {args.steps}") - print(f" Guidance scale: {args.guidance_scale}") - print(f" Seed: {args.seed}") - print(f" diffusion_kv_cache_dtype: {args.diffusion_kv_cache_dtype}") - print(f" diffusion_kv_cache_skip_steps: {args.diffusion_kv_cache_skip_steps}") - print(f" diffusion_kv_cache_skip_layers: {args.diffusion_kv_cache_skip_layers}") - if args.modality == "text2img": - print(f" Output size: {args.width}x{args.height}") - if args.image_path: - print(f" Input image: {args.image_path}") - if additional_config is not None: - print(f" Additional config: {additional_config}") - print(f" Prompts: {prompts}") - print(f"{'=' * 60}\n") - - # When --stream is set, print AR CoT text token-by-token in real time. - # Otherwise, collect and print the full AR text once when stage 0 finishes. - omni_outputs = omni.generate( - prompts=formatted_prompts, - sampling_params_list=params_list, - py_generator=True, - use_tqdm=False, - ) - img_idx = 0 - for req_output in omni_outputs: - ro = req_output - stage_id = getattr(req_output, "stage_id", None) - - # AR stage text — each CompletionOutput.text is already a delta when - # output_kind=DELTA, so we can print it directly (matching the pattern - # in serving_chat.py). - if stage_id == 0 and ro and getattr(ro, "outputs", None): - for o in ro.outputs: - text = getattr(o, "text", "") or "" - if text: - print(text, end="", flush=True) - # Non-streaming: one shot with full text — emit a trailing newline. - if not args.stream: - print(flush=True) - - # Collect images from diffusion stage - images = getattr(req_output, "images", None) - if not images and ro and hasattr(ro, "images"): - images = ro.images - if images: - for j, img in enumerate(images): - save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png") - img.save(save_path) - print(f"\n[Output] Saved image to {save_path}") - img_idx += 1 - - -if __name__ == "__main__": - main() diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index 136d99857ad..cec7014de5d 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -96,6 +96,7 @@ import torch from PIL import Image +from vllm_omni.diffusion.data import logger from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.omni import Omni @@ -103,6 +104,8 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.model_extras import ( build_image_to_image_prompt, + get_ar_input_builder, + get_ar_tokenizer_validator, get_extra_body_params, get_model_class_name, should_init_extra_args_for_non_diffusion_stages, @@ -110,6 +113,91 @@ from vllm_omni.platforms import current_omni_platform +def _apply_ar_stage_inputs( + ar_input_builder: Any, + *, + model: str, + prompt_text: str, + extra_body: dict[str, Any], + num_images: int, + height: int | None, + width: int | None, + prompt_dict: dict[str, Any], + sampling_params_list: list[Any], + text_output: bool = False, + trust_remote_code: bool = False, + validate_tokenizer: Any | None = None, + allow_tokenizer_fallback: bool = False, +) -> None: + """Apply a model's declared AR-stage inputs to the request in place. + + Mirrors the helper in ``text_to_image.py``: loads the tokenizer for + byte-for-byte HF-parity AR tokenization, asks the model's + ``ar_input_builder`` for the AR prefill + stop tokens, and writes them + onto ``prompt_dict`` and the non-diffusion (AR) stage params. + Model-agnostic -- only the declared-hook contract is assumed. + + ``trust_remote_code`` must be threaded in from the caller's own resolved + ``--trust-remote-code`` flag -- this is a separate tokenizer load outside + the main engine, so it needs the same explicit user opt-in rather than + defaulting to trusting whatever ``--model`` was passed. + + ``validate_tokenizer``, when the model declares one via + ``get_ar_tokenizer_validator``, is called on a successfully-loaded real + tokenizer -- outside the load's try/except, so a validation failure (a + model/tokenizer revision drifting from hardcoded special-token ids) + raises instead of being swallowed into the string-prompt fallback. + + ``allow_tokenizer_fallback`` (default ``False``, i.e. fail fast): a + tokenizer load failure normally raises, since silently degrading to the + string-prompt form can produce a request that completes with subtly + wrong stop tokens instead of surfacing the real problem (missing + ``--trust-remote-code``, network/cache issue, etc.). Set ``True`` only + for explicit unit tests or a deliberate offline-compat run where the + caller has already decided a degraded prompt is acceptable. + """ + try: + from transformers import AutoTokenizer + + ar_tokenizer: Any | None = AutoTokenizer.from_pretrained(model, trust_remote_code=trust_remote_code) + except Exception as exc: # noqa: BLE001 - re-raised unless the caller opts into the fallback + if not allow_tokenizer_fallback: + raise + logger.warning(f"AR tokenizer load failed ({exc}); falling back to string prompt (no BPE parity).") + ar_tokenizer = None + + if ar_tokenizer is not None and validate_tokenizer is not None: + validate_tokenizer(ar_tokenizer) + + ar_inputs = ar_input_builder( + prompt=prompt_text, + tokenizer=ar_tokenizer, + extra_body=extra_body, + num_images=num_images, + height=height, + width=width, + text_output=text_output, + ) + + if ar_inputs.prompt_token_ids is not None: + prompt_dict["prompt_token_ids"] = ar_inputs.prompt_token_ids + elif ar_inputs.prompt is not None: + prompt_dict["prompt"] = ar_inputs.prompt + if ar_inputs.use_system_prompt: + prompt_dict["use_system_prompt"] = ar_inputs.use_system_prompt + prompt_dict["modalities"] = ar_inputs.modalities + + # Apply stop_token_ids to exactly the stage(s) the builder names -- not + # "whichever stage isn't a diffusion stage," which would misfire on any + # future topology with more than one non-diffusion stage. + for stage_index in ar_inputs.stage_indices: + if stage_index >= len(sampling_params_list): + continue + stage_params = sampling_params_list[stage_index] + if not isinstance(stage_params, OmniDiffusionSamplingParams) and hasattr(stage_params, "stop_token_ids"): + stage_params.stop_token_ids = ar_inputs.stop_token_ids + + def parse_profiler_config(value: str) -> dict[str, Any]: try: config = json.loads(value) @@ -419,6 +507,22 @@ def parse_args() -> argparse.Namespace: default=None, help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', ) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Trust and execute custom modeling code from the model repo (required by e.g. HunyuanImage-3.0).", + ) + parser.add_argument( + "--allow-tokenizer-fallback", + action="store_true", + help=( + "For models with a declared ar_input_builder (e.g. HunyuanImage-3.0): if loading " + "the AR tokenizer fails, degrade to the string-prompt form (no BPE parity) instead " + "of failing the run. Off by default -- a tokenizer load failure usually means a " + "missing --trust-remote-code or a network/cache issue that's worth surfacing, not " + "silently masking with a possibly-wrong prompt." + ), + ) return parser.parse_args() @@ -492,6 +596,8 @@ def main(): ) if args.enforce_eager is not None: omni_kwargs["enforce_eager"] = args.enforce_eager + if args.trust_remote_code: + omni_kwargs["trust_remote_code"] = True if args.deploy_config: omni_kwargs["deploy_config"] = args.deploy_config omni = Omni(**omni_kwargs) @@ -584,6 +690,26 @@ def main(): if not diffusion_replaced and len(sampling_params_list) == 1: sampling_params_list = [diffusion_params] + # Models with an AR text stage (e.g. HunyuanImage3 image-editing) declare an + # ar_input_builder; build the AR prefill + stop tokens declaratively from the + # plain prompt + reference image count + extra_args. Others are untouched. + ar_input_builder = get_ar_input_builder(model_class_name) + if ar_input_builder is not None: + _apply_ar_stage_inputs( + ar_input_builder, + model=args.model, + prompt_text=args.prompt, + extra_body=extra_args_from_cli, + num_images=len(input_images), + height=args.height, + width=args.width, + prompt_dict=prompt_dict, + sampling_params_list=sampling_params_list, + trust_remote_code=args.trust_remote_code, + validate_tokenizer=get_ar_tokenizer_validator(model_class_name), + allow_tokenizer_fallback=args.allow_tokenizer_fallback, + ) + outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list) generation_end = time.perf_counter() generation_time = generation_end - generation_start diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 2ee09e1784e..adcc1066971 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -24,6 +24,8 @@ build_text_to_image_prompt as build_model_text_to_image_prompt, ) from vllm_omni.model_extras import ( + get_ar_input_builder, + get_ar_tokenizer_validator, get_extra_body_params, get_model_class_name, should_init_extra_args_for_non_diffusion_stages, @@ -377,10 +379,111 @@ def parse_args() -> argparse.Namespace: default=None, help="Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).", ) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Trust and execute custom modeling code from the model repo (required by e.g. HunyuanImage-3.0).", + ) + parser.add_argument( + "--allow-tokenizer-fallback", + action="store_true", + help=( + "For models with a declared ar_input_builder (e.g. HunyuanImage-3.0): if loading " + "the AR tokenizer fails, degrade to the string-prompt form (no BPE parity) instead " + "of failing the run. Off by default -- a tokenizer load failure usually means a " + "missing --trust-remote-code or a network/cache issue that's worth surfacing, not " + "silently masking with a possibly-wrong prompt." + ), + ) current_omni_platform.pre_register_and_update(parser) return parser.parse_args() +def _apply_ar_stage_inputs( + ar_input_builder: Any, + *, + model: str, + prompt_text: str, + extra_body: dict[str, Any], + num_images: int, + height: int | None, + width: int | None, + prompt_dict: dict[str, Any], + sampling_params_list: list[Any], + text_output: bool = False, + trust_remote_code: bool = False, + validate_tokenizer: Any | None = None, + allow_tokenizer_fallback: bool = False, +) -> None: + """Apply a model's declared AR-stage inputs to the request in place. + + Loads the model tokenizer (for byte-for-byte HF-parity segment + tokenization) and asks the model's ``ar_input_builder`` for the AR + prefill + stop tokens, then writes them onto ``prompt_dict`` and the + non-diffusion (AR) stage sampling params. Kept model-agnostic: the + example only knows the declared-hook contract. + + ``trust_remote_code`` must be threaded in from the caller's own resolved + ``--trust-remote-code`` flag -- this is a separate tokenizer load outside + the main engine, so it needs the same explicit user opt-in rather than + defaulting to trusting whatever ``--model`` was passed. + + ``validate_tokenizer``, when the model declares one via + ``get_ar_tokenizer_validator``, is called on a successfully-loaded real + tokenizer -- outside the load's try/except, so a validation failure (a + model/tokenizer revision drifting from hardcoded special-token ids) + raises instead of being swallowed into the string-prompt fallback. + + ``allow_tokenizer_fallback`` (default ``False``, i.e. fail fast): a + tokenizer load failure normally raises, since silently degrading to the + string-prompt form can produce a request that completes with subtly + wrong stop tokens instead of surfacing the real problem (missing + ``--trust-remote-code``, network/cache issue, etc.). Set ``True`` only + for explicit unit tests or a deliberate offline-compat run where the + caller has already decided a degraded prompt is acceptable. + """ + try: + from transformers import AutoTokenizer + + ar_tokenizer: Any | None = AutoTokenizer.from_pretrained(model, trust_remote_code=trust_remote_code) + except Exception as exc: # noqa: BLE001 - re-raised unless the caller opts into the fallback + if not allow_tokenizer_fallback: + raise + logger.warning(f"AR tokenizer load failed ({exc}); falling back to string prompt (no BPE parity).") + ar_tokenizer = None + + if ar_tokenizer is not None and validate_tokenizer is not None: + validate_tokenizer(ar_tokenizer) + + ar_inputs = ar_input_builder( + prompt=prompt_text, + tokenizer=ar_tokenizer, + extra_body=extra_body, + num_images=num_images, + height=height, + width=width, + text_output=text_output, + ) + + if ar_inputs.prompt_token_ids is not None: + prompt_dict["prompt_token_ids"] = ar_inputs.prompt_token_ids + elif ar_inputs.prompt is not None: + prompt_dict["prompt"] = ar_inputs.prompt + if ar_inputs.use_system_prompt: + prompt_dict["use_system_prompt"] = ar_inputs.use_system_prompt + prompt_dict["modalities"] = ar_inputs.modalities + + # Apply stop_token_ids to exactly the stage(s) the builder names -- not + # "whichever stage isn't a diffusion stage," which would misfire on any + # future topology with more than one non-diffusion stage. + for stage_index in ar_inputs.stage_indices: + if stage_index >= len(sampling_params_list): + continue + stage_params = sampling_params_list[stage_index] + if not isinstance(stage_params, OmniDiffusionSamplingParams) and hasattr(stage_params, "stop_token_ids"): + stage_params.stop_token_ids = ar_inputs.stop_token_ids + + def main(): args = parse_args() generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) @@ -469,6 +572,8 @@ def main(): omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size if args.enforce_eager is not None: omni_kwargs["enforce_eager"] = args.enforce_eager + if args.trust_remote_code: + omni_kwargs["trust_remote_code"] = True if args.deploy_config: omni_kwargs["deploy_config"] = args.deploy_config if use_nextstep: @@ -613,6 +718,27 @@ def main(): if not diffusion_replaced and len(sampling_params_list) == 1: sampling_params_list = [diffusion_params] + # Models with an AR text stage (e.g. HunyuanImage3) declare an + # ar_input_builder. When present, build the AR prefill token-ids and AR + # stop-token-ids declaratively from the plain prompt + extra_body, so this + # example stays model-agnostic. Models without one are untouched. + ar_input_builder = get_ar_input_builder(model_class_name) + if ar_input_builder is not None: + _apply_ar_stage_inputs( + ar_input_builder, + model=args.model, + prompt_text=args.prompt, + extra_body=user_extra, + num_images=0, + height=args.height, + width=args.width, + prompt_dict=prompt_dict, + sampling_params_list=sampling_params_list, + trust_remote_code=args.trust_remote_code, + validate_tokenizer=get_ar_tokenizer_validator(model_class_name), + allow_tokenizer_fallback=args.allow_tokenizer_fallback, + ) + outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list) generation_end = time.perf_counter() diff --git a/recipes/Tencent/HunyuanImage-3.0-Instruct.md b/recipes/Tencent/HunyuanImage-3.0-Instruct.md index 4614df382d5..eb5ccedf631 100644 --- a/recipes/Tencent/HunyuanImage-3.0-Instruct.md +++ b/recipes/Tencent/HunyuanImage-3.0-Instruct.md @@ -8,7 +8,7 @@ - Vendor: Tencent Hunyuan - Model: `tencent/HunyuanImage-3.0-Instruct` -- Task: Text-to-image generation; text-to-text and image-to-text understanding +- Task: Text-to-image generation, image editing; text-to-text and image-to-text understanding - Mode: Offline understanding, online serving, and performance benchmarking - Maintainer: Community @@ -31,8 +31,10 @@ FP8/NVFP4 configuration: ## References - Model: -- Existing model-specific offline example: - [`examples/offline_inference/hunyuan_image3`](../../examples/offline_inference/hunyuan_image3) +- Shared T2I example: + [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py) +- Shared IT2I example: + [`examples/offline_inference/image_to_image/image_edit.py`](../../examples/offline_inference/image_to_image/image_edit.py) - Shared T2T/I2T example: [`examples/offline_inference/x_to_text/x_to_text.py`](../../examples/offline_inference/x_to_text/x_to_text.py) - Related PRs: @@ -66,8 +68,48 @@ python examples/offline_inference/x_to_text/x_to_text.py \ ``` The AR-only default uses four GPUs. Pass `--deploy-config` to override the -layout. Image-output tasks remain documented by the existing model-specific -example and the DiT sections below. +layout. + +### Shared T2I/IT2I offline examples + +Text-to-image and image-editing route through the shared task examples, with +all HunyuanImage-3.0-specific knobs (`bot_task`, `use_system_prompt`, +`system_prompt`, `negative_prompt`) declared in +`vllm_omni/model_extras/hunyuan_image3.py` and passed via `--extra-body` / +`--extra-args`: + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model tencent/HunyuanImage-3.0-Instruct \ + --deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml \ + --trust-remote-code \ + --prompt "A cute cat sitting on a windowsill watching the sunset" \ + --height 1024 --width 1024 \ + --guidance-scale 5.0 --num-inference-steps 50 --seed 42 \ + --extra-body '{"bot_task": "think", "use_system_prompt": "en_recaption"}' \ + --output hunyuan_t2i.png +``` + +```bash +python examples/offline_inference/image_to_image/image_edit.py \ + --model tencent/HunyuanImage-3.0-Instruct \ + --deploy-config vllm_omni/deploy/hunyuan_image_3_moe.yaml \ + --trust-remote-code \ + --image /path/to/image.png \ + --prompt "Make the petals neon pink" \ + --guidance-scale 5.0 --num-inference-steps 50 --seed 42 \ + --extra-args '{"bot_task": "think"}' \ + --output hunyuan_edit.png +``` + +`image_edit.py` accepts up to 3 reference images (repeat `--image`) for +HunyuanImage-3.0's multi-image fusion. All four shared-script paths (t2i, +it2i, i2t, t2t) build the AR prefill and stop tokens through the same +declarative seam (`build_ar_stage_inputs` / `build_x_to_text_prompt`, both +wrapping `prompt_utils.build_ar_prompt_inputs`), so prompt formatting is +identical across them. The OpenAI server's `serving_chat.py` does not go +through this seam yet -- see that function's docstring for the resulting +`bot_task`-resolution divergence risk when omitted from `extra_body`. ### 4x H100/H800 80GB diff --git a/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py b/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py index 7cb7ea83ad4..7a467e35973 100644 --- a/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py +++ b/tests/diffusion/models/hunyuan_image3/test_prompt_utils.py @@ -13,9 +13,11 @@ HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, available_bot_tasks, available_tasks, + build_ar_prompt_inputs, build_prompt, build_prompt_tokens, resolve_stop_token_ids, + validate_special_token_ids, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -296,19 +298,42 @@ def _repo_root() -> pathlib.Path: return pathlib.Path(__file__).resolve().parents[4] -def test_end2end_routes_through_shared_prompt_utils(): - end2end_path = _repo_root() / "examples" / "offline_inference" / "hunyuan_image3" / "end2end.py" - tree = ast.parse(end2end_path.read_text(encoding="utf-8")) +def test_example_routes_through_shared_seam(): + """Neither the shared task examples nor the ``model_extras`` seam itself + should duplicate AR prompt tokenization. - local_func_names = {n.name for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)} - assert not (local_func_names & {"build_prompt", "build_prompt_tokens"}) + After the model_extras migration, t2i / image-editing run through the + shared ``text_to_image.py`` / ``image_edit.py`` examples (both call + ``build_ar_stage_inputs``), and t2t / i2t run through the shared + ``x_to_text.py`` example (which calls ``build_x_to_text_prompt``). This + test pins that both ``model_extras.hunyuan_image3`` seam functions + delegate to ``prompt_utils.build_ar_prompt_inputs`` / ``build_ar_stage_inputs`` + rather than re-implementing tokenization inline. + """ + seam_path = _repo_root() / "vllm_omni" / "model_extras" / "hunyuan_image3.py" + seam_tree = ast.parse(seam_path.read_text(encoding="utf-8")) + + local_func_names = {n.name for n in ast.walk(seam_tree) if isinstance(n, ast.FunctionDef)} + assert not (local_func_names & {"build_prompt", "build_prompt_tokens", "resolve_stop_token_ids"}) + + # build_x_to_text_prompt must route through build_ar_stage_inputs, not + # re-derive the AR prefill / stop-token logic itself. + build_x_to_text_prompt_node = next( + n for n in ast.walk(seam_tree) if isinstance(n, ast.FunctionDef) and n.name == "build_x_to_text_prompt" + ) + called_names = { + node.func.id + for node in ast.walk(build_x_to_text_prompt_node) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) + } + assert "build_ar_stage_inputs" in called_names - imported_from_prompt_utils: set[str] = set() - for node in ast.walk(tree): + # The model_extras seam itself must delegate to the shared prompt_utils. + routed_from_prompt_utils: set[str] = set() + for node in ast.walk(seam_tree): if isinstance(node, ast.ImportFrom) and node.module and node.module.endswith("hunyuan_image3.prompt_utils"): - imported_from_prompt_utils.update(alias.name for alias in node.names) - expected_imports = {"build_prompt_tokens", "resolve_stop_token_ids", "resolve_sys_type"} - assert expected_imports <= imported_from_prompt_utils + routed_from_prompt_utils.update(alias.name for alias in node.names) + assert "build_ar_prompt_inputs" in routed_from_prompt_utils _HUNYUAN_MODEL_ID = "tencent/HunyuanImage-3.0-Instruct" @@ -331,3 +356,87 @@ def test_segment_tokenize_diverges_from_full_string_encode(): full_ids = tok.encode(build_prompt(user_prompt, task="i2t", bot_task=None), add_special_tokens=False) assert seg_ids != full_ids assert len(seg_ids) >= len(full_ids) + + +def test_build_ar_prompt_inputs_tokenizer_path_matches_build_prompt_tokens(): + """With a tokenizer, the helper delegates to build_prompt_tokens (parity).""" + tok = FakeTokenizer() + ar = build_ar_prompt_inputs("a red panda", task="t2i", bot_task="think", tokenizer=tok) + expected = build_prompt_tokens("a red panda", FakeTokenizer(), task="t2i", bot_task="think") + assert ar.prompt is None + assert ar.prompt_token_ids == expected.token_ids + assert ar.system_prompt_type == expected.system_prompt_type + # think + no explicit size -> AR predicts ratio, so stop range is non-trivial. + assert len(ar.stop_token_ids) > 1 + + +def test_build_ar_prompt_inputs_string_fallback_without_tokenizer(): + """Without a tokenizer, the helper returns the build_prompt string form.""" + ar = build_ar_prompt_inputs("a red panda", task="t2i", bot_task="recaption") + assert ar.prompt_token_ids is None + assert ar.prompt == build_prompt("a red panda", task="t2i", bot_task="recaption") + assert ar.system_prompt_type + assert ar.stop_token_ids + + +def test_build_ar_prompt_inputs_explicit_size_stops_at_terminator(): + """Explicit image_size makes 'recaption' stop at only.""" + ar = build_ar_prompt_inputs("x", task="t2i", bot_task="recaption", image_size="1024x1024") + assert ar.stop_token_ids == [HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]] + + +def test_build_ar_prompt_inputs_comprehension_stops_on_answer(): + """Text-output tasks (t2t / i2t) stop on .""" + ar = build_ar_prompt_inputs("hi", task="t2t") + assert ar.stop_token_ids == [HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]] + + +class _MatchingTokenizer: + """Reports exactly HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS -- a "real" tokenizer + that hasn't drifted from the hardcoded table.""" + + def convert_tokens_to_ids(self, tok: str) -> int: + return HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[tok] + + +class _DriftedTokenizer: + """Reports a different id for one token -- simulates a model/tokenizer + revision bump shifting special-token ids out from under the hardcoded + table.""" + + def convert_tokens_to_ids(self, tok: str) -> int: + if tok == "": + return 999999 + return HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[tok] + + +class _BaseCheckpointTokenizer: + """Simulates tencent/HunyuanImage-3.0 (base): reports the same + model_class_name as -Instruct, but its vocab doesn't include the + Instruct-only - tokens -- convert_tokens_to_ids + falls back to unk_token_id for those, like a real HF tokenizer would.""" + + unk_token_id = 0 + _INSTRUCT_ONLY = {"", ""} + + def convert_tokens_to_ids(self, tok: str) -> int: + if tok in self._INSTRUCT_ONLY: + return self.unk_token_id + return HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[tok] + + +def test_validate_special_token_ids_passes_for_matching_tokenizer(): + validate_special_token_ids(_MatchingTokenizer()) # must not raise + + +def test_validate_special_token_ids_raises_on_drift(): + with pytest.raises(ValueError, match=".*expected 128023.*got 999999"): + validate_special_token_ids(_DriftedTokenizer()) + + +def test_validate_special_token_ids_tolerates_base_checkpoint_missing_tokens(): + """The base checkpoint's tokenizer not recognizing Instruct-only ratio + tokens (unk_token_id, not a real drifted id) must not be flagged as a + mismatch -- both checkpoints share model_class_name, so this hook fires + for either one.""" + validate_special_token_ids(_BaseCheckpointTokenizer()) # must not raise diff --git a/tests/e2e/offline_inference/test_hunyuanimage3.py b/tests/e2e/offline_inference/test_hunyuanimage3.py index 97179d96311..8655bd63fdd 100644 --- a/tests/e2e/offline_inference/test_hunyuanimage3.py +++ b/tests/e2e/offline_inference/test_hunyuanimage3.py @@ -1,5 +1,8 @@ # ruff: noqa: E501 +import importlib.util +import sys from collections.abc import Generator +from pathlib import Path import pytest import torch @@ -10,13 +13,15 @@ from tests.helpers.runtime import OmniRunner from tests.helpers.stage_config import get_deploy_config_path from vllm_omni import Omni +from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_extras import get_ar_input_builder, get_ar_tokenizer_validator, get_model_class_name from vllm_omni.platforms import current_omni_platform PROMPT = "A brown and white dog is running on the grass" MODEL_NAME = "tencent/HunyuanImage-3.0" LOCAL_CLIP_PATH = "openai/clip-vit-base-patch32" -DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image3.yaml") +DEPLOY_CONFIG_PATH = get_deploy_config_path("hunyuan_image_3_moe.yaml") pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion] @@ -328,6 +333,15 @@ def _generate_image(omni: Omni, use_system_prompt: str | None) -> Image.Image: return _extract_generated_image(outputs) +# A floor for "this is recognizably the same generation" against each +# case's precomputed reference embedding, not a quality bar -- catches a +# gross migration regression (wrong/blank/garbage image, or the wrong +# system prompt applied) without being flaky on run-to-run sampling noise. +# Matches the threshold test_shared_script_ar_path_reaches_generation uses +# for its own (seed=1234, en_recaption) comparison below. +MIN_SEMANTIC_SIMILARITY = 0.5 + + @pytest.mark.skipif(torch.accelerator.device_count() < 8, reason="Need at least 8 CUDA GPUs for this test.") @pytest.mark.parametrize("system_prompt_name,use_system_prompt,expected_embedding", SYSTEM_PROMPT_CASES) def test_system_prompt_scores( @@ -342,3 +356,109 @@ def test_system_prompt_scores( score = compare_semantic(expected_embedding, generated_image, clip_model, clip_processor) print(f"{system_prompt_name}: CLIP cosine similarity = {score:.6f}") + assert score >= MIN_SEMANTIC_SIMILARITY, ( + f"{system_prompt_name}: CLIP similarity {score:.4f} is below " + f"{MIN_SEMANTIC_SIMILARITY} against the reference (seed=1234) embedding -- " + "the migrated path may have diverged from the reference HunyuanImage-3.0 " + f"output for use_system_prompt={use_system_prompt!r}." + ) + + +def _load_text_to_image_module(): + """Dynamically load the shared example script by file path (``examples/`` + is not an installed package) so this test can call its real + ``_apply_ar_stage_inputs`` instead of re-implementing its logic.""" + repo_root = Path(__file__).resolve().parents[3] + path = repo_root / "examples/offline_inference/text_to_image/text_to_image.py" + spec = importlib.util.spec_from_file_location("_e2e_shared_text_to_image", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +@pytest.mark.skipif(torch.accelerator.device_count() < 8, reason="Need at least 8 CUDA GPUs for this test.") +def test_shared_script_ar_path_reaches_generation( + omni: Omni, + clip_bundle: tuple[CLIPModel, CLIPProcessor], +) -> None: + """Exercise the path this PR actually migrates HunyuanImage-3.0 onto: + + get_model_class_name -> get_ar_input_builder -> _apply_ar_stage_inputs + -> Omni.generate() + + using the *real* model_extras registry lookup (keyed on the engine's + live ``model_class_name``) and the shared ``text_to_image.py``'s own + helper -- not the hand-constructed ``sampling_params.extra_args`` that + ``_generate_image``/``test_system_prompt_scores`` above use, which never + touch the registry or the AR-tokenizer-loading path at all. This is the + one gap those tests leave: a regression in the registry key (the bug + this PR fixes), in ``_apply_ar_stage_inputs``, or in the + ``ar_tokenizer_validator`` wiring would not fail any of them. + + Reuses the ``en_recaption`` reference embedding from + ``test_system_prompt_scores`` so this also asserts the migrated path + produces the *same* output as the hand-constructed one, not just "didn't + crash." + """ + text_to_image = _load_text_to_image_module() + clip_model, clip_processor = clip_bundle + + model_class_name = get_model_class_name(omni) + ar_input_builder = get_ar_input_builder(model_class_name) + assert ar_input_builder is not None, ( + f"get_ar_input_builder({model_class_name!r}) returned None for the running " + "engine -- this is exactly the registry-key regression this PR fixes " + "(spec registered under the wrong model_class_name), and it would " + "silently skip AR-stage input building entirely." + ) + + generator_device = current_omni_platform.device_type or "cuda" + diffusion_params = OmniDiffusionSamplingParams( + seed=1234, + generator=torch.Generator(device=generator_device).manual_seed(1234), + num_outputs_per_prompt=1, + ) + defaults = list(omni.default_sampling_params_list or []) + sampling_params_list = [clone_sampling_params(p) for p in defaults] + for idx, params in enumerate(sampling_params_list): + if isinstance(params, OmniDiffusionSamplingParams): + sampling_params_list[idx] = diffusion_params + + prompt_dict: dict[str, object] = {"prompt": PROMPT, "modalities": ["image"]} + text_to_image._apply_ar_stage_inputs( + ar_input_builder, + model=MODEL_NAME, + prompt_text=PROMPT, + extra_body={"use_system_prompt": "en_recaption"}, + num_images=0, + height=None, + width=None, + prompt_dict=prompt_dict, + sampling_params_list=sampling_params_list, + trust_remote_code=True, + # validate_special_token_ids now tolerates tokens the loaded tokenizer + # doesn't recognize at all (unk_token_id) instead of requiring every + # entry in HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS to be present, so this is + # safe against MODEL_NAME's base checkpoint missing the Instruct-only + # / tokens. Wiring it here exercises the + # real production path (get_ar_tokenizer_validator via the registry), + # not just the isolated fake-tokenizer coverage in test_prompt_utils.py + # / test_shared_script_ar_integration.py. + validate_tokenizer=get_ar_tokenizer_validator(model_class_name), + ) + + outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list) + generated_image = _extract_generated_image(outputs) + score = compare_semantic(SYSTEM_EN_RECAPTION, generated_image, clip_model, clip_processor) + print(f"shared-script AR path (en_recaption): CLIP cosine similarity = {score:.6f}") + # test_system_prompt_scores now asserts this same (seed=1234, en_recaption) + # case through the hand-constructed extra_args path against the same + # module-level MIN_SEMANTIC_SIMILARITY floor, so this reuses it too rather + # than picking an independent threshold for the same comparison. + assert score >= MIN_SEMANTIC_SIMILARITY, ( + f"shared-script AR path (en_recaption) CLIP similarity {score:.4f} is below " + f"{MIN_SEMANTIC_SIMILARITY} -- the migrated get_ar_input_builder path may have " + "diverged from the reference HunyuanImage-3.0 output for this prompt/system-prompt." + ) diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index c45563617c2..60850248ef0 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -2314,3 +2314,72 @@ def test_image_edits_size_auto_preserves_bridge_size(async_omni_stage_configs_on assert captured_prompt["prompt"].count("") == 2, ( f"N=2 reference images must emit 2 placeholders in AR prompt; got {captured_prompt[KEY].count(IMG)} -- prompt: {captured_prompt[KEY]!r}" ) + + +@pytest.mark.xfail( + strict=True, + reason=( + "Pre-existing gap, not introduced by this change: the /v1/images/edits route in " + "api_server.py never calls resolve_stop_token_ids at all (grep: api_server.py has " + "no stop_token_ids reference), so the AR stage keeps SamplingParams' default empty " + "list. The omitted-bot_task fix landed in serving_chat.py, which is the only " + "production caller; that path is covered by " + "test_serving_chat_multistage_generation.py::" + "test_build_multistage_generation_inputs_omitted_bot_task_matches_prompt_default. " + "Wiring the images/edits route through the same seam needs a tokenizer in that " + "scope and a decision on whether HunyuanImage3 it2i should resolve AR stop tokens " + "unconditionally (today the whole AR block is gated on an explicit bot_task / " + "use_system_prompt / system_prompt), so it is tracked separately." + ), +) +def test_image_edits_omitted_bot_task_stop_tokens_match_prompt_default( + async_omni_stage_configs_only_client, +): + """Regression: an omitted bot_task must resolve identically for the AR + prompt and its stop_token_ids. + + build_prompt/build_prompt_tokens default an omitted bot_task per-task + (e.g. "think" for the base "it2i" task, since it isn't itself a key in + _TASK_PRESETS). resolve_stop_token_ids must land on that same default + to compute the matching stop set. Passing the raw (still-None) outer + bot_task variable to resolve_stop_token_ids -- instead of mirroring + build_kwargs's own omitted-or-not "bot_task" entry -- made it normalize + bot_task=None instead of "think", so it fell through to the full + stop range instead of the think/recaption-only pair. + + An explicit (non-"auto") size is required to reach this: with + need_ratio=True (size="auto"), resolve_stop_token_ids returns the full + ratio range regardless of bot_task, so the two code paths only visibly + disagree once a concrete size selects the narrower think/recaption stop + set. + """ + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, + ) + + img = make_test_image_bytes((64, 64)) + response = async_omni_stage_configs_only_client.post( + "/v1/images/edits", + files=[("image", img)], + data={ + "prompt": "make it neon", + "size": "512x512", + }, + ) + assert response.status_code == 200, response.text + + engine = async_omni_stage_configs_only_client.app.state.engine_client + captured = engine.captured_sampling_params_list + assert captured is not None + + ar_params = captured[0] + expected_stop_token_ids = [ + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""], + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""], + ] + assert ar_params.stop_token_ids == expected_stop_token_ids, ( + f"omitted bot_task with an explicit size must resolve stop_token_ids " + f"for the default 'think' bot_task ({expected_stop_token_ids}); got " + f"{ar_params.stop_token_ids} -- this is the full ratio range, meaning " + "resolve_stop_token_ids disagreed with build_prompt_tokens's default." + ) diff --git a/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py b/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py index 0dea9717957..639ddf6c650 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py +++ b/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py @@ -538,3 +538,84 @@ def test_build_multistage_generation_inputs_no_stop_token_ids_without_size(servi assert sampling_params_list[0].stop_token_ids == [], ( "Without bot_task, AR stage stop_token_ids must be the default empty list" ) + + +def test_build_multistage_generation_inputs_omitted_bot_task_matches_prompt_default(serving_chat): + """Regression: an omitted bot_task must resolve identically for the AR + prompt and for its stop_token_ids. + + ``build_prompt_tokens``/``build_prompt`` default an omitted bot_task + per-task via the ``_DEFAULT_BOT_TASK`` sentinel (``"think"`` for the base + ``it2i`` task, which isn't itself a key in ``_TASK_PRESETS``). + ``resolve_stop_token_ids`` must normalize from that same omitted-or-not + starting point. Passing the raw -- still ``None`` -- outer ``bot_task`` + variable instead of mirroring ``build_kwargs``'s own omitted-or-not + ``"bot_task"`` entry made it normalize ``bot_task=None`` (plain mode), so + it fell through to the full ```` range instead of the + think/recaption terminator pair. + + An explicit (non-"auto") size is required to observe this: with + ``need_ratio=True`` the ratio range is returned regardless of bot_task, so + the two paths only visibly disagree once a concrete size selects the + narrower stop set. + """ + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, + ) + from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + + class FakeTokenizer: + SPECIAL = { + "<|startoftext|>": 1, + "": 2, + "": 4, + } + + def convert_tokens_to_ids(self, tok: str) -> int: + return self.SPECIAL.get(tok, 0) + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + return list(range(100, 100 + len(text))) + + def _build(extra_body): + engine = SimpleNamespace( + stage_configs=[ + SimpleNamespace(stage_type="llm", is_comprehension=True), + SimpleNamespace(stage_type="diffusion", is_comprehension=False), + ], + default_sampling_params_list=[ + SamplingParams(temperature=0.0), + OmniDiffusionSamplingParams(), + ], + ) + _, sampling_params_list = OmniOpenAIServingChat._build_multistage_generation_inputs( + serving_chat, + engine=engine, + prompt="draw a cat", + extra_body=extra_body, + reference_images=[Image.new("RGB", (32, 32), color="red")], + gen_params=OmniDiffusionSamplingParams(height=768, width=1024), + tokenizer=FakeTokenizer(), + ) + return sampling_params_list[0].stop_token_ids + + # bot_task omitted entirely; use_system_prompt is what carries the request + # into the AR-prompt path. + omitted = _build({"use_system_prompt": "en_recaption"}) + # Same request with the default spelled out explicitly. + explicit = _build({"bot_task": "think", "use_system_prompt": "en_recaption"}) + + expected = [ + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""], + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""], + ] + assert omitted == expected, ( + f"omitted bot_task with an explicit size must resolve stop_token_ids for the " + f"default 'think' bot_task ({expected}); got {omitted}. A long list here is the " + f"full range, meaning resolve_stop_token_ids normalized bot_task=None " + f"(plain mode) instead of agreeing with build_prompt_tokens's default." + ) + assert omitted == explicit, ( + "omitting bot_task must resolve the same stop_token_ids as passing the default " + f"explicitly; omitted={omitted} explicit={explicit}" + ) diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index e55691e104d..dcdf13678af 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -3,6 +3,7 @@ from __future__ import annotations +import sys from typing import Any import pytest @@ -592,6 +593,193 @@ def test_declared_extra_args_apply_to_existing_sampling_params() -> None: } +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_extra_registry_declares_request_and_response_params() -> None: + assert get_extra_body_params("HunyuanImage3Pipeline") == frozenset( + { + "bot_task", + "use_system_prompt", + "system_prompt", + "negative_prompt", + } + ) + assert get_extra_output_params("HunyuanImage3Pipeline") == frozenset() + assert should_init_extra_args_for_non_diffusion_stages("HunyuanImage3Pipeline") is True + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_extra_registry_matches_runtime_model_class_name() -> None: + # od_config.model_class_name resolves to "HunyuanImage3ForCausalMM" at + # runtime; "HunyuanImage3Pipeline" is kept only as a compatibility alias. + assert get_extra_body_params("HunyuanImage3ForCausalMM") == get_extra_body_params("HunyuanImage3Pipeline") + assert get_extra_output_params("HunyuanImage3ForCausalMM") == get_extra_output_params("HunyuanImage3Pipeline") + assert should_init_extra_args_for_non_diffusion_stages( + "HunyuanImage3ForCausalMM" + ) == should_init_extra_args_for_non_diffusion_stages("HunyuanImage3Pipeline") + + +class _FakeHunyuanTokenizer: + """Minimal tokenizer with distinct ids so AR token assertions are unambiguous.""" + + SPECIAL = { + "<|startoftext|>": 1, + "": 2, + "": 3, + "": 4, + } + + def convert_tokens_to_ids(self, tok: str) -> int: + return self.SPECIAL.get(tok, 0) + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + # Offset well above the special-id range to avoid collisions. + return [500 + len(text)] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_ar_input_builder_registered() -> None: + from vllm_omni.model_extras import get_ar_input_builder + + builder = get_ar_input_builder("HunyuanImage3Pipeline") + assert builder is not None + assert get_ar_input_builder("HunyuanImage3ForCausalMM") is builder + # Models without the hook return None. + assert get_ar_input_builder("WanVACEPipeline") is None + assert get_ar_input_builder(None) is None + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_ar_tokenizer_validator_registered() -> None: + from vllm_omni.model_extras import get_ar_tokenizer_validator + + validator = get_ar_tokenizer_validator("HunyuanImage3Pipeline") + assert validator is not None + assert get_ar_tokenizer_validator("HunyuanImage3ForCausalMM") is validator + # Models without the hook return None. + assert get_ar_tokenizer_validator("WanVACEPipeline") is None + assert get_ar_tokenizer_validator(None) is None + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_build_ar_stage_inputs_t2i_string_fallback() -> None: + from vllm_omni.model_extras.hunyuan_image3 import build_ar_stage_inputs + + ar = build_ar_stage_inputs( + prompt="a red panda", + tokenizer=None, + extra_body={"bot_task": "recaption"}, + num_images=0, + height=1024, + width=1024, + ) + # No tokenizer -> string prompt form; image modality; recaption + explicit + # size stops at only. + assert ar.prompt_token_ids is None + assert ar.prompt and "a red panda" in ar.prompt + assert ar.modalities == ["image"] + assert ar.use_system_prompt + assert len(ar.stop_token_ids) == 1 + # stop_token_ids belong to stage 0 (HunyuanImage3's single AR stage) -- + # declared explicitly, not inferred by the caller from stage type. + assert ar.stage_indices == [0] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_build_ar_stage_inputs_t2i_tokenizer_parity() -> None: + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens + from vllm_omni.model_extras.hunyuan_image3 import build_ar_stage_inputs + + ar = build_ar_stage_inputs( + prompt="a red panda", + tokenizer=_FakeHunyuanTokenizer(), + extra_body={"bot_task": "think"}, + num_images=0, + ) + expected = build_prompt_tokens("a red panda", _FakeHunyuanTokenizer(), task="t2i", bot_task="think") + assert ar.prompt is None + assert ar.prompt_token_ids == expected.token_ids + assert ar.modalities == ["image"] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_build_ar_stage_inputs_it2i_inserts_image_placeholders() -> None: + from vllm_omni.model_extras.hunyuan_image3 import build_ar_stage_inputs + + ar = build_ar_stage_inputs( + prompt="edit it", + tokenizer=_FakeHunyuanTokenizer(), + extra_body={}, + num_images=2, + ) + # Two reference images -> two placeholder ids (2) in the AR prefill. + assert ar.prompt_token_ids is not None + assert ar.prompt_token_ids.count(_FakeHunyuanTokenizer.SPECIAL[""]) == 2 + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_build_ar_stage_inputs_comprehension_text_output() -> None: + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, + ) + from vllm_omni.model_extras.hunyuan_image3 import build_ar_stage_inputs + + # t2t (no image, text output) -> text modality, stop on . + ar_t2t = build_ar_stage_inputs("hi", tokenizer=None, extra_body={}, num_images=0, text_output=True) + assert ar_t2t.modalities == ["text"] + assert ar_t2t.stop_token_ids == [HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]] + + # i2t (image, text output) -> single placeholder in the prefill. + ar_i2t = build_ar_stage_inputs( + "describe", tokenizer=_FakeHunyuanTokenizer(), extra_body={}, num_images=1, text_output=True + ) + assert ar_i2t.modalities == ["text"] + assert ar_i2t.prompt_token_ids.count(_FakeHunyuanTokenizer.SPECIAL[""]) == 1 + + +class _RealIdHunyuanTokenizer: + """Mirrors HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS exactly, so validate_ar_tokenizer + accepts it -- unlike _FakeHunyuanTokenizer above, whose ids are deliberately + wrong to exercise the drift-detection failure path.""" + + unk_token_id = -1 + + def convert_tokens_to_ids(self, tok: str) -> int: + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, + ) + + return HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS.get(tok, self.unk_token_id) + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + return [900 + len(text)] + + +def _patch_hunyuan_autotokenizer(monkeypatch: pytest.MonkeyPatch, tokenizer: object) -> list[str]: + """Patch the ``AutoTokenizer`` that ``hunyuan_image3.build_x_to_text_prompt`` + imports locally, recording the models it was asked to load.""" + calls: list[str] = [] + + class _FakeAutoTokenizer: + @staticmethod + def from_pretrained(model: str, trust_remote_code: bool = False) -> object: + calls.append(model) + assert trust_remote_code is True + return tokenizer + + fake_transformers = type(sys)("transformers") + fake_transformers.AutoTokenizer = _FakeAutoTokenizer + monkeypatch.setitem(sys.modules, "transformers", fake_transformers) + return calls + + @pytest.mark.core_model @pytest.mark.cpu def test_mammothmoda2_extra_registry_declares_request_and_response_params() -> None: @@ -619,6 +807,59 @@ def test_mammothmoda2_extra_registry_declares_request_and_response_params() -> N assert wrapper_prompt["additional_information"]["omni_task"] == ["t2i"] +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_x_to_text_prompt_builder_t2t(monkeypatch: pytest.MonkeyPatch) -> None: + """t2t (no image) routes through build_ar_stage_inputs and stops on .""" + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, + ) + from vllm_omni.model_extras.hunyuan_image3 import build_x_to_text_prompt as hunyuan_build_x_to_text_prompt + + calls = _patch_hunyuan_autotokenizer(monkeypatch, _RealIdHunyuanTokenizer()) + prompt_dict, stop_token_ids = hunyuan_build_x_to_text_prompt( + "tencent/HunyuanImage-3.0-Instruct", "What is the capital of France?", has_image=False + ) + + assert calls == ["tencent/HunyuanImage-3.0-Instruct"] + assert prompt_dict["modalities"] == ["text"] + assert prompt_dict["prompt_token_ids"] is not None + assert prompt_dict["prompt_token_ids"].count(HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]) == 0 + assert stop_token_ids == [HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_x_to_text_prompt_builder_i2t(monkeypatch: pytest.MonkeyPatch) -> None: + """i2t (image) inserts exactly one placeholder and still stops on .""" + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, + ) + from vllm_omni.model_extras.hunyuan_image3 import build_x_to_text_prompt as hunyuan_build_x_to_text_prompt + + _patch_hunyuan_autotokenizer(monkeypatch, _RealIdHunyuanTokenizer()) + prompt_dict, stop_token_ids = hunyuan_build_x_to_text_prompt( + "tencent/HunyuanImage-3.0-Instruct", "Describe this image.", has_image=True + ) + + assert prompt_dict["modalities"] == ["text"] + assert prompt_dict["prompt_token_ids"].count(HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]) == 1 + assert stop_token_ids == [HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS[""]] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_hunyuan_image3_x_to_text_prompt_builder_validates_tokenizer(monkeypatch: pytest.MonkeyPatch) -> None: + """A tokenizer whose special-token ids have drifted must be rejected here too -- + x_to_text.py (unlike text_to_image.py/image_edit.py) has no registry-driven + ar_tokenizer_validator wiring of its own, so this seam must self-validate.""" + from vllm_omni.model_extras.hunyuan_image3 import build_x_to_text_prompt as hunyuan_build_x_to_text_prompt + + _patch_hunyuan_autotokenizer(monkeypatch, _FakeHunyuanTokenizer()) + with pytest.raises(ValueError, match="no longer match"): + hunyuan_build_x_to_text_prompt("tencent/HunyuanImage-3.0-Instruct", "Hello.", has_image=False) + + @pytest.mark.core_model @pytest.mark.cpu def test_mammothmoda2_text_to_image_prompt_builder() -> None: diff --git a/tests/model_extras/test_shared_script_ar_integration.py b/tests/model_extras/test_shared_script_ar_integration.py new file mode 100644 index 00000000000..f781ca5fbc4 --- /dev/null +++ b/tests/model_extras/test_shared_script_ar_integration.py @@ -0,0 +1,340 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Integration coverage for the shared task examples' ``_apply_ar_stage_inputs``. + +The unit tests in ``test_model_extras.py`` cover ``get_ar_input_builder`` and +``build_ar_stage_inputs`` in isolation, but nothing previously exercised the +*glue* in ``text_to_image.py`` / ``image_edit.py`` that wires them together: +resolving the model's declared ``ar_input_builder`` via the registry, loading +the AR tokenizer, and writing the result onto ``prompt_dict`` / +``sampling_params_list``. A regression in that glue -- e.g. the registry key +mismatch fixed in this PR (``HunyuanImage3ForCausalMM`` vs +``HunyuanImage3Pipeline``), or the ``trust_remote_code`` opt-in being +silently bypassed -- would not have been caught by the isolated unit tests +above, since they call ``build_ar_stage_inputs`` directly rather than going +through the shared script's own code path. + +These tests dynamically load the example scripts by file path (``examples/`` +is not an installed package) and call ``_apply_ar_stage_inputs`` the same way +``main()`` does, with a fake tokenizer standing in for a real HF download. +""" + +from __future__ import annotations + +import importlib.util +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import pytest + +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_REPO_ROOT = Path(__file__).resolve().parents[2] + + +def _load_module(name: str, relpath: str): + path = _REPO_ROOT / relpath + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +class _FakeHunyuanTokenizer: + """Mirrors test_model_extras.py's fixture: distinct ids, no network I/O.""" + + SPECIAL = {"<|startoftext|>": 1, "": 2, "": 3, "": 4} + + def convert_tokens_to_ids(self, tok: str) -> int: + return self.SPECIAL.get(tok, 0) + + def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: + return [500 + len(text)] + + +@dataclass +class _FakeARSamplingParams: + """Stands in for an LLM (non-diffusion) stage's sampling params.""" + + stop_token_ids: list[int] | None = None + + +def _patch_tokenizer_loader(monkeypatch: pytest.MonkeyPatch, module: Any) -> list[dict[str, Any]]: + """Patch the ``AutoTokenizer`` the module imports locally, recording calls.""" + calls: list[dict[str, Any]] = [] + + class _FakeAutoTokenizer: + @staticmethod + def from_pretrained(model: str, trust_remote_code: bool = False) -> _FakeHunyuanTokenizer: + calls.append({"model": model, "trust_remote_code": trust_remote_code}) + return _FakeHunyuanTokenizer() + + fake_transformers = type(sys)("transformers") + fake_transformers.AutoTokenizer = _FakeAutoTokenizer + monkeypatch.setitem(sys.modules, "transformers", fake_transformers) + return calls + + +def _patch_failing_tokenizer_loader(monkeypatch: pytest.MonkeyPatch) -> None: + """Patch AutoTokenizer.from_pretrained to always raise, simulating a + missing --trust-remote-code / network / cache failure.""" + + class _FailingAutoTokenizer: + @staticmethod + def from_pretrained(model: str, trust_remote_code: bool = False) -> Any: + raise OSError("simulated tokenizer load failure") + + fake_transformers = type(sys)("transformers") + fake_transformers.AutoTokenizer = _FailingAutoTokenizer + monkeypatch.setitem(sys.modules, "transformers", fake_transformers) + + +@pytest.mark.parametrize( + "module_name,relpath", + [ + ("_shared_text_to_image", "examples/offline_inference/text_to_image/text_to_image.py"), + ("_shared_image_edit", "examples/offline_inference/image_to_image/image_edit.py"), + ], +) +def test_apply_ar_stage_inputs_routes_through_real_registry( + monkeypatch: pytest.MonkeyPatch, module_name: str, relpath: str +) -> None: + """The shared script's glue must resolve HunyuanImage3's *real* registry + entry (keyed on the runtime ``model_class_name``) and write its output + onto the request -- not a stub, not a hardcoded id.""" + from vllm_omni.model_extras import get_ar_input_builder + + module = _load_module(module_name, relpath) + tokenizer_calls = _patch_tokenizer_loader(monkeypatch, module) + + ar_input_builder = get_ar_input_builder("HunyuanImage3ForCausalMM") + assert ar_input_builder is not None, ( + "get_ar_input_builder('HunyuanImage3ForCausalMM') returned None -- " + "if this regresses, the shared script silently skips AR-stage " + "input building entirely (this is exactly how the registry-key bug " + "this PR fixes went undetected)." + ) + + prompt_dict: dict[str, Any] = {} + ar_params = _FakeARSamplingParams() + diffusion_params = OmniDiffusionSamplingParams() + sampling_params_list: list[Any] = [ar_params, diffusion_params] + + module._apply_ar_stage_inputs( + ar_input_builder, + model="tencent/HunyuanImage-3.0-Instruct", + prompt_text="a red panda", + extra_body={"bot_task": "think", "use_system_prompt": "en_recaption"}, + num_images=0, + height=1024, + width=1024, + prompt_dict=prompt_dict, + sampling_params_list=sampling_params_list, + trust_remote_code=False, + ) + + # The tokenizer was loaded exactly once, for the model the user asked for. + assert len(tokenizer_calls) == 1 + assert tokenizer_calls[0]["model"] == "tencent/HunyuanImage-3.0-Instruct" + + # Byte-for-byte tokenizer path: prompt_token_ids populated, not the string form. + assert prompt_dict.get("prompt_token_ids"), "AR prefill token-ids were not written onto the request" + assert "prompt" not in prompt_dict + assert prompt_dict.get("use_system_prompt") == "en_recaption" + assert prompt_dict.get("modalities") == ["image"] + + # stop_token_ids reached the AR-stage params, and only those. The diffusion + # stage must be left untouched: OmniDiffusionSamplingParams has no + # stop_token_ids field, so the helper skipping it means the attribute is + # never created (reading it would raise AttributeError, not return None). + assert ar_params.stop_token_ids, "AR stage did not receive stop_token_ids" + assert not hasattr(diffusion_params, "stop_token_ids"), "diffusion stage params must not receive stop_token_ids" + + +@pytest.mark.parametrize( + "module_name,relpath,cli_flag", + [ + ("_shared_text_to_image_trc", "examples/offline_inference/text_to_image/text_to_image.py", False), + ("_shared_text_to_image_trc", "examples/offline_inference/text_to_image/text_to_image.py", True), + ], +) +def test_apply_ar_stage_inputs_respects_trust_remote_code_option( + monkeypatch: pytest.MonkeyPatch, module_name: str, relpath: str, cli_flag: bool +) -> None: + """The tokenizer load must use the caller's resolved --trust-remote-code + value, not silently default to True regardless of user intent.""" + from vllm_omni.model_extras import get_ar_input_builder + + module = _load_module(module_name, relpath) + tokenizer_calls = _patch_tokenizer_loader(monkeypatch, module) + ar_input_builder = get_ar_input_builder("HunyuanImage3ForCausalMM") + + module._apply_ar_stage_inputs( + ar_input_builder, + model="tencent/HunyuanImage-3.0-Instruct", + prompt_text="a red panda", + extra_body={}, + num_images=0, + height=1024, + width=1024, + prompt_dict={}, + sampling_params_list=[_FakeARSamplingParams(), OmniDiffusionSamplingParams()], + trust_remote_code=cli_flag, + ) + + assert tokenizer_calls[0]["trust_remote_code"] is cli_flag + + +@pytest.mark.parametrize( + "module_name,relpath", + [ + ("_shared_text_to_image_validate", "examples/offline_inference/text_to_image/text_to_image.py"), + ("_shared_image_edit_validate", "examples/offline_inference/image_to_image/image_edit.py"), + ], +) +def test_apply_ar_stage_inputs_runs_declared_tokenizer_validator( + monkeypatch: pytest.MonkeyPatch, module_name: str, relpath: str +) -> None: + """get_ar_tokenizer_validator's hook must run on a real tokenizer load, + and a validation failure must propagate -- not get swallowed by the + tokenizer-load try/except (which would silently fall back to the string + prompt instead of surfacing the drift).""" + from vllm_omni.model_extras import get_ar_input_builder, get_ar_tokenizer_validator + + module = _load_module(module_name, relpath) + _patch_tokenizer_loader(monkeypatch, module) + ar_input_builder = get_ar_input_builder("HunyuanImage3ForCausalMM") + validator = get_ar_tokenizer_validator("HunyuanImage3ForCausalMM") + assert validator is not None + + # The fake tokenizer's ids intentionally don't match + # HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS, so the real validator must reject it. + with pytest.raises(ValueError, match="no longer match"): + module._apply_ar_stage_inputs( + ar_input_builder, + model="tencent/HunyuanImage-3.0-Instruct", + prompt_text="a red panda", + extra_body={}, + num_images=0, + height=1024, + width=1024, + prompt_dict={}, + sampling_params_list=[_FakeARSamplingParams(), OmniDiffusionSamplingParams()], + trust_remote_code=False, + validate_tokenizer=validator, + ) + + +@pytest.mark.parametrize( + "module_name,relpath,required_argv", + [ + ("_shared_text_to_image_cli", "examples/offline_inference/text_to_image/text_to_image.py", []), + # image_edit.py's parser declares --image/--prompt as required (no defaults); + # omitting them makes parse_args() raise SystemExit(2) before the assertion + # below is ever reached. + ( + "_shared_image_edit_cli", + "examples/offline_inference/image_to_image/image_edit.py", + ["--image", "dummy.png", "--prompt", "x"], + ), + ], +) +@pytest.mark.parametrize("argv,expected", [([], False), (["--trust-remote-code"], True)]) +def test_trust_remote_code_flag_parses_from_real_cli( + monkeypatch: pytest.MonkeyPatch, + module_name: str, + relpath: str, + required_argv: list[str], + argv: list[str], + expected: bool, +) -> None: + """The other tests in this file call _apply_ar_stage_inputs directly with + an already-resolved trust_remote_code bool, which never exercises + argparse itself. Closing that specific gap doesn't need a GPU or a model + load -- parse_args() only touches sys.argv -- so cover it here instead of + the (expensive, one-model-load-per-invocation) full CLI subprocess a + stricter GPU e2e test would need.""" + module = _load_module(module_name, relpath) + monkeypatch.setattr(sys, "argv", ["prog", *required_argv, *argv]) + args = module.parse_args() + assert args.trust_remote_code is expected + + +@pytest.mark.parametrize( + "module_name,relpath", + [ + ("_shared_text_to_image_failfast", "examples/offline_inference/text_to_image/text_to_image.py"), + ("_shared_image_edit_failfast", "examples/offline_inference/image_to_image/image_edit.py"), + ], +) +def test_apply_ar_stage_inputs_tokenizer_load_failure_raises_by_default( + monkeypatch: pytest.MonkeyPatch, module_name: str, relpath: str +) -> None: + """A tokenizer load failure must surface as an error by default, not + silently degrade to the string-prompt form -- that degraded form can + complete "successfully" with subtly wrong stop tokens, masking the real + problem (e.g. a missing --trust-remote-code).""" + from vllm_omni.model_extras import get_ar_input_builder + + module = _load_module(module_name, relpath) + _patch_failing_tokenizer_loader(monkeypatch) + ar_input_builder = get_ar_input_builder("HunyuanImage3ForCausalMM") + + with pytest.raises(OSError, match="simulated tokenizer load failure"): + module._apply_ar_stage_inputs( + ar_input_builder, + model="tencent/HunyuanImage-3.0-Instruct", + prompt_text="a red panda", + extra_body={}, + num_images=0, + height=1024, + width=1024, + prompt_dict={}, + sampling_params_list=[_FakeARSamplingParams(), OmniDiffusionSamplingParams()], + trust_remote_code=False, + ) + + +@pytest.mark.parametrize( + "module_name,relpath", + [ + ("_shared_text_to_image_fallback", "examples/offline_inference/text_to_image/text_to_image.py"), + ("_shared_image_edit_fallback", "examples/offline_inference/image_to_image/image_edit.py"), + ], +) +def test_apply_ar_stage_inputs_allow_tokenizer_fallback_option( + monkeypatch: pytest.MonkeyPatch, module_name: str, relpath: str +) -> None: + """With allow_tokenizer_fallback=True (the CLI's --allow-tokenizer-fallback), + a tokenizer load failure degrades to the string-prompt form instead of + raising -- for a deliberate offline-compat run, not the default.""" + from vllm_omni.model_extras import get_ar_input_builder + + module = _load_module(module_name, relpath) + _patch_failing_tokenizer_loader(monkeypatch) + ar_input_builder = get_ar_input_builder("HunyuanImage3ForCausalMM") + + prompt_dict: dict[str, Any] = {} + module._apply_ar_stage_inputs( + ar_input_builder, + model="tencent/HunyuanImage-3.0-Instruct", + prompt_text="a red panda", + extra_body={"bot_task": "recaption"}, + num_images=0, + height=1024, + width=1024, + prompt_dict=prompt_dict, + sampling_params_list=[_FakeARSamplingParams(), OmniDiffusionSamplingParams()], + trust_remote_code=False, + allow_tokenizer_fallback=True, + ) + + # No tokenizer -> string-prompt form, not the byte-for-byte token-ids form. + assert "prompt_token_ids" not in prompt_dict + assert prompt_dict.get("prompt") and "a red panda" in prompt_dict["prompt"] diff --git a/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py b/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py index b64fe2e97cd..b4ecd7dcac5 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py @@ -60,6 +60,48 @@ "": 130106, } + +def validate_special_token_ids(tokenizer: Any) -> None: + """Fail fast if a real tokenizer's ids drift from the hardcoded table. + + ``HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS`` exists so stop-token / trigger-tag + resolution doesn't need a tokenizer round-trip, but that only holds as + long as the ids match ``tokenizer.json`` for whatever model repo/revision + is actually loaded. Call this once a real tokenizer is available (the + example scripts do, right after loading it) to catch drift immediately + instead of silently producing a request that runs to completion with the + wrong stop tokens. Not wired into ``build_prompt_tokens`` itself: callers + that intentionally pass a stand-in tokenizer with different ids (e.g. + this module's own unit tests) must not be forced through this check. + + Both ``tencent/HunyuanImage-3.0`` (base) and ``-Instruct`` report the + same ``model_class_name`` (``HunyuanImage3ForCausalMM``), so this hook + fires for either checkpoint -- but the base tokenizer's vocab doesn't + include the Instruct-only ````-```` tokens + (added for it2i's second aspect-ratio range). A token the tokenizer + doesn't recognize at all (``convert_tokens_to_ids`` returning + ``unk_token_id``, or ``None`` for tokenizers without an unk concept) is + "not part of this checkpoint's vocab," not "drifted"; only an id that + resolves to something *else* indicates the hardcoded table is stale. + """ + unk_id = getattr(tokenizer, "unk_token_id", None) + mismatches: dict[str, tuple[int, int]] = {} + for tok, expected_id in HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS.items(): + actual_id = tokenizer.convert_tokens_to_ids(tok) + if actual_id is None or (unk_id is not None and actual_id == unk_id): + continue # not in this checkpoint's vocab -- nothing to compare + if actual_id != expected_id: + mismatches[tok] = (expected_id, actual_id) + if mismatches: + details = ", ".join(f"{tok!r}: expected {exp}, got {act}" for tok, (exp, act) in mismatches.items()) + raise ValueError( + "HunyuanImage-3.0 tokenizer special-token ids no longer match " + f"HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS ({details}). The loaded " + "tokenizer/model revision has likely changed; update the hardcoded " + "table in prompt_utils.py to match." + ) + + # bot_task -> (sys_type, trigger_tag). # ``vanilla`` is special-cased downstream: it bypasses the chat template # (no ``User:`` / ``Assistant:`` framing) and is only valid with @@ -317,12 +359,106 @@ def build_prompt_tokens( ) +@dataclass +class ARPromptInputs: + """Fully-resolved AR-stage inputs for one HunyuanImage3 request. + + Consolidates the three things every caller (offline example, OpenAI + server) must compute before submitting an AR request: the AR prefill + (as segment-tokenized ``prompt_token_ids`` when a tokenizer is + available, else a formatted ``prompt`` string), the resolved DiT + system-prompt type, and the AR ``stop_token_ids``. + """ + + prompt: str | None + prompt_token_ids: list[int] | None + system_prompt_type: str + stop_token_ids: list[int] + + +def build_ar_prompt_inputs( + user_prompt: str, + task: str = "t2i", + bot_task: str | None | _DefaultBotTask = _DEFAULT_BOT_TASK, + sys_type: str | None = None, + custom_system_prompt: str | None = None, + num_images: int = 1, + tokenizer: Any | None = None, + image_size: str | None = None, +) -> ARPromptInputs: + """Build the complete AR-stage inputs for a HunyuanImage3 request. + + Wraps :func:`build_prompt_tokens` (byte-for-byte HF parity when + ``tokenizer`` is given), the :func:`build_prompt` string fallback (when + no tokenizer is available), and :func:`resolve_stop_token_ids` behind + one call, so a caller doesn't have to re-derive the AR prefill/stop-token + logic -- and, critically, doesn't have to independently resolve + ``bot_task`` for the prompt-building call vs. the stop-token call, which + can silently diverge if a caller passes the same nominal ``bot_task`` + value to both but resolves "omitted" vs. "explicit None" differently for + each (this seam always forwards the identical resolved value to both). + + Used by all four offline shared task examples: ``text_to_image.py`` / + ``image_edit.py`` (via ``model_extras.hunyuan_image3.build_ar_stage_inputs``) + and ``x_to_text.py`` (via ``build_x_to_text_prompt``, which also calls + ``build_ar_stage_inputs``). The OpenAI server's ``serving_chat.py`` does + not currently call this seam -- it still independently calls + ``build_prompt``/``build_prompt_tokens``/``resolve_stop_token_ids``, and + can hit exactly the divergence described above when ``bot_task`` is + omitted from ``extra_body`` (pre-existing, unrelated to this module). + ``image_size`` follows the same convention as :func:`resolve_stop_token_ids` + (``None``/``"auto"`` lets the AR predict the aspect ratio; an explicit + ``"{w}x{h}"`` makes it stop at the terminator). + """ + if tokenizer is not None: + result = build_prompt_tokens( + user_prompt, + tokenizer, + task=task, + bot_task=bot_task, + sys_type=sys_type, + custom_system_prompt=custom_system_prompt, + num_images=num_images, + ) + prompt_str: str | None = None + prompt_token_ids: list[int] | None = result.token_ids + system_prompt_type = result.system_prompt_type + else: + prompt_str = build_prompt( + user_prompt, + task=task, + bot_task=bot_task, + sys_type=sys_type, + custom_system_prompt=custom_system_prompt, + num_images=num_images, + ) + prompt_token_ids = None + _, resolved_bot_task = _normalize_task_and_bot_task(task, bot_task) + system_prompt_type = sys_type or resolve_sys_type(resolved_bot_task) + + stop_token_ids = resolve_stop_token_ids( + task=task, + bot_task=bot_task, + tokenizer=tokenizer, + image_size=image_size, + ) + return ARPromptInputs( + prompt=prompt_str, + prompt_token_ids=prompt_token_ids, + system_prompt_type=system_prompt_type, + stop_token_ids=stop_token_ids, + ) + + __all__ = [ "HUNYUAN_IMAGE3_SPECIAL_TOKEN_IDS", + "validate_special_token_ids", "MAX_IMAGES_PER_REQUEST", "_TASK_PRESETS", + "ARPromptInputs", "available_bot_tasks", "available_tasks", + "build_ar_prompt_inputs", "build_prompt", "build_prompt_tokens", "resolve_stop_token_ids", diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index f5cb24efd29..20bddaba67d 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -3080,11 +3080,20 @@ def _build_multistage_generation_inputs( ar_image_size: str | None = None if height is not None and width is not None: ar_image_size = f"{width}x{height}" + # Reuse build_kwargs's own (possibly-omitted) "bot_task" entry rather + # than the raw outer `bot_task` variable: build_prompt_tokens/build_prompt + # above default an omitted bot_task per-task (e.g. "think" for t2i), and + # resolve_stop_token_ids must normalize from that same omitted-or-not + # starting point to agree on which stop tokens apply -- passing the raw + # `bot_task` (still None when omitted) made the two calls disagree. + ar_stop_kwargs: dict[str, Any] = {} + if "bot_task" in build_kwargs: + ar_stop_kwargs["bot_task"] = build_kwargs["bot_task"] ar_stop_token_ids = resolve_stop_token_ids( task=ar_task, - bot_task=bot_task, tokenizer=tokenizer, image_size=ar_image_size, + **ar_stop_kwargs, ) engine_prompt: OmniTextPrompt = {"prompt": prompt} diff --git a/vllm_omni/model_extras/__init__.py b/vllm_omni/model_extras/__init__.py index 143a9b22451..2e1146a8292 100644 --- a/vllm_omni/model_extras/__init__.py +++ b/vllm_omni/model_extras/__init__.py @@ -6,6 +6,8 @@ build_image_to_video_prompt, build_text_to_image_prompt, build_x_to_text_prompt, + get_ar_input_builder, + get_ar_tokenizer_validator, get_extra_body_params, get_extra_output_params, get_model_class_name, @@ -22,6 +24,8 @@ "build_image_to_video_prompt", "build_text_to_image_prompt", "build_x_to_text_prompt", + "get_ar_input_builder", + "get_ar_tokenizer_validator", "get_extra_body_params", "get_extra_output_params", "get_model_class_name", diff --git a/vllm_omni/model_extras/hunyuan_image3.py b/vllm_omni/model_extras/hunyuan_image3.py index 8a53dfd54c8..46d48b51355 100644 --- a/vllm_omni/model_extras/hunyuan_image3.py +++ b/vllm_omni/model_extras/hunyuan_image3.py @@ -3,6 +3,8 @@ from __future__ import annotations +from collections.abc import Mapping +from dataclasses import dataclass, field from typing import Any @@ -11,27 +13,162 @@ def build_x_to_text_prompt( prompt: str, has_image: bool, ) -> tuple[dict[str, Any], list[int] | None]: - """Build HunyuanImage-3 prompt tokens and stopping rules for T2T/I2T.""" - from transformers import AutoTokenizer + """Build HunyuanImage-3 prompt tokens and stopping rules for T2T/I2T. - from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( - build_prompt_tokens, - resolve_stop_token_ids, - resolve_sys_type, - ) + Routes through :func:`build_ar_stage_inputs`, the same declarative seam + the shared text_to_image / image_edit examples use, instead of + re-deriving the AR prefill and stop-token logic here. Unlike those + examples -- generic scripts that reach ``validate_ar_tokenizer`` only + indirectly via ``get_ar_tokenizer_validator`` -- this function is already + HunyuanImage3-specific, so it calls the validator directly rather than + requiring the (model-agnostic) ``x_to_text.py`` caller to know about it. + """ + from transformers import AutoTokenizer - task = "i2t" if has_image else "t2t" tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True) - build_kwargs: dict[str, Any] = {"task": task, "bot_task": None} - if has_image: - build_kwargs["num_images"] = 1 - built = build_prompt_tokens(prompt, tokenizer, **build_kwargs) + validate_ar_tokenizer(tokenizer) + ar_inputs = build_ar_stage_inputs( + prompt=prompt, + tokenizer=tokenizer, + extra_body=None, + num_images=1 if has_image else 0, + text_output=True, + ) return ( { - "prompt": prompt, - "prompt_token_ids": built.token_ids, - "modalities": ["text"], - "use_system_prompt": resolve_sys_type(None), + "prompt": ar_inputs.prompt, + "prompt_token_ids": ar_inputs.prompt_token_ids, + "modalities": ar_inputs.modalities, + "use_system_prompt": ar_inputs.use_system_prompt, }, - resolve_stop_token_ids(task=task, bot_task=None, tokenizer=tokenizer, image_size="auto"), + ar_inputs.stop_token_ids, ) + + +@dataclass +class ARStageInputs: + """AR-stage inputs the shared example applies to a HunyuanImage3 request. + + ``prompt_token_ids`` (byte-for-byte HF parity) is preferred; ``prompt`` + is the string fallback used when no tokenizer is available. ``modalities`` + picks the AR output type (``["image"]`` for generation, ``["text"]`` for + comprehension). ``use_system_prompt`` feeds the DiT system-prompt prefix + and ``stop_token_ids`` terminate the AR decode. + + ``stage_indices`` names which stage(s) in the request's + ``sampling_params_list`` these ``stop_token_ids`` belong to -- the shared + scripts apply them by explicit index, not by scanning for "whichever + stage isn't a diffusion stage." HunyuanImage-3.0's topology always puts + its single AR stage at index 0 (see + ``vllm_omni.model_executor.models.hunyuan_image3.pipeline``); a future + model with multiple AR/understanding stages would declare its own + indices here instead of relying on type-based inference. + """ + + prompt: str | None + prompt_token_ids: list[int] | None + modalities: list[str] + use_system_prompt: str + stop_token_ids: list[int] + stage_indices: list[int] = field(default_factory=lambda: [0]) + + +def build_ar_stage_inputs( + prompt: str, + tokenizer: Any | None, + extra_body: Mapping[str, Any] | None, + *, + num_images: int = 0, + height: int | None = None, + width: int | None = None, + text_output: bool = False, +) -> ARStageInputs: + """Resolve HunyuanImage-3.0 AR-stage inputs declaratively. + + Invoked generically by the shared task examples (``text_to_image`` / + ``image_to_image`` / ``x_to_text``) whenever the model declares an + ``ar_input_builder``. All model-specific knobs (``bot_task`` / + ``use_system_prompt`` / ``system_prompt``) arrive via ``extra_body``, so + the example itself stays model-agnostic. The heavy lifting is delegated to + :func:`build_ar_prompt_inputs`. The OpenAI server's ``serving_chat.py`` + does not go through this seam -- it independently calls the same + underlying ``build_prompt``/``build_prompt_tokens``/``resolve_stop_token_ids`` + primitives (see that function's docstring for the resulting divergence risk). + """ + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + build_ar_prompt_inputs, + ) + + extra_body = extra_body or {} + has_image = num_images > 0 + if text_output: + task = "i2t" if has_image else "t2t" + else: + task = "it2i" if has_image else "t2i" + + # ar_image_size: None -> AR predicts aspect ratio; explicit "{w}x{h}" -> + # AR stops at the terminator. Comprehension (text) tasks always auto-size. + ar_image_size: str | None = None + if not text_output and height is not None and width is not None: + ar_image_size = f"{width}x{height}" + + kwargs: dict[str, Any] = { + "task": task, + "sys_type": extra_body.get("use_system_prompt"), + "custom_system_prompt": extra_body.get("system_prompt"), + "num_images": max(num_images, 1) if has_image else 1, + "tokenizer": tokenizer, + "image_size": ar_image_size, + } + # Only forward bot_task when the caller set it: present (incl. explicit + # ``None`` -> plain mode) overrides; absent lets each task pick its default + # trigger via build_ar_prompt_inputs' own sentinel default. + if "bot_task" in extra_body: + kwargs["bot_task"] = extra_body["bot_task"] + + resolved = build_ar_prompt_inputs(prompt, **kwargs) + return ARStageInputs( + prompt=resolved.prompt, + prompt_token_ids=resolved.prompt_token_ids, + modalities=["text"] if text_output else ["image"], + use_system_prompt=resolved.system_prompt_type, + stop_token_ids=resolved.stop_token_ids, + # HunyuanImage-3.0's AR stage is always stage 0 (see pipeline.py's + # HUNYUAN_IMAGE3_PIPELINE / _AR_PIPELINE topologies). + stage_indices=[0], + ) + + +def validate_ar_tokenizer(tokenizer: Any) -> None: + """Registry-declared ``ar_tokenizer_validator`` hook for HunyuanImage3. + + Called by the shared task examples right after they load a *real* + tokenizer for the AR stage, so a model/tokenizer revision bump that + silently shifts special-token ids fails loudly instead of producing a + request that completes with the wrong stop tokens. + """ + from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import ( + validate_special_token_ids, + ) + + validate_special_token_ids(tokenizer) + + +# Per-request, model-specific knobs (non-standard sampling-param fields). +# Standard knobs (guidance_scale / num_inference_steps / seed / height / width) +# stay on OmniDiffusionSamplingParams; engine-level knobs +# (diffusion_kv_cache_* / vae_use_tiling) stay as Omni() init args. +HUNYUAN_IMAGE3_EXTRA_BODY_PARAMS = frozenset( + { + "bot_task", + "use_system_prompt", + "system_prompt", + "negative_prompt", + } +) +# Text outputs for i2t/t2t are surfaced through the standard AR text path; +# no diffusion custom-output params are declared. +HUNYUAN_IMAGE3_EXTRA_OUTPUT_PARAMS = frozenset() +# HunyuanImage3 runs an AR stage before DiT; its extra_args must be initialised +# so bot_task / use_system_prompt reach the AR input path. +HUNYUAN_IMAGE3_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES = True diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index 4d2ad010aa3..f21691eab39 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -28,7 +28,18 @@ HELIOS_EXTRA_BODY_PARAMS, HELIOS_EXTRA_OUTPUT_PARAMS, ) +from vllm_omni.model_extras.hunyuan_image3 import ( + HUNYUAN_IMAGE3_EXTRA_BODY_PARAMS, + HUNYUAN_IMAGE3_EXTRA_OUTPUT_PARAMS, + HUNYUAN_IMAGE3_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES, +) +from vllm_omni.model_extras.hunyuan_image3 import ( + build_ar_stage_inputs as build_hunyuan_image3_ar_stage_inputs, +) from vllm_omni.model_extras.hunyuan_image3 import build_x_to_text_prompt as build_hunyuan_x_to_text_prompt +from vllm_omni.model_extras.hunyuan_image3 import ( + validate_ar_tokenizer as validate_hunyuan_image3_ar_tokenizer, +) from vllm_omni.model_extras.lingbot_video import LINGBOT_VIDEO_EXTRA_BODY_PARAMS from vllm_omni.model_extras.ltx2 import ( LTX_EXTRA_BODY_PARAMS, @@ -233,6 +244,18 @@ def default_image_to_image_prompt( "LTX2DistilledTwoStagePipeline", ) }, + # HunyuanImage3 keys on the model architecture name that ``get_model_class_name`` + # reports at runtime for its AR+DiT / AR-only / DiT-only deploys + # (``od_config.model_class_name`` == the diffusion registry arch key). The + # "HunyuanImage3Pipeline" alias is kept for any path that surfaces the + # pipeline class name instead. + "HunyuanImage3ForCausalMM": { + "extra_body_params": HUNYUAN_IMAGE3_EXTRA_BODY_PARAMS, + "extra_output_params": HUNYUAN_IMAGE3_EXTRA_OUTPUT_PARAMS, + "init_extra_args_for_non_diffusion_stages": HUNYUAN_IMAGE3_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES, + "ar_input_builder": build_hunyuan_image3_ar_stage_inputs, + "ar_tokenizer_validator": validate_hunyuan_image3_ar_tokenizer, + }, "SanaVideoPipeline": { "extra_body_params": SANA_VIDEO_EXTRA_BODY_PARAMS, }, @@ -263,6 +286,10 @@ def default_image_to_image_prompt( _EXTRA_SPECS[model_class_name]["transformer_config_subfolder_resolver"] = ltx_transformer_config_subfolder +# Alias: some code paths surface HunyuanImage3's pipeline class name rather than +# the architecture name; point both at the same spec. +_EXTRA_SPECS["HunyuanImage3Pipeline"] = _EXTRA_SPECS["HunyuanImage3ForCausalMM"] + # Multi-stage discovery reports the top-level wrapper rather than its DiT # submodule, so both names must resolve to the same request builders. _EXTRA_SPECS["MammothModa2ForConditionalGeneration"] = _EXTRA_SPECS["MammothModa2DiTPipeline"] @@ -358,6 +385,32 @@ def should_init_extra_args_for_non_diffusion_stages(model_class_name: str | None return bool(spec and spec.get("init_extra_args_for_non_diffusion_stages", False)) +def get_ar_input_builder(model_class_name: str | None) -> Callable[..., Any] | None: + """Return a model's AR-stage input builder, or ``None`` if undeclared. + + Models with a text/AR stage that needs template-formatted prompt tokens + and AR stop tokens (e.g. HunyuanImage3) declare an ``ar_input_builder``. + The shared task examples call it generically when present, so the example + scripts stay model-agnostic; models without one are unaffected. + """ + spec = _get_spec(model_class_name) + return spec.get("ar_input_builder") if spec is not None else None + + +def get_ar_tokenizer_validator(model_class_name: str | None) -> Callable[[Any], None] | None: + """Return a model's AR-tokenizer validator, or ``None`` if undeclared. + + Models whose AR prompt/stop-token logic depends on hardcoded special + token ids (e.g. HunyuanImage3) declare an ``ar_tokenizer_validator`` to + check those ids against whatever tokenizer actually loads at runtime. + The shared task examples call it generically, right after loading a real + tokenizer, so model/tokenizer revision drift fails loudly instead of + silently producing a request with the wrong stop tokens. + """ + spec = _get_spec(model_class_name) + return spec.get("ar_tokenizer_validator") if spec is not None else None + + def build_text_to_image_prompt( model_class_name: str | None, prompt: dict[str, Any], From 3f7c2174574f14a9767c022b01834fa2c52420df Mon Sep 17 00:00:00 2001 From: hyw <109567717+yuweih205@users.noreply.github.com> Date: Fri, 11 Sep 2026 17:09:26 +0800 Subject: [PATCH 71/93] [Model] Avoid scalar synchronizations in GLM-Image preparation (#7172) Signed-off-by: hyw --- .../glm_image/test_glm_image_quantization.py | 5 ++- .../models/glm_image/glm_image_transformer.py | 32 ++++++------------- 2 files changed, 11 insertions(+), 26 deletions(-) diff --git a/tests/diffusion/models/glm_image/test_glm_image_quantization.py b/tests/diffusion/models/glm_image/test_glm_image_quantization.py index 1a6170ca20c..dfcd15b1a86 100644 --- a/tests/diffusion/models/glm_image/test_glm_image_quantization.py +++ b/tests/diffusion/models/glm_image/test_glm_image_quantization.py @@ -399,15 +399,14 @@ def test_prepare_module_exists(self): prepare = GlmImagePrepare( image_projector=projector, rope=rope, - patch_size=2, ) # Test forward pass hidden_states = torch.randn(1, 16, 64, 64) # [B, C, H, W] result = prepare(hidden_states) - assert len(result) == 5 - hidden_out, rope_cos, rope_sin, height, width = result + assert len(result) == 3 + hidden_out, rope_cos, rope_sin = result assert hidden_out.shape[0] == 1 # batch size assert hidden_out.shape[1] == 1024 # seq_len (32 * 32) assert hidden_out.shape[2] == 2560 # hidden_dim diff --git a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py index 75ba6020381..1b08684007a 100644 --- a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py +++ b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.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 math from collections.abc import Iterable @@ -124,8 +124,8 @@ def __init__( def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: batch_size, channel, height, width = hidden_states.shape - post_patch_height = torch.tensor(height // self.patch_size, device=hidden_states.device, dtype=torch.int64) - post_patch_width = torch.tensor(width // self.patch_size, device=hidden_states.device, dtype=torch.int64) + post_patch_height = height // self.patch_size + post_patch_width = width // self.patch_size # Reshape: [B, C, H, W] -> [B, H', W', C*p*p] -> [B, H'*W', C*p*p] hidden_states = hidden_states.reshape( @@ -189,18 +189,16 @@ def __init__( self, image_projector: nn.Module, rope: GlmImageRotaryPosEmbed, - patch_size: int, ): super().__init__() self.image_projector = image_projector self.rope = rope - self.patch_size = patch_size def forward( self, hidden_states: torch.Tensor, prior_hidden_states: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Process hidden_states and compute RoPE embeddings. Args: @@ -211,14 +209,7 @@ def forward( hidden_states: Patched hidden states [B, seq_len, D] rope_cos: RoPE cos embeddings [seq_len, dim] rope_sin: RoPE sin embeddings [seq_len, dim] - post_patch_height: Scalar tensor for height after patching - post_patch_width: Scalar tensor for width after patching """ - batch_size, num_channels, height, width = hidden_states.shape - - post_patch_height = torch.tensor(height // self.patch_size, device=hidden_states.device, dtype=torch.int64) - post_patch_width = torch.tensor(width // self.patch_size, device=hidden_states.device, dtype=torch.int64) - # Compute RoPE (uses original 4D hidden_states shape) image_rotary_emb = self.rope(hidden_states) rope_cos = image_rotary_emb[0].to(hidden_states.device) @@ -231,7 +222,7 @@ def forward( if prior_hidden_states is not None: hidden_states = hidden_states + prior_hidden_states - return hidden_states, rope_cos, rope_sin, post_patch_height, post_patch_width + return hidden_states, rope_cos, rope_sin class GlmImageAdaLayerNormZero(nn.Module): @@ -961,7 +952,6 @@ class GlmImageTransformer2DModel(CachedTransformer): 1: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True), # RoPE sin: [seq_len, dim] - shard along sequence dimension 2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True), - # post_patch_height and post_patch_width are scalars, not sharded }, # Gather output at proj_out "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), @@ -1044,7 +1034,7 @@ def __init__( ) # Prepare module for SP (encapsulates patch embedding and RoPE for _sp_plan) - self.prepare = GlmImagePrepare(self.image_projector, self.rope, patch_size) + self.prepare = GlmImagePrepare(self.image_projector, self.rope) self.time_condition_embed = GlmImageCombinedTimestepSizeEmbeddings( embedding_dim=time_embed_dim, @@ -1137,12 +1127,10 @@ def forward( # 1. Prepare hidden_states and RoPE via GlmImagePrepare module # _sp_plan will shard hidden_states and RoPE together via split_output=True - hidden_states, rope_cos, rope_sin, post_patch_height_t, post_patch_width_t = self.prepare( - hidden_states, prior_hidden_states - ) + hidden_states, rope_cos, rope_sin = self.prepare(hidden_states, prior_hidden_states) image_rotary_emb = (rope_cos, rope_sin) - post_patch_height = int(post_patch_height_t.item()) - post_patch_width = int(post_patch_width_t.item()) + post_patch_height = height // self.patch_size + post_patch_width = width // self.patch_size # Timestep conditioning temb = self.time_condition_embed(timestep, target_size, crop_coords, hidden_states.dtype) @@ -1163,8 +1151,6 @@ def forward( device=hidden_states.device, ) hidden_states_mask[:, ctx.sp_original_seq_len :] = False - if hidden_states_mask.all(): - hidden_states_mask = None # 2. Transformer blocks for layer_idx, block in enumerate(self.transformer_blocks): From f62d4ad4a794bd0706340d8dcc9d92e67672af1c Mon Sep 17 00:00:00 2001 From: hyw <109567717+yuweih205@users.noreply.github.com> Date: Fri, 11 Sep 2026 17:15:32 +0800 Subject: [PATCH 72/93] [Model][ERNIE-Image] Delay AdaLN modulation broadcast (#7171) Signed-off-by: hyw --- .../models/ernie_image/ernie_image_transformer.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py b/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py index 244aeeb95f8..bd96630e8d4 100644 --- a/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py +++ b/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.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 # Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_ernie_image.py from collections.abc import Iterable @@ -599,17 +599,12 @@ def forward( hidden_states, freqs_cos, freqs_sin, attention_mask = self.unified_prepare(hidden_states, text_bth, text_lens) attention_mask = self._slice_attention_mask_for_ring(attention_mask) rotary_pos_emb = (freqs_cos, freqs_sin) - S = hidden_states.shape[0] - sample = self.time_proj(timestep) sample = sample.to(dtype=dtype) c = self.time_embedding(sample) - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = [ - t.unsqueeze(0).expand(S, -1, -1).contiguous() for t in self.adaLN_modulation(c).chunk(6, dim=-1) - ] + temb = self.adaLN_modulation(c).chunk(6, dim=-1) for layer in self.layers: - temb = [shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp] if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = self._gradient_checkpointing_func( layer, From 43b8de9b0ecdeb78ec3585222db02164a087103a Mon Sep 17 00:00:00 2001 From: hyw <109567717+yuweih205@users.noreply.github.com> Date: Fri, 11 Sep 2026 17:18:12 +0800 Subject: [PATCH 73/93] [Kernel][MiniMax-H3] Run Q/K RMSNorm-RoPE in one launch (#7167) Signed-off-by: hyw --- .../models/minimax_h3/ops/vae/qk_norm_rope.py | 76 +++++++++++-------- 1 file changed, 45 insertions(+), 31 deletions(-) diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py index 40b7069f61c..f964c15772b 100644 --- a/vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py +++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/qk_norm_rope.py @@ -11,18 +11,26 @@ @triton.jit def _qk_rms_norm_rope_exact_kernel( - x_ptr, + q_ptr, + k_ptr, cos_ptr, sin_ptr, - output_ptr, - x_stride_token, - x_stride_head, - x_stride_dim, + q_output_ptr, + k_output_ptr, + q_stride_token, + q_stride_head, + q_stride_dim, + k_stride_token, + k_stride_head, + k_stride_dim, rope_stride_token, rope_stride_dim, - output_stride_token, - output_stride_head, - output_stride_dim, + q_output_stride_token, + q_output_stride_head, + q_output_stride_dim, + k_output_stride_token, + k_output_stride_head, + k_output_stride_dim, num_heads: tl.constexpr, head_dim: tl.constexpr, rotary_dim: tl.constexpr, @@ -31,12 +39,15 @@ def _qk_rms_norm_rope_exact_kernel( ): token = tl.program_id(0) head_group = tl.program_id(1) + qk_domain = tl.program_id(2) heads = head_group * heads_per_program + tl.arange(0, heads_per_program) dims = tl.arange(0, head_dim) mask = heads[:, None] < num_heads - offsets = token * x_stride_token + heads[:, None] * x_stride_head + dims[None, :] * x_stride_dim - x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + q_offsets = token * q_stride_token + heads[:, None] * q_stride_head + dims[None, :] * q_stride_dim + k_offsets = token * k_stride_token + heads[:, None] * k_stride_head + dims[None, :] * k_stride_dim + x_ptrs = tl.where(qk_domain == 0, q_ptr + q_offsets, k_ptr + k_offsets) + x = tl.load(x_ptrs, mask=mask, other=0.0).to(tl.float32) inv_rms = tl.rsqrt(tl.sum(x * x, axis=1) / head_dim + eps) normalized = (x * inv_rms[:, None]).to(tl.float16) @@ -46,8 +57,10 @@ def _qk_rms_norm_rope_exact_kernel( dims + rotary_half, tl.where(dims < rotary_dim, dims - rotary_half, dims), ) - pair_offsets = token * x_stride_token + heads[:, None] * x_stride_head + pair_dims[None, :] * x_stride_dim - pair_x = tl.load(x_ptr + pair_offsets, mask=mask, other=0.0).to(tl.float32) + q_pair_offsets = token * q_stride_token + heads[:, None] * q_stride_head + pair_dims[None, :] * q_stride_dim + k_pair_offsets = token * k_stride_token + heads[:, None] * k_stride_head + pair_dims[None, :] * k_stride_dim + pair_x_ptrs = tl.where(qk_domain == 0, q_ptr + q_pair_offsets, k_ptr + k_pair_offsets) + pair_x = tl.load(pair_x_ptrs, mask=mask, other=0.0).to(tl.float32) pair_normalized = (pair_x * inv_rms[:, None]).to(tl.float16) rope_mask = dims < rotary_dim @@ -79,10 +92,18 @@ def _qk_rms_norm_rope_exact_kernel( ) output = tl.where(rope_mask[None, :], rotated, normalized) - output_offsets = ( - token * output_stride_token + heads[:, None] * output_stride_head + dims[None, :] * output_stride_dim + q_output_offsets = ( + token * q_output_stride_token + heads[:, None] * q_output_stride_head + dims[None, :] * q_output_stride_dim ) - tl.store(output_ptr + output_offsets, output, mask=mask) + k_output_offsets = ( + token * k_output_stride_token + heads[:, None] * k_output_stride_head + dims[None, :] * k_output_stride_dim + ) + output_ptrs = tl.where( + qk_domain == 0, + q_output_ptr + q_output_offsets, + k_output_ptr + k_output_offsets, + ) + tl.store(output_ptrs, output, mask=mask) def _supported_inputs( @@ -143,11 +164,11 @@ def try_qk_norm_rope_exact( sin = sin.reshape(tokens, sin.shape[-1]) q_output = torch.empty_like(q) k_output = torch.empty_like(k) - # This launch layout is part of the exactness contract on the validated - # SM90 and SM103 targets: changing the warp-to-row mapping changes the - # FP32 RMS reduction order. + # The first two launch dimensions are part of the exactness contract on + # the validated SM90 and SM103 targets: changing the warp-to-row mapping + # changes the FP32 RMS reduction order. heads_per_program = 8 - grid = (tokens, triton.cdiv(heads, heads_per_program)) + grid = (tokens, triton.cdiv(heads, heads_per_program), 2) launch_args = { "num_heads": heads, "head_dim": head_dim, @@ -158,29 +179,22 @@ def try_qk_norm_rope_exact( } _qk_rms_norm_rope_exact_kernel[grid]( q, + k, cos, sin, q_output, + k_output, q.stride(0), q.stride(1), q.stride(2), - cos.stride(0), - cos.stride(1), - q_output.stride(0), - q_output.stride(1), - q_output.stride(2), - **launch_args, - ) - _qk_rms_norm_rope_exact_kernel[grid]( - k, - cos, - sin, - k_output, k.stride(0), k.stride(1), k.stride(2), cos.stride(0), cos.stride(1), + q_output.stride(0), + q_output.stride(1), + q_output.stride(2), k_output.stride(0), k_output.stride(1), k_output.stride(2), From dff8f8f4cfe9ba6e6eaeeac455f8aee66f95c4ae Mon Sep 17 00:00:00 2001 From: andyluo7 <43718156+andyluo7@users.noreply.github.com> Date: Fri, 11 Sep 2026 02:41:45 -0700 Subject: [PATCH 74/93] [CI][ROCm] Align AMD image with vLLM 0.29 (#7395) Signed-off-by: andyluo7 Co-authored-by: Hongsheng Liu --- docker/Dockerfile.rocm | 24 +++--- .../installation/gpu/rocm.inc.md | 44 ++++++++++- tests/buildkite/test_rocm_dockerfile.py | 77 +++++++++++++++++++ 3 files changed, 132 insertions(+), 13 deletions(-) create mode 100644 tests/buildkite/test_rocm_dockerfile.py diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 597f5d13414..adc76769141 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -1,16 +1,18 @@ -ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.28.0 +# Keep the ROCm runtime aligned with the vLLM release targeted by +# docker/Dockerfile.ci. A stale base image can build successfully while failing +# at runtime when vLLM-Omni imports APIs added by the newer vLLM release. +ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.29.0 FROM ${BASE_IMAGE} AS base -# Declare a variable to know if we want to use the nightly build or the stable build. # NOTE: REMINDER to vLLM-Omni rebase maintainer -# Remember to set `USE_NIGHTLY_BUILD` to 0, when switching back to stable vLLM docker image -# 1. If vLLM-omni maintainer is forced to use custom commits -# during rebasing, they can change the two variables -# 2. Whenever vLLM upstream has released a stable version, -# we should swap over to use stable release ASAP. -# We should avoid relying on custom commits. +# BASE_IMAGE pins the default prebuilt runtime when USE_NIGHTLY_BUILD=0. +# VLLM_VERSION_OR_COMMIT_HASH is consumed only by the opt-in source rebuild when +# USE_NIGHTLY_BUILD=1. Keep both aligned with docker/Dockerfile.ci's +# VLLM_BASE_TAG. This is enforced by tests/buildkite/test_rocm_dockerfile.py. An +# intentional custom source pin must update that regression in the same change, +# and should return to the stable release as soon as one is available. ARG USE_NIGHTLY_BUILD=0 -ARG VLLM_VERSION_OR_COMMIT_HASH=v0.28.0 +ARG VLLM_VERSION_OR_COMMIT_HASH=v0.29.0 ARG ARG_PYTORCH_ROCM_ARCH ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}} @@ -38,6 +40,10 @@ RUN if [ "${USE_NIGHTLY_BUILD}" = "1" ]; then \ rm -rf vllm/.git; \ fi +# Fail during image construction instead of letting every AMD test fail at +# runtime when the base image and vLLM-Omni target incompatible vLLM APIs. +RUN python3 -c "import vllm; from vllm.v1.kv_cache_interface import compute_layout_strides; from vllm.v1.kv_cache_layout import KVCacheLayout; print(f'Validated vLLM {vllm.__version__}')" + # Build TorchCodec after any optional nightly vLLM reinstall so it targets the # final torch/ROCm ABI. Keep this layer before COPY . so source changes do not # invalidate the expensive native build. FFmpeg development headers are diff --git a/docs/getting_started/installation/gpu/rocm.inc.md b/docs/getting_started/installation/gpu/rocm.inc.md index d15bba2eb44..83efe0c62d5 100644 --- a/docs/getting_started/installation/gpu/rocm.inc.md +++ b/docs/getting_started/installation/gpu/rocm.inc.md @@ -1,11 +1,23 @@ + + # --8<-- [start:requirements] - GPU: Validated on gfx942 (It should be supported on the AMD GPUs that are supported by vLLM.) # --8<-- [end:requirements] + # --8<-- [start:set-up-using-python] -vLLM-Omni current recommends the steps in under setup through Docker Images. +For ROCm, vLLM-Omni currently recommends the setup steps through Docker Images. + +vLLM-Omni depends on the matching major/minor release of vLLM. The 0.29 +development line uses vLLM 0.29.x. Published 0.28.0 wheels and images use vLLM +0.28.x. + +The Dockerfile's `BASE_IMAGE` pin applies only to Docker builds. The +`vllm-omni` package does not install vLLM as a dependency, so non-Docker source +installs must install the matching ROCm vLLM release explicitly before +installing vLLM-Omni, as shown below. # --8<-- [start:pre-built-wheels] @@ -14,6 +26,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images. These pre-built wheel instructions install the published vLLM-Omni 0.28.0 release. For the 0.29 development line, use the source-install instructions below. vLLM-Omni is built based on vLLM. Please install it with command below. + ```bash uv pip install vllm==0.28.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.28.0/rocm723 ``` @@ -36,6 +49,7 @@ uv pip install onnxruntime-rocm # --8<-- [start:build-wheel-from-source] #### Installation of vLLM + If you do not need to modify source code of vLLM, you can directly install the stable 0.29.0 release version of the library ```bash @@ -45,7 +59,9 @@ uv pip install vllm==0.29.0+rocm723 --extra-index-url https://wheels.vllm.ai/roc The pre-built 0.29.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead. #### Installation of vLLM-Omni + Since vllm-omni is rapidly evolving, it's recommended to install it from source + ```bash git clone https://github.com/vllm-project/vllm-omni.git cd vllm-omni @@ -71,15 +87,31 @@ python3 setup.py develop #### Build docker image +The source build defaults to the published upstream base image +`vllm/vllm-openai-rocm:v0.29.0`, aligned with the vLLM release used by CI. +Older or custom bases must provide the vLLM APIs checked by the Dockerfile's +image-build canary. This upstream base is distinct from the prebuilt +`vllm/vllm-omni-rocm` images discussed below; their published-tag availability +is tracked in [#7405](https://github.com/vllm-project/vllm-omni/issues/7405). + ```bash DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.rocm -t vllm-omni-rocm . ``` +To select the upstream ROCm base explicitly, pass `BASE_IMAGE`: + +```bash +DOCKER_BUILDKIT=1 docker build \ + -f docker/Dockerfile.rocm \ + --build-arg BASE_IMAGE=vllm/vllm-openai-rocm:v0.29.0 \ + -t vllm-omni-rocm . +``` + #### Launch the docker image ##### Launch with OpenAI API Server -``` +```bash docker run --rm \ --group-add=video \ --ipc=host \ @@ -97,7 +129,7 @@ vllm-omni-rocm \ ##### Launch with interactive session for development -``` +```bash docker run --rm -it \ --network=host \ --group-add=video \ @@ -119,7 +151,9 @@ vllm-omni-rocm vLLM-Omni offers an official docker image for deployment. These images are built on top of vLLM docker images and available on Docker Hub as [vllm/vllm-omni-rocm](https://hub.docker.com/r/vllm/vllm-omni-rocm/tags). The version of vLLM-Omni indicates which release of vLLM it is based on. #### Launch vLLM-Omni Server + Here's an example deployment command that has been verified on 2 x MI300's: + ```bash docker run --rm \ --group-add=video \ @@ -136,8 +170,10 @@ docker run --rm \ --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` -#### Launch an interactive terminal with prebuilt docker image. +#### Launch an interactive terminal with prebuilt docker image + If you want to run in dev environment you can launch the docker image as follows: + ```bash docker run --rm -it \ --network=host \ diff --git a/tests/buildkite/test_rocm_dockerfile.py b/tests/buildkite/test_rocm_dockerfile.py new file mode 100644 index 00000000000..9171b3d53ac --- /dev/null +++ b/tests/buildkite/test_rocm_dockerfile.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +REPO_ROOT = Path(__file__).resolve().parents[2] +AMD_TEMPLATE = REPO_ROOT / ".buildkite/amd/test-template-amd-omni.j2" +CI_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.ci" +ROCM_DOCKERFILE = REPO_ROOT / "docker/Dockerfile.rocm" + + +def _docker_arg(path: Path, name: str) -> str: + prefix = f"ARG {name}=" + matches = [ + line.removeprefix(prefix) for line in path.read_text(encoding="utf-8").splitlines() if line.startswith(prefix) + ] + assert len(matches) == 1, f"expected one {name} declaration in {path}, found {len(matches)}" + return matches[0] + + +def _line_index(lines: list[str], prefix: str) -> int: + matches = [index for index, line in enumerate(lines) if line.startswith(prefix)] + assert len(matches) == 1, f"expected one line starting with {prefix!r}, found {len(matches)}" + return matches[0] + + +def test_rocm_base_tracks_ci_vllm_release() -> None: + ci_release = _docker_arg(CI_DOCKERFILE, "VLLM_BASE_TAG") + rocm_base = _docker_arg(ROCM_DOCKERFILE, "BASE_IMAGE") + + image_ref, separator, image_tag = rocm_base.rpartition(":") + assert separator, f"expected a tagged ROCm base image, got {rocm_base}" + assert image_ref.rsplit("/", 1)[-1] == "vllm-openai-rocm", image_ref + assert image_tag == ci_release + + +def test_rocm_defaults_to_prebuilt_base_image() -> None: + assert _docker_arg(ROCM_DOCKERFILE, "USE_NIGHTLY_BUILD") == "0" + + +def test_amd_build_uses_rocm_dockerfile_defaults() -> None: + template = AMD_TEMPLATE.read_text(encoding="utf-8") + build_commands = [line.strip() for line in template.splitlines() if '"docker build ' in line] + + assert len(build_commands) == 1, f"expected one AMD image build command, found {len(build_commands)}" + build_command = build_commands[0] + assert "-f docker/Dockerfile.rocm" in build_command + for arg_name in ("BASE_IMAGE", "USE_NIGHTLY_BUILD"): + assert f"--build-arg {arg_name}" not in build_command + assert f"--build-arg={arg_name}" not in build_command + + +def test_rocm_source_ref_tracks_ci_vllm_release() -> None: + ci_release = _docker_arg(CI_DOCKERFILE, "VLLM_BASE_TAG") + rocm_source_ref = _docker_arg(ROCM_DOCKERFILE, "VLLM_VERSION_OR_COMMIT_HASH") + + assert rocm_source_ref == ci_release + + +def test_rocm_dockerfile_contains_vllm_api_canary() -> None: + dockerfile = ROCM_DOCKERFILE.read_text(encoding="utf-8") + lines = dockerfile.splitlines() + nightly_start = _line_index(lines, 'RUN if [ "${USE_NIGHTLY_BUILD}" = "1" ]; then') + canary = ( + 'RUN python3 -c "import vllm; ' + "from vllm.v1.kv_cache_interface import compute_layout_strides; " + "from vllm.v1.kv_cache_layout import KVCacheLayout;" + ) + canary_index = _line_index(lines, canary) + + assert any(line.strip() == "fi" for line in lines[nightly_start + 1 : canary_index]), ( + "vLLM API canary must follow the optional nightly source reinstall block" + ) From de9a133d48cbcb35c7165c772524be20e0fd1bdd Mon Sep 17 00:00:00 2001 From: Guangjian Dong <163994576+Hiro208@users.noreply.github.com> Date: Fri, 11 Sep 2026 18:05:12 +0800 Subject: [PATCH 75/93] [Bugfix] Add embed_multimodal to MiniCPM-o 4.5 omni LLM class (#7384) Signed-off-by: Guangjian --- .../test_minicpmo_4_5_omni_llm_embed.py | 33 +++++++++++++++++++ .../minicpmo_4_5/minicpmo_4_5_omni_llm.py | 4 +++ 2 files changed, 37 insertions(+) create mode 100644 tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_embed.py diff --git a/tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_embed.py b/tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_embed.py new file mode 100644 index 00000000000..9f53fb7ece3 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_minicpmo_4_5_omni_llm_embed.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + MiniCPMO45OmniLLMForConditionalGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_embed_multimodal_delegates_to_get_multimodal_embeddings(monkeypatch): + """vLLM V1 encoder profiling calls embed_multimodal; without an override the + SupportsMultiModal protocol stub returns None and startup fails downstream.""" + sentinel = (object(),) + seen: dict[str, object] = {} + + def fake_get_multimodal_embeddings(self, **kwargs): + seen.update(kwargs) + return sentinel + + monkeypatch.setattr( + MiniCPMO45OmniLLMForConditionalGeneration, + "get_multimodal_embeddings", + fake_get_multimodal_embeddings, + ) + model = object.__new__(MiniCPMO45OmniLLMForConditionalGeneration) + + out = model.embed_multimodal(pixel_values="pv", audio_features="af") + + assert out is sentinel + assert seen == {"pixel_values": "pv", "audio_features": "af"} diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index d70bc30d0ab..6d1c5de320d 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -4580,6 +4580,10 @@ def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings: multimodal_embeddings += tuple(audio_embeddings) return multimodal_embeddings + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + """vLLM V1 encoder profiling calls this; the inherited Protocol stub returns None.""" + return self.get_multimodal_embeddings(**kwargs) + def get_input_embeddings( self, input_ids: torch.Tensor, From 02aaa34059478280eb3e6a5ee05a9484f0fa23c4 Mon Sep 17 00:00:00 2001 From: wtz2333 <2955110911@qq.com> Date: Fri, 11 Sep 2026 18:37:12 +0800 Subject: [PATCH 76/93] [Model] Add LingBot World Ulysses sequence parallelism (#6841) Signed-off-by: wtz2333 <2955110911@qq.com> Co-authored-by: Zhou Taichang --- docs/design/feature/realtime_ar_diffusion.md | 18 + .../diffusion/lingbot_world_v2.py | 6 +- .../test_lingbot_world_attention.py | 84 ++++- .../test_lingbot_world_transformer.py | 312 +++++++++++++++++- .../test_pipeline_lingbot_world.py | 51 ++- .../test_lingbot_world_v2.py | 24 ++ .../models/lingbot_world/pipeline.py | 42 ++- .../models/lingbot_world/transformer.py | 167 +++++++++- 8 files changed, 659 insertions(+), 45 deletions(-) create mode 100644 tests/examples/offline_inference/test_lingbot_world_v2.py diff --git a/docs/design/feature/realtime_ar_diffusion.md b/docs/design/feature/realtime_ar_diffusion.md index 5963befe7f6..426ee5afd73 100644 --- a/docs/design/feature/realtime_ar_diffusion.md +++ b/docs/design/feature/realtime_ar_diffusion.md @@ -178,6 +178,24 @@ For LingBot World v2: - `lingbot_world/pipeline.py` constructs conditioning, owns small non-KV session state, and produces one block plus the standard metadata envelope. +## Ulysses sequence parallelism + +LingBot supports pure Ulysses sequence parallelism for both direct and +AR-Diffusion execution. With Ulysses degree greater than one, hidden tokens, +camera features, token-expanded timestep modulation, and RoPE tables are sharded +together. Without SP, timestep modulation retains the frame-broadcast path. +Self-attention performs the sequence-to-head all-to-all before reading or writing +paged KV, while static text K/V uses the same local head shard. Text K/V shards +own compact storage; cross-attention exchanges query/output layouts to use these +shards. This retains the shared cache geometry and reduces text K/V storage at +the cost of two all-to-all calls per layer. + +The output head projects local tokens before gathering the flow values. For +the 14B model this reduces the gathered width from 5120 to 64; frame modulation +uses each shard's global token offset, including shards that split a frame. Only +`ulysses_mode="strict"` is supported; `advanced_uaa`, Ring, and AllGather-KV modes +remain unsupported for this model. + ## Non-goals This contract does not currently provide: diff --git a/examples/offline_inference/diffusion/lingbot_world_v2.py b/examples/offline_inference/diffusion/lingbot_world_v2.py index 590442c4d84..4bbea2f4f69 100644 --- a/examples/offline_inference/diffusion/lingbot_world_v2.py +++ b/examples/offline_inference/diffusion/lingbot_world_v2.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 """Generate a LingBot-World v2 video from an image and camera trajectory. The official checkpoint is licensed separately under CC BY-NC-SA and is @@ -99,6 +99,7 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: default=1, help="Number of GPUs used for tensor parallelism inside the DiT.", ) + parser.add_argument("--ulysses-degree", type=int, default=1, help="Pure Ulysses sequence parallel degree.") parser.add_argument("--flow-shift", type=float, default=5.0, help="Positive FlowUniPC scheduler shift.") parser.add_argument("--fps", type=int, default=16, help="Frames per second in the exported MP4.") parser.add_argument("--output", default="lingbot_world_v2.mp4", help="Output MP4 path.") @@ -158,6 +159,8 @@ def build_omni_kwargs( if args.tensor_parallel_size <= 0: raise ValueError("--tensor-parallel-size must be a positive integer.") + if args.ulysses_degree <= 0: + raise ValueError("--ulysses-degree must be a positive integer.") flow_shift = _positive_finite(args.flow_shift, "--flow-shift") model_path = Path(args.model).expanduser() model = str(model_path.resolve()) if model_path.exists() else args.model @@ -165,6 +168,7 @@ def build_omni_kwargs( "model": model, "flow_shift": flow_shift, "tensor_parallel_size": args.tensor_parallel_size, + "ulysses_degree": args.ulysses_degree, "enforce_eager": args.enforce_eager, "model_config": {"lingbot_action_root": str(paths.action_root)}, } diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py index f5567a706d1..3d883414e1e 100644 --- a/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py +++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.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 __future__ import annotations @@ -7,7 +7,7 @@ import math import sys from pathlib import Path -from types import ModuleType +from types import ModuleType, SimpleNamespace import pytest import torch @@ -33,6 +33,10 @@ "vllm_omni.diffusion", "vllm_omni.diffusion.attention", "vllm_omni.diffusion.attention.layer", + "vllm_omni.diffusion.distributed", + "vllm_omni.diffusion.distributed.comm", + "vllm_omni.diffusion.distributed.parallel_state", + "vllm_omni.diffusion.distributed.sp_plan", "vllm_omni.diffusion.layers", "vllm_omni.diffusion.layers.norm", "vllm_omni.diffusion.layers.rope", @@ -58,6 +62,47 @@ def _install_vllm_stubs() -> None: distributed.get_tensor_model_parallel_world_size = lambda: 1 distributed.tensor_model_parallel_all_reduce = lambda value: value + class _SeqAllToAll4D: + @staticmethod + def apply(group, value, scatter_idx, gather_idx, use_sync=False): + del group, scatter_idx, gather_idx, use_sync + return value + + setattr( + sys.modules["vllm_omni.diffusion.distributed.comm"], + "SeqAllToAll4D", + _SeqAllToAll4D, + ) + + def get_sp_group(): + return SimpleNamespace( + ulysses_world_size=1, + ulysses_rank=0, + ulysses_group=None, + ) + + setattr( + sys.modules["vllm_omni.diffusion.distributed.parallel_state"], + "get_sp_group", + get_sp_group, + ) + + class _SequenceParallelInput: + def __init__(self, split_dim, expected_dims=None, split_output=False, auto_pad=False): + self.split_dim = split_dim + self.expected_dims = expected_dims + self.split_output = split_output + self.auto_pad = auto_pad + + class _SequenceParallelOutput: + def __init__(self, gather_dim, expected_dims=None): + self.gather_dim = gather_dim + self.expected_dims = expected_dims + + sp_plan = sys.modules["vllm_omni.diffusion.distributed.sp_plan"] + setattr(sp_plan, "SequenceParallelInput", _SequenceParallelInput) + setattr(sp_plan, "SequenceParallelOutput", _SequenceParallelOutput) + def set_weight_attrs(weight: torch.Tensor, attrs: dict) -> None: for name, value in attrs.items(): setattr(weight, name, value) @@ -541,3 +586,38 @@ def test_tp_rmsnorm_weight_loader_selects_rank_shard(monkeypatch: pytest.MonkeyP norm.weight.weight_loader(norm.weight, torch.tensor([10.0, 20.0, 30.0, 40.0])) torch.testing.assert_close(norm.weight, torch.tensor([30.0, 40.0])) + + +def test_single_token_text_kv_shard_owns_compact_storage(monkeypatch): + module = _load_module() + monkeypatch.setattr( + module, + "get_sp_group", + lambda: SimpleNamespace( + ulysses_world_size=2, + ulysses_rank=1, + ulysses_group=None, + ), + ) + attention = module.LingBotCrossAttention(dim=8, num_heads=4) + full = torch.arange(8, dtype=torch.float32).reshape(1, 1, 4, 2) + shard = attention.shard_kv_heads(full) + torch.testing.assert_close(shard, full[:, :, 2:], rtol=0, atol=0) + assert shard.is_contiguous() + assert shard.untyped_storage().nbytes() == shard.numel() * shard.element_size() + + +@pytest.mark.parametrize("attention_class", ["LingBotSelfAttention", "LingBotCrossAttention"]) +def test_ulysses_rejects_non_divisible_head_count(monkeypatch, attention_class): + module = _load_module() + monkeypatch.setattr( + module, + "get_sp_group", + lambda: SimpleNamespace( + ulysses_world_size=3, + ulysses_rank=0, + ulysses_group=None, + ), + ) + with pytest.raises(ValueError, match="heads must be divisible"): + getattr(module, attention_class)(dim=8, num_heads=4) diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py index e2e2a08b6a6..b4bf7805fe3 100644 --- a/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py +++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.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 __future__ import annotations @@ -7,14 +7,17 @@ import inspect import json import math +from datetime import timedelta from pathlib import Path +from types import SimpleNamespace import pytest import torch from tests.diffusion.models.lingbot_world import test_lingbot_world_attention as attention_tests +from tests.helpers.mark import hardware_test -pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] _FIXTURE_PATH = Path(__file__).with_name("fixtures") / "lingbot_world_weight_index_names.json.fixture" _SHAPE_FIXTURE_PATH = Path(__file__).with_name("fixtures") / "lingbot_world_official_shapes.json.fixture" @@ -117,6 +120,7 @@ def _assert_self_cache_unchanged(cache, snapshot) -> None: attention_tests._assert_cache_unchanged(layer_cache, layer_snapshot) +@pytest.mark.cpu def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera_path() -> None: torch.manual_seed(7) module = attention_tests._load_module() @@ -191,6 +195,7 @@ def test_tiny_transformer_runs_four_chunks_with_explicit_cache_commit_and_camera assert not torch.equal(camera_output, alternate_output) +@pytest.mark.cpu def test_partial_cross_attention_cache_projects_text_only_for_missing_layers() -> None: module = attention_tests._load_module() model = _tiny_model(module, num_layers=2).eval() @@ -229,6 +234,7 @@ def test_partial_cross_attention_cache_projects_text_only_for_missing_layers() - @pytest.mark.parametrize("frames", [1, 6], ids=("partial", "multiple_blocks")) +@pytest.mark.cpu def test_forward_rejects_chunks_that_do_not_equal_configured_block_size(frames: int) -> None: module = attention_tests._load_module() model = _tiny_model( @@ -257,6 +263,7 @@ def test_forward_rejects_chunks_that_do_not_equal_configured_block_size(frames: assert cache.cross_attention == [None] +@pytest.mark.cpu def test_forward_accepts_exactly_one_configured_frame_block() -> None: module = attention_tests._load_module() model = _tiny_model( @@ -279,6 +286,7 @@ def test_forward_accepts_exactly_one_configured_frame_block() -> None: assert output.shape == (1, 2, 3, 4, 4) +@pytest.mark.cpu def test_transformer_allocates_request_cache_from_its_configured_geometry() -> None: module = attention_tests._load_module() model = _tiny_model( @@ -303,6 +311,7 @@ def test_transformer_allocates_request_cache_from_its_configured_geometry() -> N assert all(layer.value.shape == (2, 6 * 4 * 6, 2, 2) for layer in cache.self_attention) +@pytest.mark.cpu def test_video_patch_embedding_uses_temporal_height_width_token_order() -> None: module = attention_tests._load_module() model = _tiny_model(module, num_layers=1, num_frames_per_block=2) @@ -323,6 +332,7 @@ def test_video_patch_embedding_uses_temporal_height_width_token_order() -> None: torch.testing.assert_close(tokens[0, :, 0], torch.tensor([1.0, 3.0, 9.0, 11.0])) +@pytest.mark.cpu def test_unpatchify_restores_two_frame_channel_and_spatial_order() -> None: module = attention_tests._load_module() model = _tiny_model(module, num_layers=1, num_frames_per_block=2) @@ -363,6 +373,7 @@ def test_unpatchify_restores_two_frame_channel_and_spatial_order() -> None: torch.testing.assert_close(output, expected) +@pytest.mark.cpu def test_head_modulation_broadcasts_distinct_condition_per_frame() -> None: module = attention_tests._load_module() model = _tiny_model(module, num_layers=1, num_frames_per_block=2) @@ -381,6 +392,7 @@ def test_head_modulation_broadcasts_distinct_condition_per_frame() -> None: torch.testing.assert_close(output, expected) +@pytest.mark.cpu def test_constructor_defaults_match_official_checkpoint_config() -> None: module = attention_tests._load_module() parameters = inspect.signature(module.CausalLingBotWorldTransformer3DModel.__init__).parameters @@ -406,6 +418,7 @@ def test_constructor_defaults_match_official_checkpoint_config() -> None: assert {name: parameters[name].default for name in expected} == expected +@pytest.mark.cpu def test_transformer_exposes_parameter_dtype_for_pipeline_runtime() -> None: module = attention_tests._load_module() model = _tiny_model(module) @@ -413,12 +426,14 @@ def test_transformer_exposes_parameter_dtype_for_pipeline_runtime() -> None: assert model.dtype == next(model.parameters()).dtype +@pytest.mark.cpu def test_transformer_declares_regional_compile_block() -> None: module = attention_tests._load_module() assert module.CausalLingBotWorldTransformer3DModel._repeated_blocks == ["LingBotAttentionBlock"] +@pytest.mark.cpu def test_lingbot_rms_norm_uses_global_tp_square_mean(monkeypatch) -> None: module = attention_tests._load_module() reduced_values: list[torch.Tensor] = [] @@ -440,6 +455,7 @@ def all_reduce(value: torch.Tensor) -> torch.Tensor: torch.testing.assert_close(reduced_values[0], torch.tensor([[25.0]])) +@pytest.mark.cpu def test_attention_rejects_heads_not_divisible_by_tp_size(monkeypatch) -> None: module = attention_tests._load_module() monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: 3) @@ -448,6 +464,7 @@ def test_attention_rejects_heads_not_divisible_by_tp_size(monkeypatch) -> None: module.LingBotSelfAttention(dim=4, num_heads=2) +@pytest.mark.cpu def test_constructor_rejects_unsupported_qk_norm_with_config_error() -> None: module = attention_tests._load_module() @@ -456,6 +473,7 @@ def test_constructor_rejects_unsupported_qk_norm_with_config_error() -> None: @pytest.mark.parametrize("field", ["image_dim", "added_kv_proj_dim", "pos_embed_seq_len"]) +@pytest.mark.cpu def test_constructor_rejects_non_null_image_embedding_fields(field: str) -> None: module = attention_tests._load_module() @@ -463,6 +481,7 @@ def test_constructor_rejects_non_null_image_embedding_fields(field: str) -> None module.CausalLingBotWorldTransformer3DModel(**{field: 4}) +@pytest.mark.cpu def test_constructor_rejects_unsupported_quantization_with_runtime_error() -> None: module = attention_tests._load_module() @@ -470,6 +489,7 @@ def test_constructor_rejects_unsupported_quantization_with_runtime_error() -> No module.CausalLingBotWorldTransformer3DModel(quant_config=object()) +@pytest.mark.cpu def test_from_config_accepts_diffusers_metadata_and_normalizes_patch_size() -> None: module = attention_tests._load_module() @@ -522,6 +542,7 @@ def __init__(self, **kwargs) -> None: ("sink_size", 0), ], ) +@pytest.mark.cpu def test_from_config_rejects_checkpoint_topology_drift(field: str, value: object) -> None: module = attention_tests._load_module() config = { @@ -554,6 +575,7 @@ def test_from_config_rejects_checkpoint_topology_drift(field: str, value: object module.CausalLingBotWorldTransformer3DModel.from_config(config) +@pytest.mark.cpu def test_from_config_ignores_non_semantic_checkpoint_metadata() -> None: module = attention_tests._load_module() config = { @@ -569,6 +591,7 @@ def test_from_config_ignores_non_semantic_checkpoint_metadata() -> None: assert model.config.num_layers == 40 +@pytest.mark.cpu def test_load_weights_uses_parameter_loaders_and_rejects_unknown_model_keys() -> None: module = attention_tests._load_module() model = _tiny_model(module, num_layers=1) @@ -606,6 +629,7 @@ def record_loader(param: torch.Tensor, loaded_weight: torch.Tensor, shard_id: st model.load_weights([("unexpected_model.weight", torch.ones(1))]) +@pytest.mark.cpu def test_load_weights_consumes_checkpoint_iterator_incrementally() -> None: module = attention_tests._load_module() model = _tiny_model(module, num_layers=1) @@ -619,6 +643,7 @@ def weights(): assert model.load_weights(weights()) == {first_name} +@pytest.mark.cpu def test_real_auto_weights_loader_covers_fused_model_parameter_namespace() -> None: utils = pytest.importorskip("vllm.model_executor.models.utils") module = attention_tests._load_module() @@ -638,6 +663,7 @@ def __init__(self) -> None: assert set(dict(pipeline.named_parameters())) <= loaded +@pytest.mark.cpu def test_checkpoint_weight_index_fixture_matches_model_namespaces() -> None: module = attention_tests._load_module() fixture = json.loads(_FIXTURE_PATH.read_text()) @@ -667,6 +693,7 @@ def test_checkpoint_weight_index_fixture_matches_model_namespaces() -> None: assert all(len(shard["parameter_names_sha256"]) == 64 for shard in fixture["shards"].values()) +@pytest.mark.cpu def test_official_default_shapes_match_public_safetensors_header_fixture() -> None: module = attention_tests._load_module() fixture = json.loads(_SHAPE_FIXTURE_PATH.read_text()) @@ -683,3 +710,284 @@ def test_official_default_shapes_match_public_safetensors_header_fixture() -> No shape, dtype = parameter_specs[name] assert tuple(shape) == tuple(expected_shape), name assert dtype == expected_dtype + + +@pytest.mark.cpu +def test_timestep_expansion_only_runs_with_ulysses(monkeypatch): + dtype = torch.bfloat16 + module = attention_tests._load_module() + hidden = torch.randn(2, 12, 4).to(dtype) + camera = torch.randn_like(hidden) + table = torch.randn(2, 3, 6, 4).to(dtype) + rotary = (torch.randn(12, 2), torch.randn(12, 2)) + prepare = module._LingBotSPPrepare() + assert prepare(hidden, camera, table, rotary)[2] is table + monkeypatch.setattr(module, "get_sp_group", lambda: SimpleNamespace(ulysses_world_size=2)) + expanded = prepare(hidden, camera, table, rotary)[2] + assert expanded.shape == (2, 12, 6, 4) + torch.testing.assert_close(expanded, table.repeat_interleave(4, dim=1), rtol=0, atol=0) + + +@pytest.mark.parametrize("unsharded", [None, 3], ids=["missing-hook", "unsharded-rope"]) +@pytest.mark.cpu +def test_missing_or_partial_sp_split_fails_before_attention(monkeypatch, unsharded): + module = attention_tests._load_module() + model = _tiny_model(module, num_frames_per_block=3, sliding_window_num_frames=6) + cache = _cache(module, model) + for block in model.blocks: + block.self_attn.ulysses_world_size = 2 + monkeypatch.setattr(module, "get_sp_group", lambda: SimpleNamespace(ulysses_world_size=2)) + original = model.sp_prepare.forward + + def partial_split(*args): + values = original(*args) + if unsharded is None: + return values + return tuple( + value if i == unsharded else value.chunk(2, dim=1 if i < 3 else 0)[0] for i, value in enumerate(values) + ) + + monkeypatch.setattr(model.sp_prepare, "forward", partial_split) + + def forbidden(*args, **kwargs): + raise AssertionError("attention/collectives must not run before shard validation") + + monkeypatch.setattr(model.blocks[0], "forward", forbidden) + with pytest.raises(RuntimeError, match="SP input hooks"): + model( + torch.randn(1, 36, 3, 4, 4), + torch.tensor([1.0]), + torch.randn(1, 3, 6), + torch.randn(1, 384, 3, 4, 4), + cache=cache, + start_frame=0, + update_cache=False, + ) + + +# Real multi-rank regression: synthetic weights, native SP/TP and attention. +_HEADS, _HEAD_DIM, _LAYERS = 8, 32, 2 +_FRAMES, _SIDE, _TOKENS_PER_FRAME = 3, 32, 256 + + +def _model(dtype): + from vllm_omni.diffusion.models.lingbot_world.transformer import CausalLingBotWorldTransformer3DModel + + model = CausalLingBotWorldTransformer3DModel( + num_attention_heads=_HEADS, + attention_head_dim=_HEAD_DIM, + num_layers=_LAYERS, + in_channels=4, + out_channels=4, + text_dim=8, + freq_dim=16, + ffn_dim=512, + patch_size=(1, 2, 2), + sink_size=1, + num_frames_per_block=_FRAMES, + sliding_window_num_frames=6, + rope_max_seq_len=32, + ).eval() + return model.to(device=torch.device("cuda", torch.accelerator.current_device_index()), dtype=dtype) + + +def _checkpoint(model): + weights = {} + generator = torch.Generator().manual_seed(17) + for name, param in sorted(model.named_parameters()): + data = torch.randn(param.shape, generator=generator) * 0.02 + if "norm" in name and name.endswith("weight"): + data.fill_(1) + if name.endswith("modulation") and data.shape[-2] == 6: + data[..., 2, :] = 1 + data[..., 5, :] = 1 + data = data.to(param.dtype) + param.copy_(data) + if ".self_attn.qkv." in name: + for component, chunk in zip(("q", "k", "v"), data.chunk(3), strict=True): + weights[name.replace(".self_attn.qkv.", f".self_attn.{component}.")] = chunk + else: + weights[name] = data + return weights + + +def _rollout(model, mode, dtype, batch): + from vllm_omni.diffusion.models.lingbot_world.transformer import LingBotTransformerCache + from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec + from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVCache, ARDiffusionKVConfig + from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState + + device = next(model.parameters()).device + state = None + if mode == "paged": + kv = ARDiffusionKVCache( + ARDiffusionKVConfig(enable=True, chunk_size=_TOKENS_PER_FRAME, window_chunks=5, sink_chunks=1), + num_layers=_LAYERS, + num_kv_heads=model.blocks[0].self_attn.num_sp_heads, + head_size=_HEAD_DIM, + dtype=dtype, + block_size=_TOKENS_PER_FRAME, + max_model_len=4096, + available_bytes=1 << 27, + kv_branches=(ARDiffusionKVBranchSpec("main", 0),), + session_capacity=1, + frames_per_block=_FRAMES, + max_scratch_tokens_per_branch=_FRAMES * _TOKENS_PER_FRAME, + cross_attention_lengths={"text": 5}, + device=device, + ) + state = ARDiffusionKVState(kv, "numeric", {"main": kv.begin_request("numeric")}, num_layers=_LAYERS) + cache = LingBotTransformerCache(self_attention=[], cross_attention=[None] * _LAYERS) + else: + cache = model.allocate_cache( + batch_size=batch, latent_height=_SIDE, latent_width=_SIDE, device=device, dtype=dtype + ) + generator = torch.Generator().manual_seed(123) + text = torch.randn(batch, 5, 8, generator=generator).to(device, dtype) + outputs = [] + try: + for start in range(0, 12, _FRAMES): + latent = torch.randn(batch, 4, _FRAMES, _SIDE, _SIDE, generator=generator).to(device, dtype) + camera = torch.randn(batch, 384, _FRAMES, _SIDE, _SIDE, generator=generator).to(device, dtype) + for commit in (False, False, True): + if state is not None: + cache.self_attention = state.get_kv_caches( + "main", seq_len=_FRAMES * _TOKENS_PER_FRAME, commit_current=commit + ) + output = model( + latent, + torch.tensor([100.0, 400.0, 900.0], device=device).expand(batch, -1), + text, + camera, + cache=cache, + start_frame=start, + update_cache=commit, + ) + outputs.append(output.float().cpu()) + if state is not None: + if not state.is_cross_attention_populated("main", "text"): + state.populate_cross_attention( + "main", "text", [(layer.key, layer.value) for layer in cache.cross_attention] + ) + for layer, pool in zip( + cache.cross_attention, state.get_cross_attention_kv("main", "text"), strict=True + ): + layer.key, layer.value = pool["k"], pool["v"] + state.commit_paged_context("main") + cross = [ + (layer.key.detach().float().cpu(), layer.value.detach().float().cpu()) for layer in cache.cross_attention + ] + return torch.stack(outputs), cross + finally: + if state is not None: + state.close() + + +def _worker(rank, world_size, sp_size, tp_size, mode, dtype, batch, rendezvous): + from vllm.config import VllmConfig + from vllm.config.vllm import set_current_vllm_config + from vllm.distributed import get_tensor_model_parallel_rank + + from vllm_omni.diffusion.config import set_current_diffusion_config + from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, DiffusionParallelConfig, OmniDiffusionConfig + from vllm_omni.diffusion.distributed.parallel_state import ( + destroy_distributed_env, + destroy_model_parallel, + get_sp_group, + init_distributed_environment, + initialize_model_parallel, + ) + from vllm_omni.diffusion.forward_context import set_forward_context + from vllm_omni.diffusion.registry import _apply_sequence_parallel_if_enabled + from vllm_omni.platforms import current_omni_platform + + current_omni_platform.set_device(torch.device("cuda", rank)) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.distributed.init_process_group( + "nccl", init_method=f"file://{rendezvous}", world_size=world_size, rank=rank, timeout=timedelta(seconds=90) + ) + init_distributed_environment(world_size=world_size, rank=rank, local_rank=rank, backend="nccl") + try: + with torch.inference_mode(), set_current_vllm_config(VllmConfig()): + for baseline in (True, False): + sp, tp = (1, 1) if baseline else (sp_size, tp_size) + initialize_model_parallel( + data_parallel_size=world_size // (sp * tp), + sequence_parallel_size=sp, + ulysses_degree=sp, + tensor_parallel_size=tp, + ) + config = OmniDiffusionConfig( + model=str(Path(rendezvous).parent), + dtype=dtype, + enforce_eager=True, + parallel_config=DiffusionParallelConfig( + data_parallel_size=world_size // (sp * tp), + sequence_parallel_size=sp, + ulysses_degree=sp, + tensor_parallel_size=tp, + ), + diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="TORCH_SDPA")), + ) + with set_current_diffusion_config(config), set_forward_context(omni_diffusion_config=config): + model = _model(dtype) + if baseline: + weights = _checkpoint(model) + else: + model.load_weights(iter(weights.items())) + _apply_sequence_parallel_if_enabled(SimpleNamespace(transformer=model), config) + result, cross = _rollout(model, mode, dtype, batch) + if baseline: + expected, expected_cross = result, cross + else: + assert expected.abs().max() > 1e-3 + bound = 1e-5 if dtype == torch.float32 else 1e-2 + error = (result - expected).double().norm() / expected.double().norm() + assert error <= bound, f"rank={rank}, TP{tp} SP{sp}: relative L2 {error.item():.3g} > {bound}" + local_heads = _HEADS // tp // sp + first = ( + get_tensor_model_parallel_rank() * (_HEADS // tp) + + get_sp_group().ulysses_rank * local_heads + ) + for (k, v), (ek, ev) in zip(cross, expected_cross, strict=True): + tolerance = 1e-5 if dtype == torch.float32 else 2e-2 + torch.testing.assert_close( + k, ek[:, :, first : first + local_heads], rtol=tolerance, atol=tolerance + ) + torch.testing.assert_close( + v, ev[:, :, first : first + local_heads], rtol=tolerance, atol=tolerance + ) + if rank == 0: + print(f"{mode} {dtype} TP{tp} SP{sp}: relative L2={error.item():.3g}", flush=True) + del model + torch.distributed.barrier() + destroy_model_parallel() + finally: + destroy_distributed_env() + + +def _run(tmp_path: Path, sp, tp, mode, dtype, batch): + if torch.accelerator.device_count() < sp * tp: + pytest.skip(f"requires {sp * tp} GPUs") + torch.multiprocessing.spawn( + _worker, args=(sp * tp, sp, tp, mode, dtype, batch, str(tmp_path / "rendezvous")), nprocs=sp * tp + ) + + +@pytest.mark.parallel +@hardware_test(res={"cuda": "L4"}, num_cards=2) +def test_sp2_direct_matches_sp1_fp32(tmp_path): + _run(tmp_path, 2, 1, "direct", torch.float32, 2) + + +@pytest.mark.parallel +@hardware_test(res={"cuda": "L4"}, num_cards=4) +def test_sp4_direct_matches_sp1_bf16(tmp_path): + _run(tmp_path, 4, 1, "direct", torch.bfloat16, 1) + + +@pytest.mark.parallel +@hardware_test(res={"cuda": "L4"}, num_cards=4) +def test_tp2_sp2_paged_matches_sp1_bf16(tmp_path): + _run(tmp_path, 2, 2, "paged", torch.bfloat16, 1) diff --git a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py index 52113208a74..a58eee935db 100644 --- a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py +++ b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py @@ -126,8 +126,6 @@ def __init__(self, *, raise_on_call: int | None = None, dtype: torch.dtype = tor sink_size=3, ) self.blocks = nn.ModuleList([nn.Identity(), nn.Identity()]) - for block in self.blocks: - block.self_attn = SimpleNamespace(num_local_heads=2, head_dim=4) self.calls: list[dict] = [] self.cache_allocations: list[dict] = [] self.raise_on_call = raise_on_call @@ -387,6 +385,9 @@ def _od_config(**overrides): "parallel_config": SimpleNamespace( pipeline_parallel_size=1, sequence_parallel_size=1, + ulysses_degree=1, + ring_degree=1, + allgather_degree=1, cfg_parallel_size=1, vae_patch_parallel_size=1, use_hsdp=False, @@ -438,15 +439,15 @@ def test_pipeline_respects_loader_managed_component_placement(offload_field: str assert getattr(pipeline.vae, "to_calls", []) == [] -def test_ar_diffusion_capability_uses_fixed_tp_local_lingbot_geometry() -> None: +def test_ar_diffusion_capability_uses_transformer_local_head_geometry() -> None: module = _load_pipeline_module() - module.get_tensor_model_parallel_world_size = lambda: 1 pipeline = _pipeline(module) - + # Use the constructed head count even when the config describes a different geometry. + pipeline.transformer.blocks[0].self_attn = SimpleNamespace(num_sp_heads=1) spec = pipeline.ar_diffusion_kv_cache_spec() assert spec.num_layers == 2 - assert spec.num_kv_heads == 2 + assert spec.num_kv_heads == 1 assert spec.head_size == 4 assert spec.tokens_per_frame == 1 assert spec.frames_per_block == 3 @@ -604,7 +605,6 @@ def test_component_discovery_uses_official_checkpoint_contract() -> None: ("field", "value", "feature"), [ ("pipeline_parallel_size", 2, "pipeline parallelism"), - ("sequence_parallel_size", 2, "sequence parallelism"), ("cfg_parallel_size", 2, "CFG parallelism"), ("vae_patch_parallel_size", 2, "VAE parallelism"), ("use_hsdp", True, "HSDP"), @@ -622,6 +622,43 @@ def test_unsupported_parallel_modes_fail_before_component_loading(field: str, va assert module._loader_state.prefetch_calls == [] +def test_pure_ulysses_parallel_config_is_supported() -> None: + module = _load_pipeline_module() + parallel_config = _od_config().parallel_config + parallel_config.sequence_parallel_size = 2 + parallel_config.ulysses_degree = 2 + + pipeline = module.LingBotWorldCausalDMDPipeline(od_config=_od_config(parallel_config=parallel_config)) + + assert pipeline.transformer is not None + + +@pytest.mark.parametrize( + "overrides", + [ + {"sequence_parallel_size": 4, "ring_degree": 2}, # Normalized hybrid. + {"ulysses_degree": 1, "allgather_degree": 2}, # Normalized AllGather-KV. + {"ring_degree": 2}, # Isolate each clause from the SP-size mismatch. + {"allgather_degree": 2}, + {"ulysses_mode": "advanced_uaa"}, + {"ulysses_a2a_permute": True}, + {"ulysses_degree": None}, # Missing degree must not imply pure Ulysses. + ], +) +def test_unsupported_sp_config_fails_before_component_loading(overrides): + module = _load_pipeline_module() + config = _od_config().parallel_config + config.sequence_parallel_size = config.ulysses_degree = 2 + for name, value in overrides.items(): + if value is None: + delattr(config, name) + else: + setattr(config, name, value) + with pytest.raises(NotImplementedError, match="pure Ulysses"): + module.LingBotWorldCausalDMDPipeline(od_config=_od_config(parallel_config=config)) + assert module._loader_state.prefetch_calls == [] + + def test_unsupported_quantization_fails_before_component_loading() -> None: module = _load_pipeline_module() diff --git a/tests/examples/offline_inference/test_lingbot_world_v2.py b/tests/examples/offline_inference/test_lingbot_world_v2.py new file mode 100644 index 00000000000..41d2e73b712 --- /dev/null +++ b/tests/examples/offline_inference/test_lingbot_world_v2.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from pathlib import Path + +import pytest + +from examples.offline_inference.diffusion import lingbot_world_v2 as example + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_offline_ulysses_argument(tmp_path): + argv = ["--image", "frame.jpg", "--action-dir", "forward", "--prompt", "A lake"] + paths = example.LingBotPaths( + tmp_path / "frame.jpg", tmp_path / "forward", tmp_path, Path("forward"), 9, tmp_path / "out.mp4" + ) + assert example.parse_args(argv).ulysses_degree == 1 + args = example.parse_args([*argv, "--ulysses-degree", "4"]) + kwargs = example.build_omni_kwargs(args, paths) + assert kwargs["ulysses_degree"] == 4 and kwargs["tensor_parallel_size"] == 1 + args.ulysses_degree = 0 + with pytest.raises(ValueError, match="--ulysses-degree"): + example.build_omni_kwargs(args, paths) diff --git a/vllm_omni/diffusion/models/lingbot_world/pipeline.py b/vllm_omni/diffusion/models/lingbot_world/pipeline.py index b41abd861d6..d2dfde3d4e9 100644 --- a/vllm_omni/diffusion/models/lingbot_world/pipeline.py +++ b/vllm_omni/diffusion/models/lingbot_world/pipeline.py @@ -18,7 +18,6 @@ from diffusers.utils.torch_utils import randn_tensor from torch import nn from transformers import AutoTokenizer, UMT5EncoderModel -from vllm.distributed import get_tensor_model_parallel_world_size from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig @@ -203,7 +202,6 @@ def _validate_parallel_config(od_config: OmniDiffusionConfig) -> None: return unsupported_sizes = { "pipeline_parallel_size": "pipeline parallelism", - "sequence_parallel_size": "sequence parallelism", "cfg_parallel_size": "CFG parallelism", "vae_patch_parallel_size": "VAE parallelism", } @@ -211,6 +209,25 @@ def _validate_parallel_config(od_config: OmniDiffusionConfig) -> None: size = getattr(parallel_config, field, 1) or 1 if size > 1: raise NotImplementedError(f"LingBot World v1 does not support {feature} ({field}={size}).") + sequence_parallel_size = getattr(parallel_config, "sequence_parallel_size", 1) or 1 + ulysses_degree = getattr(parallel_config, "ulysses_degree", 1) or 1 + ring_degree = getattr(parallel_config, "ring_degree", 1) or 1 + allgather_degree = getattr(parallel_config, "allgather_degree", 1) or 1 + ulysses_mode = getattr(parallel_config, "ulysses_mode", "strict") + ulysses_a2a_permute = bool(getattr(parallel_config, "ulysses_a2a_permute", False)) + if ( + sequence_parallel_size != ulysses_degree + or ring_degree != 1 + or allgather_degree != 1 + or ulysses_mode != "strict" + or ulysses_a2a_permute + ): + raise NotImplementedError( + "LingBot World sequence parallelism requires pure Ulysses with ulysses_a2a_permute disabled: " + f"sequence_parallel_size={sequence_parallel_size}, ulysses_degree={ulysses_degree}, " + f"ring_degree={ring_degree}, allgather_degree={allgather_degree}, ulysses_mode={ulysses_mode!r}, " + f"ulysses_a2a_permute={ulysses_a2a_permute}." + ) if getattr(parallel_config, "use_hsdp", False): raise NotImplementedError("LingBot World v1 does not support HSDP.") if getattr(parallel_config, "enable_expert_parallel", False): @@ -557,8 +574,6 @@ def ar_diffusion_kv_cache_spec(self) -> ARDiffusionKVCacheSpec: latent_height = self._ar_height // spatial latent_width = self._ar_width // spatial tokens_per_frame = (latent_height // patch_height) * (latent_width // patch_width) - tp_size = get_tensor_model_parallel_world_size() - num_local_heads = int(self.transformer.config.num_attention_heads) // tp_size total_window_frames = ( int(self.transformer.config.local_attn_size) if int(self.transformer.config.local_attn_size) != -1 @@ -579,7 +594,7 @@ def ar_diffusion_kv_cache_spec(self) -> ARDiffusionKVCacheSpec: ) return ARDiffusionKVCacheSpec( num_layers=int(self.transformer.config.num_layers), - num_kv_heads=num_local_heads, + num_kv_heads=int(self.transformer.blocks[0].self_attn.num_sp_heads), head_size=int(self.transformer.config.attention_head_dim), tokens_per_frame=tokens_per_frame, frames_per_block=int(self.transformer.config.num_frames_per_block), @@ -960,20 +975,11 @@ def _ar_text_caches( def layer_kv() -> Iterator[tuple[torch.Tensor, torch.Tensor]]: for block in self.transformer.blocks: cross_attention = block.cross_attn - key = cross_attention.norm_k(cross_attention.k(projected_text)).unflatten( - 2, - ( - cross_attention.num_local_heads, - cross_attention.head_dim, - ), - ) - value = cross_attention.v(projected_text).unflatten( - 2, - ( - cross_attention.num_local_heads, - cross_attention.head_dim, - ), + shape = (cross_attention.num_local_heads, cross_attention.head_dim) + key = cross_attention.shard_kv_heads( + cross_attention.norm_k(cross_attention.k(projected_text)).unflatten(2, shape) ) + value = cross_attention.shard_kv_heads(cross_attention.v(projected_text).unflatten(2, shape)) yield key, value state.populate_cross_attention( diff --git a/vllm_omni/diffusion/models/lingbot_world/transformer.py b/vllm_omni/diffusion/models/lingbot_world/transformer.py index 35ec1880112..32931db0c9e 100644 --- a/vllm_omni/diffusion/models/lingbot_world/transformer.py +++ b/vllm_omni/diffusion/models/lingbot_world/transformer.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 """Checkpoint-compatible causal DiT for the LingBot World v2 model package.""" from __future__ import annotations @@ -24,6 +24,9 @@ from vllm.model_executor.utils import set_weight_attrs from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D +from vllm_omni.diffusion.distributed.parallel_state import get_sp_group +from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput from vllm_omni.diffusion.layers.norm import LayerNorm from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import ( @@ -127,6 +130,37 @@ def _projection_prefix(prefix: str, name: str) -> str: return f"{prefix}.{name}" if prefix else name +def _ulysses_state() -> tuple[int, int, torch.distributed.ProcessGroup | None]: + coordinator = get_sp_group() + return ( + int(coordinator.ulysses_world_size), + int(coordinator.ulysses_rank), + coordinator.ulysses_group, + ) + + +class _LingBotSPPrepare(nn.Module): + """Expand frame conditioning only when Ulysses shards the token sequence.""" + + def forward( + self, + hidden_states: torch.Tensor, + camera_hidden_states: torch.Tensor, + timestep_projection: torch.Tensor, + rotary_emb: tuple[torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + num_frames = timestep_projection.shape[1] + if hidden_states.shape[1] % num_frames: + raise ValueError("LingBot token count must be divisible by the latent frame count.") + tokens_per_frame = hidden_states.shape[1] // num_frames + if get_sp_group().ulysses_world_size > 1: + timestep_projection = ( + timestep_projection.unsqueeze(2).expand(-1, -1, tokens_per_frame, -1, -1).flatten(1, 2) + ) + cosine, sine = rotary_emb + return hidden_states, camera_hidden_states, timestep_projection, cosine, sine + + class LingBotSelfAttention(nn.Module): """Block-causal self-attention over retained history and one full chunk.""" @@ -157,6 +191,13 @@ def __init__( prefix=_projection_prefix(prefix, "qkv"), ) self.num_local_heads = self.qkv.num_heads + self.ulysses_world_size, self.ulysses_rank, self.ulysses_group = _ulysses_state() + if self.num_local_heads % self.ulysses_world_size: + raise ValueError( + "LingBot local attention heads must be divisible by the Ulysses degree: " + f"heads={self.num_local_heads}, ulysses={self.ulysses_world_size}." + ) + self.num_sp_heads = self.num_local_heads // self.ulysses_world_size self.tp_inner_dim = self.num_local_heads * self.head_dim self.o = RowParallelLinear( dim, @@ -170,9 +211,9 @@ def __init__( self.norm_k = _LingBotRMSNorm(self.tp_inner_dim, eps) self.rotary_embedding = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True) self.attn = Attention( - num_heads=self.num_local_heads, + num_heads=self.num_sp_heads, head_size=self.head_dim, - num_kv_heads=self.num_local_heads, + num_kv_heads=self.num_sp_heads, softmax_scale=self.head_dim**-0.5, causal=False, role="self", @@ -287,6 +328,11 @@ def forward( query = self.rotary_embedding(query, cos, sin) key = self.rotary_embedding(key, cos, sin) + if self.ulysses_world_size > 1: + query = SeqAllToAll4D.apply(self.ulysses_group, query, 2, 1, False) + key = SeqAllToAll4D.apply(self.ulysses_group, key, 2, 1, False) + value = SeqAllToAll4D.apply(self.ulysses_group, value, 2, 1, False) + if isinstance(cache, ARDiffusionPagedLayerInputs): if query.shape[0] != 1: raise RuntimeError("LingBot AR-Diffusion paged attention requires batch_size=1.") @@ -347,6 +393,8 @@ def forward( ) else: output = self.attn(query, visible_key, visible_value) + if self.ulysses_world_size > 1: + output = SeqAllToAll4D.apply(self.ulysses_group, output, 1, 2, False) return self.o(output.flatten(2, 3)) @@ -373,6 +421,13 @@ def __init__( self.num_heads = num_heads self.head_dim = dim // num_heads self.num_local_heads = num_heads // tp_size + self.ulysses_world_size, self.ulysses_rank, self.ulysses_group = _ulysses_state() + if self.num_local_heads % self.ulysses_world_size: + raise ValueError( + "LingBot local cross-attention heads must be divisible by the Ulysses degree: " + f"heads={self.num_local_heads}, ulysses={self.ulysses_world_size}." + ) + self.num_sp_heads = self.num_local_heads // self.ulysses_world_size self.tp_inner_dim = self.num_local_heads * self.head_dim self.q = ColumnParallelLinear( @@ -410,9 +465,9 @@ def __init__( self.norm_q = _LingBotRMSNorm(self.tp_inner_dim, eps) self.norm_k = _LingBotRMSNorm(self.tp_inner_dim, eps) self.attn = Attention( - num_heads=self.num_local_heads, + num_heads=self.num_sp_heads, head_size=self.head_dim, - num_kv_heads=self.num_local_heads, + num_kv_heads=self.num_sp_heads, softmax_scale=self.head_dim**-0.5, causal=False, role="cross", @@ -422,6 +477,14 @@ def __init__( disable_kv_quant=True, ) + def shard_kv_heads(self, value: torch.Tensor) -> torch.Tensor: + """Select this Ulysses rank's projected K/V heads.""" + + if self.ulysses_world_size == 1: + return value + start = self.ulysses_rank * self.num_sp_heads + return value[:, :, start : start + self.num_sp_heads].clone(memory_format=torch.contiguous_format) + def forward( self, hidden_states: torch.Tensor, @@ -431,6 +494,8 @@ def forward( ) -> tuple[torch.Tensor, LingBotAttentionCache]: query = self.norm_q(self.q(hidden_states)) query = query.unflatten(2, (self.num_local_heads, self.head_dim)) + if self.ulysses_world_size > 1: + query = SeqAllToAll4D.apply(self.ulysses_group, query, 2, 1, False) # Text K/V is constant within a request and is projected once per layer. if cache is None: @@ -438,8 +503,8 @@ def forward( raise ValueError("encoder_hidden_states are required when the cross-attention cache is empty.") key = self.norm_k(self.k(encoder_hidden_states)) value = self.v(encoder_hidden_states) - key = key.unflatten(2, (self.num_local_heads, self.head_dim)) - value = value.unflatten(2, (self.num_local_heads, self.head_dim)) + key = self.shard_kv_heads(key.unflatten(2, (self.num_local_heads, self.head_dim))) + value = self.shard_kv_heads(value.unflatten(2, (self.num_local_heads, self.head_dim))) cache = LingBotAttentionCache( key=key, value=value, @@ -452,6 +517,8 @@ def forward( value = cache.value[:, : cache.end] output = self.attn(query, key, value) + if self.ulysses_world_size > 1: + output = SeqAllToAll4D.apply(self.ulysses_group, output, 1, 2, False) return self.o(output.flatten(2, 3)), cache @@ -540,7 +607,18 @@ def forward( rotary_emb: tuple[torch.Tensor, torch.Tensor], ) -> tuple[torch.Tensor, LingBotAttentionCache]: batch_size, token_count, dim = hidden_states.shape + if ( + timestep_projection.ndim != 4 + or timestep_projection.shape[0] != batch_size + or timestep_projection.shape[2:] != (6, dim) + ): + raise ValueError("LingBot timestep projection must have shape (batch, frames or tokens, 6, dim).") num_frames = timestep_projection.shape[1] + if num_frames == 0 or token_count % num_frames: + raise ValueError("LingBot timestep projection must align with hidden tokens.") + if self.self_attn.ulysses_world_size > 1 and num_frames != token_count: + raise ValueError("LingBot Ulysses requires one timestep projection per local token.") + # Without SP, broadcast per frame. With SP, each entry represents one token. tokens_per_frame = token_count // num_frames # Timestep, camera, and text remain separate conditioning paths. modulation = self.modulation.unsqueeze(1) + timestep_projection.float() @@ -635,13 +713,26 @@ def forward( self, hidden_states: torch.Tensor, timestep_embedding: torch.Tensor, + *, + tokens_per_frame: int | None = None, + token_offset: int = 0, ) -> torch.Tensor: num_frames = timestep_embedding.shape[1] - tokens_per_frame = hidden_states.shape[1] // num_frames modulation = self.modulation.unsqueeze(1) + timestep_embedding.unsqueeze(2).float() shift, scale = modulation.chunk(2, dim=2) - normalized = self.norm(hidden_states.float()).unflatten(1, (num_frames, tokens_per_frame)) - normalized = (normalized * (1 + scale) + shift).flatten(1, 2).to(hidden_states.dtype) + normalized = self.norm(hidden_states.float()) + if tokens_per_frame is None: + normalized = normalized.unflatten(1, (num_frames, -1)) + normalized = (normalized * (1 + scale) + shift).flatten(1, 2) + else: + # A sequence shard can begin/end inside a frame. + frames = ( + torch.arange(hidden_states.shape[1], device=hidden_states.device) + token_offset + ) // tokens_per_frame + scale = scale.squeeze(2).index_select(1, frames) + shift = shift.squeeze(2).index_select(1, frames) + normalized = normalized * (1 + scale) + shift + normalized = normalized.to(hidden_states.dtype) return self.head(normalized) @@ -678,6 +769,16 @@ class CausalLingBotWorldTransformer3DModel(nn.Module): _repeated_blocks = ["LingBotAttentionBlock"] packed_modules_mapping = {"qkv": ["q", "k", "v"]} _layerwise_offload_blocks_attrs = ["blocks"] + _sp_plan = { + "sp_prepare": { + 0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True), + 1: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True), + 2: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), + 3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True), + 4: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True), + }, + "sp_output_gather": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } def __init__( self, @@ -768,6 +869,8 @@ def __init__( stride=patch_size, ) self.patch_embedding_wancamctrl = _LingBotCameraPatchEmbedding(6 * 8 * 8, dim, patch_size) + self.sp_prepare = _LingBotSPPrepare() + self.sp_output_gather = nn.Identity() self.c2ws_hidden_states_layer1 = ColumnParallelLinear( dim, dim, @@ -904,8 +1007,7 @@ def allocate_cache( self.config.local_attn_size if self.config.local_attn_size != -1 else self.config.sliding_window_num_frames ) max_tokens = int(window_frames * post_patch_height * post_patch_width) - tp_size = get_tensor_model_parallel_world_size() - num_local_heads = self.config.num_attention_heads // tp_size + num_local_heads = self.blocks[0].self_attn.num_sp_heads return allocate_lingbot_cache( batch_size=batch_size, num_layers=self.config.num_layers, @@ -1138,6 +1240,33 @@ def forward( query_len=hidden_states.shape[1], ) cache.self_attention = [layer_context.to_layer_inputs() for layer_context in cache.self_attention] + sp_size = self.blocks[0].self_attn.ulysses_world_size + global_tokens = patched_frames * tokens_per_frame + hidden_dim = hidden_states.shape[-1] + rope_dim = rotary_emb[0].shape[-1] + if global_tokens % sp_size: + raise ValueError("LingBot Ulysses requires the token count to be divisible by its degree.") + hidden_states, camera_hidden_states, timestep_projection, cosine, sine = self.sp_prepare( + hidden_states, + camera_hidden_states, + timestep_projection, + rotary_emb, + ) + if sp_size > 1: + local_tokens = global_tokens // sp_size + expected_hidden = (batch_size, local_tokens, hidden_dim) + if ( + hidden_states.shape != expected_hidden + or camera_hidden_states.shape != expected_hidden + or timestep_projection.shape != (batch_size, local_tokens, 6, hidden_dim) + or cosine.shape != (local_tokens, rope_dim) + or sine.shape != (local_tokens, rope_dim) + ): + raise RuntimeError( + "LingBot SP input hooks did not shard all conditioning tensors consistently; " + f"expected {local_tokens} tokens per rank. Check SP hook registration and input dimensions." + ) + rotary_emb = (cosine, sine) # Phase 3: each layer receives its own cache entry. Text K/V is passed # only when absent; the returned cache is stored for subsequent DMD # steps and causal blocks in this request. @@ -1156,9 +1285,17 @@ def forward( ) cache.cross_attention[index] = cross_cache - # Phase 4: map tokens to per-patch 16-channel flow values and restore - # [B, C, F, H, W] for the Pipeline's sampler update. - hidden_states = self.head(hidden_states, timestep_embedding) + # Phase 4: project local tokens before gathering the much narrower flow values. + hidden_states = self.head( + hidden_states, + timestep_embedding, + tokens_per_frame=tokens_per_frame if sp_size > 1 else None, + token_offset=get_sp_group().ulysses_rank * (global_tokens // sp_size) if sp_size > 1 else 0, + ) + hidden_states = self.sp_output_gather(hidden_states) + projected_dim = self.config.out_channels * math.prod(self.config.patch_size) + if sp_size > 1 and hidden_states.shape != (batch_size, global_tokens, projected_dim): + raise RuntimeError("LingBot SP output hook did not gather the full projected token sequence.") return self._unpatchify( hidden_states, batch_size=batch_size, From 33fff2155e0bf9bc8688f7ffb92ef508efdf6954 Mon Sep 17 00:00:00 2001 From: GXIN <37653830+gxxx-hum@users.noreply.github.com> Date: Fri, 11 Sep 2026 21:22:15 +0800 Subject: [PATCH 77/93] [Feature][TTS] Add Speech API streaming metrics (#6853) Signed-off-by: XIN GAO <1037396230@qq.com> Co-authored-by: Hongsheng Liu --- docs/design/metrics.md | 45 +- docs/usage/metrics.md | 25 +- .../openai_api/test_serving_speech.py | 3 +- .../openai_api/test_serving_speech_metrics.py | 465 ++++++++++++++++++ .../openai_api/test_serving_speech_stream.py | 73 ++- tests/metrics/test_modality.py | 46 +- .../entrypoints/openai/serving_speech.py | 281 ++++++++--- .../openai/serving_speech_stream.py | 21 +- vllm_omni/metrics/definitions.py | 2 + vllm_omni/metrics/modality.py | 26 + 10 files changed, 895 insertions(+), 92 deletions(-) create mode 100644 tests/entrypoints/openai_api/test_serving_speech_metrics.py diff --git a/docs/design/metrics.md b/docs/design/metrics.md index 85eb7ac3876..874cbd310c4 100644 --- a/docs/design/metrics.md +++ b/docs/design/metrics.md @@ -60,7 +60,7 @@ Upstream per-engine metrics retain the `vllm:` prefix but are now registered by There are five independent paths for metric collection. -**Path 1: Pipeline-level metrics (`vllm_omni:*`)** +#### Path 1: Pipeline-level metrics (`vllm_omni:*`) `OmniPrometheusMetrics` registers the Gauge / Counter / Histogram collectors at import time. It is instantiated once per entrypoint, labeled with the model name. The entrypoint calls its methods as requests progress: @@ -69,7 +69,7 @@ There are five independent paths for metric collection. - `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` on the existing completion Counter. - `inc_requests_failed(reason)` — increments the dedicated failure-attribution Counter once per request. Reasons are normalized to the bounded set `client_abort`, `client_disconnect`, `stage_error`, and `unknown`. -**Path 2: Image and diffusion metrics** +#### Path 2: Image and diffusion metrics Image and diffusion metrics use two cooperating collectors. `OmniPrometheusMetrics` owns pipeline/stage workload families, while `OmniModalityMetrics` owns per-replica diffusion breakdowns. Finished stage messages are consumed once in `OmniBase._process_single_result`, which prevents replayed terminal messages from incrementing Counters or Histograms twice. @@ -81,15 +81,15 @@ Image and diffusion metrics use two cooperating collectors. `OmniPrometheusMetri Zero and missing values have different meanings. Valid zero-duration queue observations are recorded, while unavailable optional measurements are omitted instead of being exported as synthetic zeroes. `peak_memory_mb`, `diffusion_forward_s`, `diffusion_kv_load_s`, `vae_decode_s`, and `kv_wait_s` therefore appear only when their data source is active. -**Path 3: Audio modality metrics (`vllm_omni:audio_*`)** +#### Path 3: Audio modality metrics (`vllm_omni:audio_*`) `OmniModalityMetrics` registers seven audio families with `{model_name, stage, replica}` (plus an extra `threshold_ms` / `reason` label on the two extra-cardinality Counters). Three observation sites: - `observe_modality_at_finalize(...)` — called from `omni_base._process_single_result` inside the existing `e2e_done` finalize guard. For `output_type == "audio"` it emits `audio_frames_total`, `audio_duration_s`, `audio_rtf` (or `audio_skipped_requests_total{reason="no_audio_data"}` when no audio was produced). Sample rate is resolved from `engine_outputs.multimodal_output` via `definitions.resolve_audio_sample_rate(...)` (fallback chain mirrors `serving_chat.py`'s audio response path). -- `observe_audio_first_packet(...)` — called from the OpenAI SSE audio branch in `serving_chat.py` on the first audio packet for a request. The once-per-request guard is held by `ClientRequestState.first_audio_ts`. The `request_arrival_ts` anchor is stored in `ClientRequestState` by `async_omni.generate()`, computed at request entry. -- `observe_audio_streaming_finalize(...)` — called from `serving_chat.py` after the streaming chunk loop exhausts. It runs the per-chunk player simulation from `vllm_omni/benchmarks/audio_continuity.py` to compute the worst-case underrun and emits `audio_underrun_s` plus (when the request stayed below the threshold) `audio_continuity_ok_total{threshold_ms}`. Per-chunk PCM byte counts and arrival timestamps are recorded by the same audio branch that updates `first_audio_ts`. +- `observe_audio_first_packet(...)` — called from the OpenAI SSE audio branch in `serving_chat.py` and the streaming TTS output paths in `serving_speech.py` on the first non-empty audio payload for a request. Streaming Speech observes the first PCM payload rather than a header-only WAV chunk. HTTP carries the middleware request timestamp through response generation; WebSocket uses the start of each sentence synthesis request. Non-streaming Speech does not emit TTFP and uses `e2e_request_latency_s` for overall latency instead. +- `observe_audio_streaming_finalize(...)` — called from `serving_chat.py` and `serving_speech.py` after a streaming chunk loop exhausts normally. It runs the per-chunk player simulation from `vllm_omni/benchmarks/audio_continuity.py` to compute the worst-case underrun and emits `audio_underrun_s` plus (when the request stayed below the threshold) `audio_continuity_ok_total{threshold_ms}`. Per-chunk PCM byte counts and arrival timestamps are recorded by the same output paths that update TTFP. Speech records raw PCM payloads before raw/SSE/WebSocket encoding, excludes WAV headers and empty chunks, and does not emit continuity samples for failed, cancelled, or non-streaming requests. -**Path 4: Cross-stage transfer metrics (`vllm_omni:transfer_*`)** +#### Path 4: Cross-stage transfer metrics (`vllm_omni:transfer_*`) `OmniTransferMetrics` registers four Histogram families with `{model_name, from_stage, from_replica, to_stage, to_replica}` labels. Each observation corresponds to one physical transfer hop (one chunk between adjacent stages), not the per-request accumulated total — so the histograms track per-transfer distribution. @@ -100,7 +100,7 @@ The hook lives in `OrchestratorAggregator.record_transfer_tx` and `record_transf Defensive fail-safe: if `transfer_emitter` or `replica_resolver` is missing, or the resolver returns `None` for either side, the emit is skipped silently (the underlying `TransferEdgeStats` accumulation is unaffected). -**Path 5: Per-engine metrics (`vllm:*`, stage/replica wrap)** +#### Path 5: Per-engine metrics (`vllm:*`, stage/replica wrap) The Orchestrator instantiates `OmniPrometheusStatLogger` (a thin subclass of upstream `vllm.v1.metrics.loggers.PrometheusStatLogger`) and feeds it scheduler stats and iteration stats after processing each batch of engine outputs. This populates the standard ~37 vLLM metric families (TTFT, ITL, TPOT, KV cache usage, etc.) using the same upstream code path — but with the `engine` label reshaped into `stage` + `replica` so multi-replica deployments produce distinct series per replica. See the next section for the wrap mechanics. @@ -174,7 +174,7 @@ second. ### Pipeline (4) | Metric | Type | Labels | Description | -|--------|------|--------|-------------| +| -------- | ------ | -------- | ------------- | | `vllm_omni:num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages | | `vllm_omni:num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled | | `vllm_omni:requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` | @@ -183,7 +183,7 @@ second. ### Image and diffusion service-level metrics | Metric | Type | Labels | Description | -|--------|------|--------|-------------| +| -------- | ------ | -------- | ------------- | | `vllm_omni:stage_gen_time_s` | Histogram | `model_name`, `stage`, `stage_type` | Stage submit to finished output; includes in-stage queueing | | `vllm_omni:request_queue_wait_s` | Histogram | `model_name` | Orchestration-layer queue wait; a present zero is recorded | | `vllm_omni:stage_waiting_requests` | Gauge | `model_name`, `stage` | Sum of the latest waiting snapshots across live replicas in the stage | @@ -200,7 +200,7 @@ second. Labels: `{model_name, stage, replica}`. | Metric | Type | Description | -|--------|------|-------------| +| -------- | ------ | ------------- | | `vllm_omni:diffusion_exec_s` | Histogram | Core diffusion-step execution time per request | | `vllm_omni:diffusion_exec_per_step_s` | Histogram | Core diffusion execution divided by inference-step count | | `vllm_omni:diffusion_preprocess_s` | Histogram | Diffusion input preprocessing time | @@ -218,7 +218,7 @@ The optional breakdown and memory families are sparse by design: when an engine Labels: `{model_name, stage, replica}` plus the listed extra label. | Metric | Type | Extra label | Description | -|--------|------|-------------|-------------| +| -------- | ------ | ------------- | ------------- | | `vllm_omni:audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame | | `vllm_omni:audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) | | `vllm_omni:audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` | @@ -227,12 +227,29 @@ Labels: `{model_name, stage, replica}` plus the listed extra label. | `vllm_omni:audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` | | `vllm_omni:audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio | +### Speech streaming (2) + +Labels: `{model_name}` plus the listed extra label. + +| Metric | Type | Extra label | Description | +| -------- | ------ | ------------- | ------------- | +| `vllm_omni:speech_stream_aborted_total` | Counter | `reason` | Interrupted audio generators, including before first PCM; reasons: `cancelled`, `closed`, `engine_dead`, `error` | +| `vllm_omni:speech_stream_completed_total` | Counter | — | Normally completed audio generators; does not confirm client receipt | + +Scope: started Speech audio generators (raw/SSE/WebSocket), counted once per +generation; WebSocket counts each sentence. Chat, non-streaming Speech, and +failures before generator execution are excluded. Interrupted streams do not +emit continuity or underrun samples. + +Interruption ratio: `aborted / (aborted + completed)`, using increments over the +same time window and summing all reasons per model. + ### Cross-stage transfer (4) Labels: `{model_name, from_stage, from_replica, to_stage, to_replica}`. | Metric | Type | Description | -|--------|------|-------------| +| -------- | ------ | ------------- | | `vllm_omni:transfer_size_bytes` | Histogram | Per-transfer payload size in bytes | | `vllm_omni:transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) | | `vllm_omni:transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) | @@ -245,8 +262,8 @@ After the wrap, every upstream `vllm:*` family — TTFT, ITL, TPOT, e2e latency, ## Naming Convention - All time-bearing metrics use the `_s` suffix (values in seconds). Two bucket families are used: - - `SECONDS_BUCKETS` (0.05 s – 300 s) for e2e / generation / TTFP style values. - - `SECONDS_FAST_BUCKETS` (0.001 s – 60 s) for fine-grained cross-stage transfer and audio-underrun values that need millisecond-level resolution. + - `SECONDS_BUCKETS` (0.05 s – 300 s) for e2e / generation / TTFP style values. + - `SECONDS_FAST_BUCKETS` (0.001 s – 60 s) for fine-grained cross-stage transfer and audio-underrun values that need millisecond-level resolution. - Counters use the `_total` suffix (auto-appended by `prometheus_client`). - Sizes use the `_bytes` suffix. - All omni-specific families are prefixed `vllm_omni:`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm_omni:` / `vllm_omni`. diff --git a/docs/usage/metrics.md b/docs/usage/metrics.md index 184a76eefad..1e0c1d06876 100644 --- a/docs/usage/metrics.md +++ b/docs/usage/metrics.md @@ -12,7 +12,7 @@ curl http://localhost:8000/metrics ## Metric Namespaces | Prefix | Source | Present when | -|--------|--------|--------------| +| -------- | -------- | -------------- | | `vllm_omni:` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent | | `vllm:` | Upstream vLLM engine, wrapped by `OmniPrometheusStatLogger` to expose `{stage, replica}` | Pipeline includes an LLM (AR) stage | | `http_` / `process_` | Uvicorn / Python runtime | Always | @@ -24,7 +24,7 @@ Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the ### Request counts | Metric | Type | Labels | Description | -|--------|------|--------|-------------| +| -------- | ------ | -------- | ------------- | | `vllm_omni:num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) | | `vllm_omni:num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue | | `vllm_omni:requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` | @@ -32,15 +32,15 @@ Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the ### Latency | Metric | Type | Labels | Description | -|--------|------|--------|-------------| +| -------- | ------ | -------- | ------------- | | `vllm_omni:e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds | ## Audio Modality Metrics (`vllm_omni:`) -Emitted at request finalize, except for `audio_ttfp_s` (streaming-hook at the first audio packet) and `audio_underrun_s` / `audio_continuity_ok_total` (streaming finalize, after the chunk stream is exhausted). All carry `{model_name, stage, replica}` plus the listed extra label. +Emitted at request finalize, except for `audio_ttfp_s` (the first audio packet/frame from streaming Chat or Speech APIs) and `audio_underrun_s` / `audio_continuity_ok_total` (Chat or Speech streaming finalize, after the chunk stream is exhausted). Non-streaming Speech latency is represented by `e2e_request_latency_s`, not TTFP. All audio metrics carry `{model_name, stage, replica}` plus the listed extra label. | Metric | Type | Extra label | Description | -|--------|------|-------------|-------------| +| -------- | ------ | ------------- | ------------- | | `vllm_omni:audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame | | `vllm_omni:audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) | | `vllm_omni:audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` | @@ -50,13 +50,22 @@ Emitted at request finalize, except for `audio_ttfp_s` (streaming-hook at the fi | `vllm_omni:audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio | The continuity math comes from `vllm_omni/benchmarks/audio_continuity.py::compute_continuity_stats` so the server-side observation aligns with the bench-side definition. +For Speech, continuity applies to raw audio, SSE, and WebSocket streaming. WAV headers and empty chunks are excluded; non-streaming Speech requests do not emit continuity samples. + +Continuity rate for the default 100 ms threshold can be derived from the successful-continuity counter and the underrun Histogram's request count: + +```promql +sum by (model_name, stage, replica) (rate(vllm_omni:audio_continuity_ok_total{threshold_ms="100"}[5m])) +/ +sum by (model_name, stage, replica) (rate(vllm_omni:audio_underrun_s_count[5m])) +``` ## Diffusion Metrics (`vllm_omni:`) Per-request timing breakdowns for diffusion (image/video) stages. Emitted at request finalize when `stage_metrics.diffusion_metrics` is present. All carry `{model_name, stage, replica}`. | Metric | Type | Description | -|--------|------|-------------| +| -------- | ------ | ------------- | | `vllm_omni:diffusion_exec_s` | Histogram | DiT forward pass execution time per request in seconds | | `vllm_omni:diffusion_exec_per_step_s` | Histogram | DiT forward pass execution time per denoising step in seconds (`exec_s / num_inference_steps`) | | `vllm_omni:diffusion_preprocess_s` | Histogram | Diffusion input preprocessing time per request in seconds | @@ -80,7 +89,7 @@ rate(vllm_omni:diffusion_preprocess_s_sum[5m]) Per-physical-transfer histograms tracking the data hop between adjacent stages. Labels `{model_name, from_stage, from_replica, to_stage, to_replica}` let dashboards attribute latency to specific replica edges. `from_replica` / `to_replica` are resolved from the orchestrator's sticky-routing binding (`stage_pool.get_bound_replica_id(request_id)`), so no extra plumbing through `TransferEdgeStats` is needed. | Metric | Type | Description | -|--------|------|-------------| +| -------- | ------ | ------------- | | `vllm_omni:transfer_size_bytes` | Histogram | Per-transfer payload size in bytes | | `vllm_omni:transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) | | `vllm_omni:transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) | @@ -106,7 +115,7 @@ For the full list of upstream metrics, see [the vLLM docs](https://github.com/vl ## Metric Availability by Pipeline Type | Metric group | Multi-stage LLM (Qwen3-Omni) | -|---|---| +| --- | --- | | `vllm_omni:` request tracking + latency | With `--log-stats` | | `vllm_omni:` audio modality | With `--log-stats`, if pipeline has a talker stage | | `vllm_omni:` transfer | With `--log-stats`, if pipeline has ≥ 2 stages | diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 6040bede35e..732f86f2f53 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -3287,7 +3287,7 @@ def test_raw_stream_passes_tts_params_to_common_guard(self, streaming_app, monke finalized_tts_params = {"_qwen3_tts_effective_max_tokens": [192]} captured: dict = {} - async def prepare(_request, request_id=None): + async def prepare(_request, request_id=None, arrival_time=None): return request_id, object(), finalized_tts_params async def generate_chunks( @@ -3296,6 +3296,7 @@ async def generate_chunks( _response_format="pcm", raw_request=None, request_start_s=None, + request_arrival_ts=None, include_sample_rate=False, usage_acc=None, tts_params=None, diff --git a/tests/entrypoints/openai_api/test_serving_speech_metrics.py b/tests/entrypoints/openai_api/test_serving_speech_metrics.py new file mode 100644 index 00000000000..a5e284a0177 --- /dev/null +++ b/tests/entrypoints/openai_api/test_serving_speech_metrics.py @@ -0,0 +1,465 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Prometheus metric coverage for the Speech API audio path.""" + +import asyncio +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from vllm_omni.entrypoints.openai import serving_speech as speech_module +from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest +from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _MetricsStub: + def __init__(self) -> None: + self.ttfp_calls: list[tuple[str, str, float]] = [] + self.underrun_calls: list[tuple[str, str, float]] = [] + self.continuity_calls: list[tuple[str, str, int]] = [] + self.abort_calls: list[str] = [] + self.completed = 0 + + def inc_speech_stream_completed(self) -> None: + self.completed += 1 + + def inc_speech_stream_aborted(self, reason: str) -> None: + self.abort_calls.append(reason) + + def observe_audio_ttfp(self, stage: str, replica: str, seconds: float) -> None: + self.ttfp_calls.append((stage, replica, seconds)) + + def observe_audio_underrun(self, stage: str, replica: str, seconds: float) -> None: + self.underrun_calls.append((stage, replica, seconds)) + + def inc_audio_continuity_ok(self, stage: str, replica: str, threshold_ms: int) -> None: + self.continuity_calls.append((stage, replica, threshold_ms)) + + +def _serving(metrics: _MetricsStub, *, adapter=None) -> OmniOpenAIServingSpeech: + serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech) + serving._tts_model_type = "qwen3_tts" + serving.engine_client = SimpleNamespace(mod_metrics=metrics, request_states={}) + serving._get_tts_adapter = lambda: adapter + serving.create_audio = lambda audio_obj: SimpleNamespace( + audio_data=b"\0\0" * int(audio_obj.audio_tensor.size), + media_type="audio/pcm", + ) + serving._mark_ref_audio_artifact_ready_for_request = lambda request_id: None + serving._discard_ref_audio_artifact_warmup = lambda request_id: None + return serving + + +def _result(samples: int = 320, *, stage_id: int = 1, replica_id: int | None = 2) -> SimpleNamespace: + return SimpleNamespace( + multimodal_output={"audio": torch.zeros(samples), "sr": 16000}, + stage_id=stage_id, + replica_id=replica_id, + ) + + +async def _generate(*results): + for result in results: + yield result + + +@pytest.mark.asyncio +async def test_streaming_speech_observes_ttfp_once_on_first_pcm_payload(monkeypatch): + metrics = _MetricsStub() + serving = _serving(metrics) + monkeypatch.setattr("vllm_omni.entrypoints.openai.serving_speech.time.time", lambda: 100.25) + + chunks = serving._generate_audio_chunks( + _generate(_result(), _result()), + request_id="speech-test", + request_arrival_ts=100.0, + ) + assert len([chunk async for chunk in chunks]) == 2 + assert metrics.ttfp_calls == [("1", "2", pytest.approx(0.25))] + assert metrics.underrun_calls == [("1", "2", pytest.approx(0.0))] + assert metrics.continuity_calls == [("1", "2", 100)] + assert metrics.abort_calls == [] + assert metrics.completed == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("before_pcm", [True, False]) +@pytest.mark.parametrize( + ("error_type", "reason"), + [(asyncio.CancelledError, "cancelled"), (speech_module.EngineDeadError, "engine_dead"), (RuntimeError, "error")], +) +async def test_speech_stream_abort_counted_once_without_continuity(before_pcm, error_type, reason): + metrics = _MetricsStub() + serving = _serving(metrics) + + async def results(): + if not before_pcm: + yield _result() + raise error_type() + + chunks = serving._generate_audio_chunks(results(), request_id="speech-test", request_arrival_ts=100.0) + with pytest.raises(error_type): + _ = [chunk async for chunk in chunks] + + assert metrics.abort_calls == [reason] + assert metrics.completed == 0 + assert len(metrics.ttfp_calls) == (0 if before_pcm else 1) + assert metrics.underrun_calls == [] + assert metrics.continuity_calls == [] + + +@pytest.mark.asyncio +async def test_speech_stream_close_counted_once_without_continuity(): + metrics = _MetricsStub() + serving = _serving(metrics) + chunks = serving._generate_audio_chunks( + _generate(_result(), _result()), request_id="speech-test", request_arrival_ts=100.0 + ) + assert await anext(chunks) + await chunks.aclose() + await chunks.aclose() + assert metrics.abort_calls == ["closed"] + assert metrics.completed == 0 + assert metrics.underrun_calls == [] + assert metrics.continuity_calls == [] + + +@pytest.mark.asyncio +async def test_streaming_speech_retries_ttfp_labels_without_moving_first_packet_time(monkeypatch): + metrics = _MetricsStub() + serving = _serving(metrics) + clock = {"now": 100.25, "perf": 0.25} + monkeypatch.setattr(speech_module.time, "time", lambda: clock["now"]) + monkeypatch.setattr(speech_module.time, "perf_counter", lambda: clock["perf"]) + + async def results(): + yield _result(replica_id=None) + clock["now"] = 101.0 + clock["perf"] = 0.26 + yield _result(replica_id=2) + clock["now"] = 101.25 + clock["perf"] = 0.27 + yield _result(replica_id=2) + + chunks = serving._generate_audio_chunks( + results(), + request_id="speech-test", + request_arrival_ts=100.0, + ) + assert len([chunk async for chunk in chunks]) == 3 + assert metrics.ttfp_calls == [("1", "2", pytest.approx(0.25))] + assert metrics.underrun_calls == [("1", "2", pytest.approx(0.0))] + assert metrics.continuity_calls == [("1", "2", 100)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["raw", "sse", "pcm"]) +async def test_speech_disconnect_during_send_closes_entire_generator_chain(mode): + metrics = _MetricsStub() + serving = _serving(metrics) + closed = [] + + async def source(): + try: + yield _result() + yield _result() + finally: + closed.append(True) + + # Retain the source to prove cleanup does not depend on garbage collection. + generator = source() + method = { + "raw": serving._generate_audio_chunks, + "sse": serving._generate_audio_sse_events, + "pcm": serving._generate_pcm_chunks, + }[mode] + stream = method(generator, request_id="speech-test", request_arrival_ts=100.0) + + async def send(message): + if message["type"] == "http.response.body": + raise OSError("client disconnected") + + if mode == "pcm": + # WebSocket handlers already explicitly close their PCM iterator. + assert await anext(stream) + await stream.aclose() + else: + response = speech_module._SpeechStreamingResponse(stream) + with pytest.raises(OSError, match="client disconnected"): + await response.stream_response(send) + assert closed == [True] + assert metrics.abort_calls == ["closed"] + assert metrics.completed == 0 + assert metrics.underrun_calls == [] + assert metrics.continuity_calls == [] + + +@pytest.mark.parametrize("arrival_ts", [0.0, -1.0]) +def test_speech_ttfp_invalid_arrival_does_not_consume_guard(arrival_ts): + metrics = _MetricsStub() + serving = _serving(metrics) + state = SimpleNamespace( + external_request_id="speech-test", + request_arrival_ts=arrival_ts, + first_audio_ts=None, + audio_emit_stage_id=None, + audio_emit_replica_id=None, + ) + serving.engine_client.request_states = {"internal-test": state} + + assert serving._observe_speech_audio_ttfp(request_id="speech-test", result=_result(), first_packet_ts=100.25) == ( + 1, + 2, + False, + ) + assert metrics.ttfp_calls == [] + assert state.first_audio_ts is None + assert state.audio_emit_stage_id is None + assert state.audio_emit_replica_id is None + + state.request_arrival_ts = 100.0 + assert serving._observe_speech_audio_ttfp(request_id="speech-test", result=_result(), first_packet_ts=100.25) == ( + 1, + 2, + True, + ) + assert metrics.ttfp_calls == [("1", "2", pytest.approx(0.25))] + assert state.first_audio_ts == 100.25 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("during_send", [False, True]) +async def test_speech_response_task_cancellation_closes_source(during_send): + metrics = _MetricsStub() + serving = _serving(metrics) + waiting = asyncio.Event() + blocker = asyncio.Event() + closed = [] + + async def source(): + try: + if not during_send: + waiting.set() + await blocker.wait() + yield _result() + finally: + closed.append(True) + + async def send(message): + if during_send and message["type"] == "http.response.body": + waiting.set() + await blocker.wait() + + generator = source() + stream = serving._generate_audio_chunks(generator, request_id="speech-test", request_arrival_ts=100.0) + response = speech_module._SpeechStreamingResponse(stream) + task = asyncio.create_task(response.stream_response(send)) + try: + await asyncio.wait_for(waiting.wait(), timeout=5) + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert closed == [True] + assert metrics.abort_calls == ["closed" if during_send else "cancelled"] + assert metrics.completed == 0 + assert metrics.underrun_calls == [] + + +@pytest.mark.asyncio +async def test_streaming_speech_does_not_count_empty_payload_as_first_packet(monkeypatch): + metrics = _MetricsStub() + serving = _serving(metrics) + monkeypatch.setattr("vllm_omni.entrypoints.openai.serving_speech.time.time", lambda: 100.25) + + chunks = serving._generate_audio_chunks( + _generate(_result(samples=0), _result()), + request_id="speech-test", + request_arrival_ts=100.0, + ) + assert len([chunk async for chunk in chunks]) == 2 + assert metrics.ttfp_calls == [("1", "2", pytest.approx(0.25))] + assert metrics.underrun_calls == [("1", "2", pytest.approx(0.0))] + assert metrics.continuity_calls == [("1", "2", 100)] + + +@pytest.mark.asyncio +async def test_streaming_speech_does_not_finalize_continuity_on_error(monkeypatch): + metrics = _MetricsStub() + serving = _serving(metrics) + monkeypatch.setattr("vllm_omni.entrypoints.openai.serving_speech.time.time", lambda: 100.25) + + async def failing_stream(): + yield _result() + raise RuntimeError("stream failed") + + chunks = serving._generate_audio_chunks( + failing_stream(), + request_id="speech-test", + request_arrival_ts=100.0, + ) + with pytest.raises(RuntimeError, match="stream failed"): + _ = [chunk async for chunk in chunks] + + assert len(metrics.ttfp_calls) == 1 + assert metrics.underrun_calls == [] + assert metrics.continuity_calls == [] + + +@pytest.mark.asyncio +async def test_streaming_speech_does_not_finalize_continuity_on_validation_error(monkeypatch): + metrics = _MetricsStub() + + def reject_generation(_tts_params, **_kwargs): + raise RuntimeError("generation validation failed") + + adapter = SimpleNamespace(validates_generation=True, validate_generation=reject_generation) + serving = _serving(metrics, adapter=adapter) + monkeypatch.setattr("vllm_omni.entrypoints.openai.serving_speech.time.time", lambda: 100.25) + + chunks = serving._generate_audio_chunks( + _generate(_result()), + request_id="speech-test", + request_arrival_ts=100.0, + tts_params={"task_type": ["Base"]}, + ) + with pytest.raises(RuntimeError, match="generation validation failed"): + _ = [chunk async for chunk in chunks] + + assert len(metrics.ttfp_calls) == 1 + assert metrics.underrun_calls == [] + assert metrics.continuity_calls == [] + + +@pytest.mark.asyncio +async def test_streaming_speech_reports_late_chunk_as_underrun(monkeypatch): + metrics = _MetricsStub() + serving = _serving(metrics) + monkeypatch.setattr("vllm_omni.entrypoints.openai.serving_speech.time.time", lambda: 100.25) + perf_times = iter((0.0, 0.1, 0.1, 2.0, 2.0)) + monkeypatch.setattr("vllm_omni.entrypoints.openai.serving_speech.time.perf_counter", lambda: next(perf_times)) + + chunks = serving._generate_audio_chunks( + _generate(_result(), _result()), + request_id="speech-test", + request_arrival_ts=100.0, + ) + assert len([chunk async for chunk in chunks]) == 2 + assert metrics.underrun_calls[0][:2] == ("1", "2") + assert metrics.underrun_calls[0][2] > 0.1 + assert metrics.continuity_calls == [] + + +@pytest.mark.asyncio +async def test_non_streaming_speech_does_not_observe_ttfp(): + metrics = _MetricsStub() + serving = _serving(metrics) + serving._audio_encode_speed = lambda _request: 1.0 + + async def prepare(_request, **_kwargs): + result = _result() + result.metrics = {} + return "speech-test", _generate(result), {} + + serving._prepare_speech_generation = prepare + request = OpenAICreateSpeechRequest(input="hello", response_format="pcm") + + audio_data, media_type = await serving._generate_audio_bytes(request, request_arrival_ts=100.0) + + assert audio_data + assert media_type == "audio/pcm" + assert metrics.ttfp_calls == [] + assert metrics.completed == 0 + assert metrics.abort_calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("flush_only", [False, True]) +@pytest.mark.parametrize("response_format", ["pcm", "wav"]) +async def test_resampled_speech_records_flush_with_audio_producer(monkeypatch, flush_only, response_format): + metrics = _MetricsStub() + serving = _serving(metrics) + clock = {"now": 0.0} + finalized = [] + original_finalize = speech_module.observe_audio_streaming_finalize + + def capture_finalize(*args, **kwargs): + finalized.append(kwargs) + original_finalize(*args, **kwargs) + + class BufferedResampler: + def __init__(self, source_rate, target_rate): + assert (source_rate, target_rate) == (16000, 24000) + + def process(self, chunk, *, final=False): + clock["now"] = 0.5 if final else 0.25 + if not final and flush_only: + return np.empty(0, dtype=np.float32) + return np.zeros(2400, dtype=np.float32) + + monkeypatch.setattr(speech_module, "StreamingAudioResampler", BufferedResampler) + monkeypatch.setattr(speech_module, "observe_audio_streaming_finalize", capture_finalize) + monkeypatch.setattr(speech_module.time, "time", lambda: 100.0 + clock["now"]) + monkeypatch.setattr(speech_module.time, "perf_counter", lambda: clock["now"]) + # The last result is not audio and must not supply the flush metric labels. + non_audio = SimpleNamespace(multimodal_output={"timestamps": []}, stage_id=9, replica_id=8) + chunks = [ + chunk + async for chunk in serving._generate_audio_chunks( + _generate(_result(), non_audio), + request_id="speech-test", + response_format=response_format, + request_start_s=0.0, + request_arrival_ts=100.0, + target_sample_rate=24000, + ) + ] + + if response_format == "wav": + assert chunks.pop(0).startswith(b"RIFF") + expected_arrivals = [0.5] if flush_only else [0.25, 0.5] + assert [len(chunk) for chunk in chunks] == [4800] * len(expected_arrivals) + assert metrics.ttfp_calls == [("1", "2", pytest.approx(expected_arrivals[0]))] + assert len(finalized) == 1 + assert finalized[0]["sample_rate"] == 24000 + assert finalized[0]["channels"] == 1 + assert finalized[0]["chunk_bytes"] == [4800] * len(expected_arrivals) + assert finalized[0]["chunk_arrival_times_s"] == expected_arrivals + assert metrics.underrun_calls == [("1", "2", pytest.approx(0.0 if flush_only else 0.15))] + assert metrics.continuity_calls == ([("1", "2", 100)] if flush_only else []) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("empty_pcm", [False, True]) +async def test_resampled_speech_without_pcm_does_not_emit_metrics(monkeypatch, empty_pcm): + metrics = _MetricsStub() + serving = _serving(metrics) + + class EmptyResampler: + def __init__(self, source_rate, target_rate): + pass + + def process(self, chunk, *, final=False): + # Also cover a nonempty flush waveform whose encoder emits no PCM. + return np.zeros(10 if final and empty_pcm else 0, dtype=np.float32) + + monkeypatch.setattr(speech_module, "StreamingAudioResampler", EmptyResampler) + serving.create_audio = lambda audio_obj: SimpleNamespace(audio_data=b"", media_type="audio/pcm") + chunks = [ + chunk + async for chunk in serving._generate_audio_chunks( + _generate(_result()), + request_id="speech-test", + request_arrival_ts=100.0, + target_sample_rate=24000, + ) + ] + assert not any(chunks) + assert metrics.ttfp_calls == [] + assert metrics.underrun_calls == [] + assert metrics.continuity_calls == [] diff --git a/tests/entrypoints/openai_api/test_serving_speech_stream.py b/tests/entrypoints/openai_api/test_serving_speech_stream.py index 6414ab4ff8b..39c9e66d950 100644 --- a/tests/entrypoints/openai_api/test_serving_speech_stream.py +++ b/tests/entrypoints/openai_api/test_serving_speech_stream.py @@ -48,7 +48,14 @@ def _build_test_app( speech_service.forced_aligner_enabled = False async def mock_generate_pcm_chunks( - _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None + _generator, + _request_id, + *, + request_start_s=None, + request_arrival_ts=None, + include_sample_rate=False, + tts_params=None, + collect=None, ): for chunk in (b"\x01\x02", b"\x03\x04\x05"): yield (chunk, 24000) if include_sample_rate else chunk @@ -249,6 +256,7 @@ def test_idle_timeout_closes_reused_connection(self, mocker: MockerFixture): def test_streaming_multiple_binary_frames(self, mocker: MockerFixture): captured_requests = [] + captured_timing = {} captured_tts_params = [] speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech) @@ -257,13 +265,27 @@ def test_streaming_multiple_binary_frames(self, mocker: MockerFixture): speech_service.engine_client.abort = mocker.AsyncMock() speech_service.forced_aligner_enabled = False - async def mock_prepare_speech_generation(request): + async def mock_prepare_speech_generation(request, *, arrival_time=None): captured_requests.append(request) + assert arrival_time is not None + captured_timing["prepare_arrival"] = arrival_time return "req-stream", object(), {"_qwen3_tts_effective_max_tokens": [192]} speech_service._prepare_speech_generation = mock_prepare_speech_generation - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None): + async def mock_generate_pcm_chunks( + _generator, + _request_id, + *, + request_start_s=None, + request_arrival_ts=None, + include_sample_rate=False, + tts_params=None, + ): + assert request_start_s is not None + assert request_arrival_ts is not None + captured_timing["chunk_start"] = request_start_s + captured_timing["chunk_arrival"] = request_arrival_ts captured_tts_params.append(tts_params) for chunk in (b"\x01\x02", b"\x03\x04\x05", b"\x06"): yield (chunk, 24000) if include_sample_rate else chunk @@ -309,6 +331,8 @@ async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_ra assert captured_requests[0].stream is True assert captured_requests[0].response_format == "pcm" assert captured_requests[0].initial_codec_chunk_frames == 12 + assert captured_timing["prepare_arrival"] == captured_timing["chunk_arrival"] + assert captured_timing["chunk_start"] > 0 assert captured_tts_params == [{"_qwen3_tts_effective_max_tokens": [192]}] assert speech_service._generate_audio_bytes.await_count == 0 @@ -341,8 +365,9 @@ def test_word_timestamps_emit_pipeline_json_frame(self, mocker: MockerFixture): speech_service.engine_client.abort = mocker.AsyncMock() speech_service.forced_aligner_enabled = True - async def mock_prepare_speech_generation(request): + async def mock_prepare_speech_generation(request, *, arrival_time=None): captured_requests.append(request) + assert arrival_time is not None return "req-stream", object(), {} speech_service._prepare_speech_generation = mock_prepare_speech_generation @@ -353,8 +378,17 @@ async def mock_prepare_speech_generation(request): # The forced-aligner stage rides the same generator: its pooling output # is surfaced via the ``collect`` channel once the audio has streamed. async def mock_generate_pcm_chunks( - _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None + _generator, + _request_id, + *, + request_start_s=None, + request_arrival_ts=None, + include_sample_rate=False, + tts_params=None, + collect=None, ): + assert request_start_s is not None + assert request_arrival_ts is not None for chunk in (first_chunk, second_chunk): yield (chunk, 1000) if include_sample_rate else chunk if collect is not None: @@ -438,7 +472,14 @@ def test_word_timestamps_emit_word_dicts(self, mocker: MockerFixture): speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {})) async def mock_generate_pcm_chunks( - _generator, _request_id, *, include_sample_rate=False, tts_params=None, collect=None + _generator, + _request_id, + *, + request_start_s=None, + request_arrival_ts=None, + include_sample_rate=False, + tts_params=None, + collect=None, ): chunk = b"\x01" * 1000 yield (chunk, 1000) if include_sample_rate else chunk @@ -620,7 +661,15 @@ def test_streaming_generation_error_marks_audio_done(self, mocker: MockerFixture speech_service.engine_client.abort = mocker.AsyncMock() speech_service.forced_aligner_enabled = False - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None): + async def mock_generate_pcm_chunks( + _generator, + _request_id, + *, + request_start_s=None, + request_arrival_ts=None, + include_sample_rate=False, + tts_params=None, + ): yield b"\x01\x02" raise RuntimeError("stream boom") @@ -716,7 +765,15 @@ def test_disconnect_aborts_streaming_request(self, mocker: MockerFixture): speech_service.engine_client.abort = mocker.AsyncMock() speech_service.forced_aligner_enabled = False - async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, tts_params=None): + async def mock_generate_pcm_chunks( + _generator, + _request_id, + *, + request_start_s=None, + request_arrival_ts=None, + include_sample_rate=False, + tts_params=None, + ): yield b"\x01\x02" speech_service._generate_pcm_chunks = mock_generate_pcm_chunks diff --git a/tests/metrics/test_modality.py b/tests/metrics/test_modality.py index 14bb4f309ca..a918dce53c7 100644 --- a/tests/metrics/test_modality.py +++ b/tests/metrics/test_modality.py @@ -48,6 +48,8 @@ def _sample_value(output: str, line_prefix: str) -> float | None: defs.AUDIO_UNDERRUN_S, defs.AUDIO_CONTINUITY_OK_METRIC, defs.AUDIO_SKIPPED_REQUESTS_METRIC, + defs.SPEECH_STREAM_ABORTED_METRIC, + defs.SPEECH_STREAM_COMPLETED_METRIC, defs.DIFFUSION_EXEC_S, defs.DIFFUSION_EXEC_PER_STEP_S, defs.DIFFUSION_PREPROCESS_S, @@ -70,6 +72,8 @@ def test_all_locked_families_present(self, mod: OmniModalityMetrics) -> None: mod.observe_audio_underrun("s", "r", 0.01) mod.inc_audio_continuity_ok("s", "r", 100) mod.inc_audio_skipped("s", "r", "malformed_codec") + mod.inc_speech_stream_aborted("error") + mod.inc_speech_stream_completed() mod.observe_diffusion_exec("s", "r", 0.5) mod.observe_diffusion_exec_per_step("s", "r", 0.01) mod.observe_diffusion_preprocess("s", "r", 0.01) @@ -91,6 +95,23 @@ def test_all_locked_families_present(self, mod: OmniModalityMetrics) -> None: class TestAudio: + def test_speech_stream_abort_counter_and_disabled_logging(self): + model = "test-speech-stream-abort-counter" + metrics = OmniModalityMetrics(model_name=model) + disabled = OmniModalityMetrics(model_name=model, log_stats=False) + metrics.inc_speech_stream_completed() + disabled.inc_speech_stream_completed() + assert REGISTRY.get_sample_value(defs.SPEECH_STREAM_COMPLETED_METRIC + "_total", {"model_name": model}) == 1.0 + for reason in ("cancelled", "closed", "engine_dead", "error"): + metrics.inc_speech_stream_aborted(reason) + disabled.inc_speech_stream_aborted(reason) + assert ( + REGISTRY.get_sample_value( + defs.SPEECH_STREAM_ABORTED_METRIC + "_total", {"model_name": model, "reason": reason} + ) + == 1.0 + ) + def test_audio_ttfp_observed(self, mod: OmniModalityMetrics) -> None: stage, replica = "talker_ttfp", "0" mod.observe_audio_ttfp(stage, replica, 0.42) @@ -242,6 +263,9 @@ def observe_audio_duration(self, s, r, d): def observe_audio_rtf(self, s, r, rtf): self.calls.append(("observe_audio_rtf", s, r, rtf)) + def observe_audio_ttfp(self, s, r, t): + self.calls.append(("observe_audio_ttfp", s, r, t)) + def observe_audio_underrun(self, s, r, u): self.calls.append(("observe_audio_underrun", s, r, u)) @@ -599,7 +623,6 @@ def test_audio_diffusion_stage_fires_both_paths(self): class TestObserveAudioFirstPacket: def test_observes_with_valid_inputs(self): stub = _StubModMetrics() - stub.observe_audio_ttfp = lambda s, r, t: stub.calls.append(("observe_audio_ttfp", s, r, t)) observe_audio_first_packet( stub, @@ -612,19 +635,16 @@ def test_observes_with_valid_inputs(self): def test_replica_none_skipped(self): stub = _StubModMetrics() - stub.observe_audio_ttfp = lambda s, r, t: stub.calls.append(("observe_audio_ttfp", s, r, t)) observe_audio_first_packet(stub, stage_id=1, replica_id=None, arrival_ts=100.0, now_ts=100.5) assert stub.calls == [] def test_arrival_ts_zero_skipped(self): stub = _StubModMetrics() - stub.observe_audio_ttfp = lambda s, r, t: stub.calls.append(("observe_audio_ttfp", s, r, t)) observe_audio_first_packet(stub, stage_id=1, replica_id=0, arrival_ts=0.0, now_ts=100.5) assert stub.calls == [] def test_clock_skew_clamped_to_zero(self): stub = _StubModMetrics() - stub.observe_audio_ttfp = lambda s, r, t: stub.calls.append(("observe_audio_ttfp", s, r, t)) observe_audio_first_packet(stub, stage_id=1, replica_id=0, arrival_ts=100.5, now_ts=100.0) assert stub.calls == [("observe_audio_ttfp", "1", "0", 0.0)] @@ -668,6 +688,24 @@ def test_late_chunk_emits_nonzero_underrun_and_no_continuity_inc(self): assert underrun_calls[0][-1] > 0.1 assert not any(c[0] == "inc_audio_continuity_ok" for c in stub.calls) + def test_stereo_channel_count_is_used_for_playback_rate(self): + stub = _StubModMetrics() + # At 100 Hz, s16le stereo consumes 400 bytes/s. The first 200-byte + # chunk buffers 0.5s, so a second chunk arriving at 1.0s underruns by 0.5s. + observe_audio_streaming_finalize( + stub, + stage_id=1, + replica_id=0, + chunk_arrival_times_s=[0.0, 1.0], + chunk_bytes=[200, 200], + sample_rate=100, + channels=2, + threshold_s=0.1, + ) + underrun_calls = [c for c in stub.calls if c[0] == "observe_audio_underrun"] + assert underrun_calls == [("observe_audio_underrun", "1", "0", pytest.approx(0.5))] + assert not any(c[0] == "inc_audio_continuity_ok" for c in stub.calls) + def test_empty_arrivals_skipped(self): stub = _StubModMetrics() observe_audio_streaming_finalize( diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 2f4e472f824..d4aa72818ae 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -13,12 +13,14 @@ import time from collections import OrderedDict from concurrent.futures import ThreadPoolExecutor +from contextlib import aclosing from http import HTTPStatus from pathlib import Path -from typing import Any +from typing import Any, cast from urllib.parse import urlparse from urllib.request import url2pathname +import anyio import numpy as np import soundfile as sf import torch @@ -59,6 +61,7 @@ tts_entry_stage_archs, ) from vllm_omni.entrypoints.utils import coerce_param_message_types +from vllm_omni.metrics.modality import observe_audio_first_packet, observe_audio_streaming_finalize from vllm_omni.outputs import OmniRequestOutput from vllm_omni.utils.speaker_cache import get_speaker_cache @@ -194,6 +197,17 @@ def _validate_path_within_directory(file_path: Path, directory: Path) -> bool: return False +class _SpeechStreamingResponse(StreamingResponse): + async def stream_response(self, send) -> None: + try: + await super().stream_response(send) + finally: + # A disconnect during send leaves the iterator suspended at yield. + # Close it explicitly, even inside Starlette's cancelled scope. + with anyio.CancelScope(shield=True): + await cast(Any, self.body_iterator).aclose() + + class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): _diffusion_mode: bool = False _media_connector: MediaConnector | None = None @@ -451,9 +465,7 @@ def _find_tts_stage(self): """Find and return the TTS stage config, or None if not found.""" tts_stage_keys = all_tts_stage_keys() entry_stage_archs = tts_entry_stage_archs() - all_stages = frozenset( - getattr(stage.engine_args, "model_stage", None) for stage in self.engine_client.stage_configs - ) + all_stages = frozenset(stage.engine_args.model_stage for stage in self.engine_client.stage_configs) for stage in self.engine_client.stage_configs: engine_args = stage.engine_args model_stage = engine_args.model_stage @@ -1083,7 +1095,7 @@ def _validate_speech_sample_rate(self, request: OpenAICreateSpeechRequest) -> st return "sample_rate is not supported by the current TTS model" return None - def _validate_ref_audio_format(self, ref_audio: str) -> str | None: + def _validate_ref_audio_format(self, ref_audio: str | list[str] | None) -> str | None: """Validate ref_audio is a supported URI format. Returns error or None.""" if not isinstance(ref_audio, str): return "ref_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)" @@ -1362,6 +1374,57 @@ async def _resolve_ref_audio_many(self, ref_audio_list: list[str]) -> list[tuple resolved.append((wav_list, sr)) return resolved + def _observe_speech_audio_ttfp( + self, + *, + request_id: str, + result: Any, + request_arrival_ts: float | None = None, + first_packet_ts: float | None = None, + ) -> tuple[int | None, int | None, bool]: + """Emit the Speech API TTFP sample once the first PCM payload exists.""" + engine_client = getattr(self, "engine_client", None) + mod_metrics = getattr(engine_client, "mod_metrics", None) + if mod_metrics is None: + return None, None, False + + req_state = next( + ( + state + for state in getattr(engine_client, "request_states", {}).values() + if state.external_request_id == request_id + ), + None, + ) + if req_state is not None and req_state.first_audio_ts is not None: + return req_state.audio_emit_stage_id, req_state.audio_emit_replica_id, True + + arrival_ts = request_arrival_ts + if arrival_ts is None and req_state is not None: + arrival_ts = req_state.request_arrival_ts + stage_id = getattr(result, "stage_id", None) + replica_id = getattr(result, "replica_id", None) + if replica_id is None and req_state is not None and stage_id is not None: + stage_pools = getattr(getattr(engine_client, "engine", None), "stage_pools", None) + if stage_pools is not None and 0 <= stage_id < len(stage_pools): + replica_id = stage_pools[stage_id].get_bound_replica_id(req_state.request_id) + if arrival_ts is None or arrival_ts <= 0 or stage_id is None or replica_id is None: + return stage_id, replica_id, False + + now_ts = first_packet_ts if first_packet_ts is not None else time.time() + observe_audio_first_packet( + mod_metrics, + stage_id=stage_id, + replica_id=replica_id, + arrival_ts=arrival_ts, + now_ts=now_ts, + ) + if req_state is not None: + req_state.first_audio_ts = now_ts + req_state.audio_emit_stage_id = stage_id + req_state.audio_emit_replica_id = replica_id + return stage_id, replica_id, True + async def _generate_audio_chunks( self, generator, @@ -1369,6 +1432,7 @@ async def _generate_audio_chunks( response_format: str = "pcm", raw_request: Request | None = None, request_start_s: float | None = None, + request_arrival_ts: float | None = None, include_sample_rate: bool = False, usage_acc: SpeechOutputTokenCounter | None = None, tts_params: dict[str, Any] | None = None, @@ -1395,10 +1459,42 @@ async def _generate_audio_chunks( sample_rate_val = 24000 first_chunk = True first_audio_chunk_s: float | None = None + first_audio_packet_ts: float | None = None + ttfp_observed = False stream_start_s = request_start_s if request_start_s is not None else time.perf_counter() artifact_ready = False + audio_chunk_arrivals_s: list[float] = [] + audio_chunk_bytes: list[int] = [] + audio_stage_id: int | None = None + audio_replica_id: int | None = None + audio_channels = 1 source_sample_rate: int | None = None resampler: StreamingAudioResampler | None = None + last_audio_result: Any = None + + def record_stream_abort(reason: str) -> None: + mod_metrics = getattr(getattr(self, "engine_client", None), "mod_metrics", None) + if mod_metrics is not None: + mod_metrics.inc_speech_stream_aborted(reason) + + def record_audio_chunk(audio_bytes: bytes, chunk_np: np.ndarray, result: Any) -> None: + nonlocal first_audio_chunk_s, first_audio_packet_ts, ttfp_observed + nonlocal audio_stage_id, audio_replica_id, audio_channels + if not audio_bytes: + return + if first_audio_chunk_s is None: + first_audio_chunk_s = time.perf_counter() + first_audio_packet_ts = time.time() + audio_channels = _infer_audio_num_channels(np.asarray(chunk_np)) + if not ttfp_observed: + audio_stage_id, audio_replica_id, ttfp_observed = self._observe_speech_audio_ttfp( + request_id=request_id, + result=result, + request_arrival_ts=request_arrival_ts, + first_packet_ts=first_audio_packet_ts, + ) + audio_chunk_arrivals_s.append(max(time.perf_counter() - stream_start_s, 0.0)) + audio_chunk_bytes.append(len(audio_bytes)) # SSE supplies an accumulator for usage output. Raw-audio and WebSocket # streams retain terminal metrics only when their model adapter needs @@ -1419,6 +1515,7 @@ async def _generate_audio_chunks( if collect is not None and self._is_timestamps_output(res): collect["aligner_res"] = res continue + audio_output = cast(dict, audio_output) sr_raw = audio_output.get("sr") if sr_raw is not None: @@ -1452,6 +1549,9 @@ async def _generate_audio_chunks( output_sample_rate = target_sample_rate or sample_rate_val for chunk_tensor in new_chunks: + # Flush may run after a non-audio (e.g. aligner) result. + # Retain the producer of the audio buffered by the resampler. + last_audio_result = res chunk_np = ( chunk_tensor.float().detach().cpu().numpy() if hasattr(chunk_tensor, "float") else chunk_tensor ) @@ -1466,6 +1566,7 @@ async def _generate_audio_chunks( # as first audio; the post-loop guard below needs to # see an audio-less stream to fail the request. continue + wav_header: bytes | None = None # For WAV format, emit header before first audio chunk if response_format == "wav" and first_chunk: num_channels = _infer_audio_num_channels(np.asarray(chunk_np)) @@ -1474,7 +1575,6 @@ async def _generate_audio_chunks( num_channels=num_channels, bits_per_sample=16, ) - yield wav_header first_chunk = False # Convert audio to PCM bytes @@ -1485,9 +1585,10 @@ async def _generate_audio_chunks( speed=1.0, base64_encode=False, ) - if first_audio_chunk_s is None: - first_audio_chunk_s = time.perf_counter() - audio_bytes = self.create_audio(audio_obj).audio_data + audio_bytes = cast(bytes, self.create_audio(audio_obj).audio_data) + record_audio_chunk(audio_bytes, chunk_np, res) + if wav_header is not None: + yield wav_header if include_sample_rate: yield audio_bytes, output_sample_rate else: @@ -1497,8 +1598,9 @@ async def _generate_audio_chunks( final_chunk = resampler.process(np.empty((0,), dtype=np.float32), final=True) if final_chunk.size: output_sample_rate = target_sample_rate or sample_rate_val + wav_header = None if response_format == "wav" and first_chunk: - yield _create_wav_header( + wav_header = _create_wav_header( sample_rate=output_sample_rate, num_channels=1, bits_per_sample=16, @@ -1511,9 +1613,10 @@ async def _generate_audio_chunks( speed=1.0, base64_encode=False, ) - if first_audio_chunk_s is None: - first_audio_chunk_s = time.perf_counter() - audio_bytes = self.create_audio(audio_obj).audio_data + audio_bytes = cast(bytes, self.create_audio(audio_obj).audio_data) + record_audio_chunk(audio_bytes, final_chunk, last_audio_result) + if wav_header is not None: + yield wav_header if include_sample_rate: yield audio_bytes, output_sample_rate else: @@ -1527,8 +1630,21 @@ async def _generate_audio_chunks( # bytes, but they must terminate as an error rather than cleanly. if tts_params is not None and usage_acc is not None: self._validate_tts_generation(tts_params, usage_acc) + mod_metrics = getattr(getattr(self, "engine_client", None), "mod_metrics", None) + if mod_metrics is not None and audio_stage_id is not None and audio_replica_id is not None: + observe_audio_streaming_finalize( + mod_metrics, + stage_id=audio_stage_id, + replica_id=audio_replica_id, + chunk_arrival_times_s=audio_chunk_arrivals_s, + chunk_bytes=audio_chunk_bytes, + sample_rate=target_sample_rate or sample_rate_val, + channels=audio_channels, + ) self._mark_ref_audio_artifact_ready_for_request(request_id) artifact_ready = True + if mod_metrics is not None: + mod_metrics.inc_speech_stream_completed() total_ms = (time.perf_counter() - stream_start_s) * 1000.0 if first_audio_chunk_s is not None: first_chunk_ms = (first_audio_chunk_s - stream_start_s) * 1000.0 @@ -1544,7 +1660,11 @@ async def _generate_audio_chunks( request_id, total_ms, ) + except GeneratorExit: + record_stream_abort("closed") + raise except asyncio.CancelledError: + record_stream_abort("cancelled") total_ms = (time.perf_counter() - stream_start_s) * 1000.0 logger.info( "[SpeechE2E] request_id=%s stream=true status=cancelled total_ms=%.2f", @@ -1554,6 +1674,7 @@ async def _generate_audio_chunks( logger.info("Streaming request %s cancelled by client", request_id) raise except EngineDeadError as e: + record_stream_abort("engine_dead") total_ms = (time.perf_counter() - stream_start_s) * 1000.0 logger.error( "[SpeechE2E] request_id=%s stream=true status=engine_dead total_ms=%.2f", @@ -1573,6 +1694,7 @@ async def _generate_audio_chunks( ) raise except Exception as e: + record_stream_abort("error") total_ms = (time.perf_counter() - stream_start_s) * 1000.0 logger.exception( "[SpeechE2E] request_id=%s stream=true status=error total_ms=%.2f error=%s", @@ -1585,6 +1707,10 @@ async def _generate_audio_chunks( finally: if not artifact_ready: self._discard_ref_audio_artifact_warmup(request_id) + close = getattr(generator, "aclose", None) + if close is not None: + with anyio.CancelScope(shield=True): + await close() async def _generate_audio_sse_events( self, @@ -1593,6 +1719,7 @@ async def _generate_audio_sse_events( response_format: str = "pcm", raw_request: Request | None = None, request_start_s: float | None = None, + request_arrival_ts: float | None = None, request: OpenAICreateSpeechRequest | None = None, tts_params: dict[str, Any] | None = None, ): @@ -1613,24 +1740,28 @@ async def _generate_audio_sse_events( usage_acc = SpeechOutputTokenCounter() emitted_audio = False try: - async for chunk in self._generate_audio_chunks( - generator, - request_id, - response_format, - raw_request=raw_request, - request_start_s=request_start_s, - usage_acc=usage_acc, - tts_params=tts_params, - target_sample_rate=request.sample_rate if request is not None else None, - ): - payload = { - "type": "speech.audio.delta", - "audio": base64.b64encode(chunk).decode("ascii"), - "response_format": response_format, - } - data = json.dumps(payload, separators=(",", ":")) - emitted_audio = True - yield f"event: speech.audio.delta\ndata: {data}\n\n" + async with aclosing( + self._generate_audio_chunks( + generator, + request_id, + response_format, + raw_request=raw_request, + request_start_s=request_start_s, + request_arrival_ts=request_arrival_ts, + usage_acc=usage_acc, + tts_params=tts_params, + target_sample_rate=request.sample_rate if request is not None else None, + ) + ) as chunks: + async for chunk in chunks: + payload = { + "type": "speech.audio.delta", + "audio": base64.b64encode(chunk).decode("ascii"), + "response_format": response_format, + } + data = json.dumps(payload, separators=(",", ":")) + emitted_audio = True + yield f"event: speech.audio.delta\ndata: {data}\n\n" done_payload: dict[str, Any] = {"type": "speech.audio.done"} if request is not None: # Streaming path: output_tokens = sum of stage-0 deltas. @@ -1695,6 +1826,7 @@ async def _prepare_speech_generation( request: OpenAICreateSpeechRequest, request_id: str | None = None, has_inline_ref_audio: bool | None = None, + arrival_time: float | None = None, ) -> tuple[str, Any, dict[str, Any]]: if self.engine_client.errored: raise self.engine_client.dead_error @@ -1823,6 +1955,7 @@ async def _prepare_speech_generation( request_id=request_id, sampling_params_list=sampling_params_list, output_modalities=output_modalities, + arrival_time=arrival_time, ) self._track_ref_audio_artifact_warmup( request_id, @@ -1836,6 +1969,8 @@ async def _generate_pcm_chunks( generator, request_id: str, *, + request_start_s: float | None = None, + request_arrival_ts: float | None = None, include_sample_rate: bool = False, tts_params: dict[str, Any] | None = None, collect: dict | None = None, @@ -1848,28 +1983,36 @@ async def _generate_pcm_chunks( ``collect`` (when given) receives the forced-aligner stage's pooling output under ``"aligner_res"`` for downstream word-timestamp extraction. """ - async for chunk in self._generate_audio_chunks( - generator, - request_id, - response_format="pcm", - include_sample_rate=include_sample_rate, - tts_params=tts_params, - collect=collect, - target_sample_rate=target_sample_rate, - ): - yield chunk + async with aclosing( + self._generate_audio_chunks( + generator, + request_id, + response_format="pcm", + request_start_s=request_start_s, + request_arrival_ts=request_arrival_ts, + include_sample_rate=include_sample_rate, + tts_params=tts_params, + collect=collect, + target_sample_rate=target_sample_rate, + ) + ) as chunks: + async for chunk in chunks: + yield chunk async def _iter_pcm_audio_bytes(self, request: OpenAICreateSpeechRequest): """Yield raw PCM bytes for a speech request as soon as chunks are decoded.""" request_id, generator, tts_params = await self._prepare_speech_generation(request) try: - async for chunk in self._generate_pcm_chunks( - generator, - request_id, - tts_params=tts_params, - target_sample_rate=request.sample_rate, - ): - yield chunk + async with aclosing( + self._generate_pcm_chunks( + generator, + request_id, + tts_params=tts_params, + target_sample_rate=request.sample_rate, + ) + ) as chunks: + async for chunk in chunks: + yield chunk finally: self._discard_ref_audio_artifact_warmup(request_id) @@ -1881,14 +2024,19 @@ async def _generate_audio_bytes( usage_out: list[SpeechTokenUsage] | None = None, has_inline_ref_audio: bool | None = None, collect: dict | None = None, + request_arrival_ts: float | None = None, ) -> tuple[bytes | str, str]: # ``usage_out`` is an opt-in output channel: when a list is passed, the # computed SpeechTokenUsage is appended to it. The return stays a # 2-tuple so existing callers (and their test mocks) are unaffected; # batch and non-streaming response-header paths opt in when surfacing # usage outside the raw audio body. + request_arrival_ts = request_arrival_ts if request_arrival_ts is not None else time.time() request_id, generator, bytes_tts_params = await self._prepare_speech_generation( - request, request_id=request_id, has_inline_ref_audio=has_inline_ref_audio + request, + request_id=request_id, + has_inline_ref_audio=has_inline_ref_audio, + arrival_time=request_arrival_ts, ) artifact_ready = False @@ -1927,6 +2075,7 @@ async def _generate_audio_bytes( continue if step_key is None: continue + step_audio = cast(dict, step_audio) chunk = step_audio[step_key] candidates = chunk if isinstance(chunk, list) else [chunk] for cand in candidates: @@ -1948,6 +2097,7 @@ async def _generate_audio_bytes( audio_output, audio_key = self._extract_audio_output(audio_source) if audio_key is None: raise ValueError("TTS model did not produce audio output.") + audio_output = cast(dict, audio_output) # Surface forced-aligner word timestamps to the caller (set as a # response header) when requested and an aligner stage produced them. @@ -2074,7 +2224,7 @@ async def _create_diffusion_speech( request_id = f"speech-{random_uuid()}" prompt: dict[str, Any] = {"input": request.input} if request.ref_audio: - wav, sr, _ = await self._resolve_ref_audio(request.ref_audio) + wav, sr, _ = await self._resolve_ref_audio(cast(str, request.ref_audio)) prompt["ref_audio"] = (np.asarray(wav, dtype=np.float32), sr) if request.ref_text: prompt["ref_text"] = request.ref_text @@ -2151,6 +2301,7 @@ async def _create_diffusion_speech( audio_output, audio_key = self._extract_audio_output(final_output) if audio_key is None: raise ValueError("TTS model did not produce audio output.") + audio_output = cast(dict, audio_output) audio_tensor = audio_output[audio_key] sr_raw = audio_output.get("sr", 24000) @@ -2265,6 +2416,11 @@ async def create_speech( request_id = f"speech-{random_uuid()}" request_start_s = time.perf_counter() + request_arrival_ts = ( + float(getattr(raw_request.state, "request_timestamp", time.time())) + if raw_request is not None + else time.time() + ) if raw_request: raw_request.state.request_metadata = RequestResponseMetadata( request_id=request_id, @@ -2292,14 +2448,19 @@ async def create_speech( return error media_type = "audio/wav" if response_format == "wav" else "audio/pcm" - _, generator, raw_tts_params = await self._prepare_speech_generation(request, request_id=request_id) - return StreamingResponse( + _, generator, raw_tts_params = await self._prepare_speech_generation( + request, + request_id=request_id, + arrival_time=request_arrival_ts, + ) + return _SpeechStreamingResponse( self._generate_audio_chunks( generator, request_id, response_format, raw_request=raw_request, request_start_s=request_start_s, + request_arrival_ts=request_arrival_ts, tts_params=raw_tts_params, target_sample_rate=request.sample_rate, ), @@ -2314,14 +2475,19 @@ async def create_speech( if error is not None: return error - _, generator, sse_tts_params = await self._prepare_speech_generation(request, request_id=request_id) - return StreamingResponse( + _, generator, sse_tts_params = await self._prepare_speech_generation( + request, + request_id=request_id, + arrival_time=request_arrival_ts, + ) + return _SpeechStreamingResponse( self._generate_audio_sse_events( generator, request_id, response_format, raw_request=raw_request, request_start_s=request_start_s, + request_arrival_ts=request_arrival_ts, request=request, tts_params=sse_tts_params, ), @@ -2332,7 +2498,11 @@ async def create_speech( usage_box: list[SpeechTokenUsage] = [] try: audio_bytes, media_type = await self._generate_audio_bytes( - request, request_id=request_id, usage_out=usage_box, collect=collect + request, + request_id=request_id, + usage_out=usage_box, + collect=collect, + request_arrival_ts=request_arrival_ts, ) except TTSGenerationError as error: # An adapter can reject otherwise completed audio. Retry only @@ -2358,6 +2528,7 @@ async def create_speech( request_id=retry_request_id, usage_out=usage_box, collect=collect, + request_arrival_ts=request_arrival_ts, ) total_ms = (time.perf_counter() - request_start_s) * 1000.0 logger.info( diff --git a/vllm_omni/entrypoints/openai/serving_speech_stream.py b/vllm_omni/entrypoints/openai/serving_speech_stream.py index 32e9fa40d01..ceeb208c92d 100644 --- a/vllm_omni/entrypoints/openai/serving_speech_stream.py +++ b/vllm_omni/entrypoints/openai/serving_speech_stream.py @@ -55,6 +55,7 @@ import asyncio import base64 import json +import time from contextlib import aclosing from fastapi import WebSocket, WebSocketDisconnect @@ -284,6 +285,8 @@ async def _generate_and_send( ``utterance_index`` identifies the flush this sentence belongs to and ``sentence_index`` its position inside that flush. """ + request_arrival_ts = time.time() + request_start_s = time.perf_counter() response_format = config.response_format or "wav" # Reject unmet word-timestamps preconditions early with a clear reason. @@ -346,7 +349,10 @@ async def _generate_and_send( request_id = None try: if config.stream_audio: - request_id, generator, tts_params = await self._speech_service._prepare_speech_generation(request) + request_id, generator, tts_params = await self._speech_service._prepare_speech_generation( + request, + arrival_time=request_arrival_ts, + ) if config.word_timestamps: total_bytes = await self._stream_audio_with_alignments( websocket=websocket, @@ -356,6 +362,8 @@ async def _generate_and_send( utterance_index=utterance_index, sentence_index=sentence_index, language=config.language, + request_start_s=request_start_s, + request_arrival_ts=request_arrival_ts, tts_params=tts_params, ) else: @@ -363,6 +371,8 @@ async def _generate_and_send( self._speech_service._generate_pcm_chunks( generator, request_id, + request_start_s=request_start_s, + request_arrival_ts=request_arrival_ts, tts_params=tts_params, ) ) as stream: @@ -370,7 +380,10 @@ async def _generate_and_send( total_bytes += len(chunk) await websocket.send_bytes(chunk) else: - audio_bytes, _ = await self._speech_service._generate_audio_bytes(request) + audio_bytes, _ = await self._speech_service._generate_audio_bytes( + request, + request_arrival_ts=request_arrival_ts, + ) total_bytes = len(audio_bytes) await websocket.send_bytes(audio_bytes) except WebSocketDisconnect: @@ -416,6 +429,8 @@ async def _stream_audio_with_alignments( sentence_text: str, utterance_index: int, sentence_index: int, + request_start_s: float, + request_arrival_ts: float, language: str | None = None, tts_params: dict | None = None, ) -> int: @@ -463,6 +478,8 @@ async def send_chunk( self._speech_service._generate_pcm_chunks( generator, request_id, + request_start_s=request_start_s, + request_arrival_ts=request_arrival_ts, include_sample_rate=True, tts_params=tts_params, collect=collect, diff --git a/vllm_omni/metrics/definitions.py b/vllm_omni/metrics/definitions.py index 6cc2825c1ea..a09fbcc2ff2 100644 --- a/vllm_omni/metrics/definitions.py +++ b/vllm_omni/metrics/definitions.py @@ -160,6 +160,8 @@ AUDIO_UNDERRUN_S = METRIC_PREFIX + AUDIO_UNDERRUN + "_s" AUDIO_CONTINUITY_OK_METRIC = METRIC_PREFIX + AUDIO_CONTINUITY_OK AUDIO_SKIPPED_REQUESTS_METRIC = METRIC_PREFIX + AUDIO_SKIPPED_REQUESTS +SPEECH_STREAM_ABORTED_METRIC = METRIC_PREFIX + "speech_stream_aborted" +SPEECH_STREAM_COMPLETED_METRIC = METRIC_PREFIX + "speech_stream_completed" # Realtime Server VAD serving metrics. REALTIME_VAD_ACTIVE_SESSIONS = METRIC_PREFIX + "realtime_vad_active_sessions" diff --git a/vllm_omni/metrics/modality.py b/vllm_omni/metrics/modality.py index e0d51e2b0b4..bb7842e36ab 100644 --- a/vllm_omni/metrics/modality.py +++ b/vllm_omni/metrics/modality.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """OmniModalityMetrics — per-modality Prometheus families (audio path only). 7 audio business-semantic metric families. Text-path metrics (TTFT / ITL / @@ -76,6 +79,16 @@ "Silent-loss counter — code2wav rejected malformed codec input and returned 200 OK with empty audio.", labelnames=list(defs.AUDIO_SKIPPED_LABELS), ) +_speech_stream_aborted_family = Counter( + defs.SPEECH_STREAM_ABORTED_METRIC, + "Speech audio generators terminated before normal completion, including before the first PCM payload.", + labelnames=["model_name", "reason"], +) +_speech_stream_completed_family = Counter( + defs.SPEECH_STREAM_COMPLETED_METRIC, + "Speech audio generators that completed normally; does not confirm client receipt.", + labelnames=["model_name"], +) # ---------------------------------------------------------------------------- @@ -153,6 +166,17 @@ def __init__(self, model_name: str, log_stats: bool = True) -> None: # ---- Audio ------------------------------------------------------------ + def inc_speech_stream_aborted(self, reason: str) -> None: + if not self._log_stats: + return + if reason not in {"cancelled", "closed", "engine_dead", "error"}: + reason = "error" + _speech_stream_aborted_family.labels(model_name=self._model_name, reason=reason).inc() + + def inc_speech_stream_completed(self) -> None: + if self._log_stats: + _speech_stream_completed_family.labels(model_name=self._model_name).inc() + def observe_audio_ttfp(self, stage: str, replica: str, ttfp_seconds: float) -> None: if not self._log_stats: return @@ -335,6 +359,7 @@ def observe_audio_streaming_finalize( chunk_arrival_times_s: list[float], chunk_bytes: list[int], sample_rate: int, + channels: int = defs.DEFAULT_AUDIO_CHANNELS, threshold_s: float = defs.AUDIO_CONTINUITY_DEFAULT_THRESHOLD_S, ) -> None: """Emit audio_underrun_s + audio_continuity_ok_total at request end. @@ -353,6 +378,7 @@ def observe_audio_streaming_finalize( chunk_arrival_times_s=chunk_arrival_times_s, chunk_bytes=chunk_bytes, sample_rate=sample_rate, + channels=channels, threshold_s=threshold_s, ) stage_label = str(stage_id) From 78e1a0ae5ce6bec260043a085df2d35d42ce1d03 Mon Sep 17 00:00:00 2001 From: Qi Jia Date: Fri, 11 Sep 2026 22:18:07 +0800 Subject: [PATCH 78/93] [Bugfix][Model] Fix FLUX.2 Klein multi-image edit metadata (#7430) Signed-off-by: QI JIA Co-authored-by: QI JIA Co-authored-by: Cursor --- .../diffusion/test_diffusion_config_propagation.py | 13 +++++++++++++ vllm_omni/diffusion/model_metadata.py | 6 ++++++ 2 files changed, 19 insertions(+) diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py index 394f7c04ff6..23ff796e208 100644 --- a/tests/diffusion/test_diffusion_config_propagation.py +++ b/tests/diffusion/test_diffusion_config_propagation.py @@ -17,6 +17,7 @@ ) from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.model_metadata import ( + FLUX2_KLEIN_MAX_INPUT_IMAGES, HUNYUAN_IMAGE3_MAX_INPUT_IMAGES, QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES, ) @@ -196,6 +197,18 @@ def test_qwen_image_edit_plus_sets_generic_multimodal_limit(): assert od_config.max_multimodal_image_inputs == QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES +def test_flux2_klein_sets_generic_multimodal_limit(): + od_config = OmniDiffusionConfig( + model="black-forest-labs/FLUX.2-klein-9B", + model_class_name="Flux2KleinPipeline", + ) + + od_config.update_multimodal_support() + + assert od_config.supports_multimodal_inputs is True + assert od_config.max_multimodal_image_inputs == FLUX2_KLEIN_MAX_INPUT_IMAGES + + def test_task_type_roundtrip(): od = _roundtrip_diffusion_config(model="x", task_type="model-defined-task") assert od.task_type == "model-defined-task" diff --git a/vllm_omni/diffusion/model_metadata.py b/vllm_omni/diffusion/model_metadata.py index c91d2ead5ee..2ee1b1f4b0a 100644 --- a/vllm_omni/diffusion/model_metadata.py +++ b/vllm_omni/diffusion/model_metadata.py @@ -19,6 +19,8 @@ class DiffusionModelMetadata: final_output_type: str | None = None +# FLUX.2 Klein supports up to four reference images. +FLUX2_KLEIN_MAX_INPUT_IMAGES = 4 QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES = 4 # Upstream HunyuanImage-3.0 "Multi-Image Fusion" caps reference images at 3. HUNYUAN_IMAGE3_MAX_INPUT_IMAGES = 3 @@ -27,6 +29,10 @@ class DiffusionModelMetadata: _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = { + "Flux2KleinPipeline": DiffusionModelMetadata( + supports_multimodal_inputs=True, + max_multimodal_image_inputs=FLUX2_KLEIN_MAX_INPUT_IMAGES, + ), "QwenImageEditPlusPipeline": DiffusionModelMetadata( supports_multimodal_inputs=True, max_multimodal_image_inputs=QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES, From d3493384cfc6c4b9ba1fc4e534e91fe588ecb4f1 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Fri, 11 Sep 2026 10:28:55 -0400 Subject: [PATCH 79/93] [BugFix] Fix leftovers of the legacy OpenAI realtime API event names (#7426) Signed-off-by: Nick Cao Co-authored-by: Codex --- docs/serving/realtime_api.md | 6 +-- docs/serving/video_stream_api.md | 4 +- examples/online_serving/qwen3_omni/README.md | 36 ++++++++------ .../qwen3_omni/openai_realtime_client.py | 14 +++--- .../openai_api/test_duplex_handler.py | 4 +- .../test_qwen3_omni_realtime_websocket.py | 16 +++--- .../openai_api/test_serving_video_stream.py | 49 +++++++++++++++++++ .../entrypoints/openai/realtime_connection.py | 8 +-- .../openai/serving_video_stream.py | 4 +- .../entrypoints/openai/video_stream_base.py | 10 ++-- 10 files changed, 103 insertions(+), 48 deletions(-) diff --git a/docs/serving/realtime_api.md b/docs/serving/realtime_api.md index 15356414079..e44ff7ed0ae 100644 --- a/docs/serving/realtime_api.md +++ b/docs/serving/realtime_api.md @@ -43,8 +43,8 @@ event is base64-encoded raw PCM16. | Client to server | `input_audio_buffer.commit` with `final: true` | Marks the end of input | | Server to client | `transcription.delta` | Carries incremental response text | | Server to client | `transcription.done` | Carries final text and token usage | -| Server to client | `response.audio.delta` | Carries incremental PCM16 response audio | -| Server to client | `response.audio.done` | Marks the end of response audio | +| Server to client | `response.output_audio.delta` | Carries incremental PCM16 response audio | +| Server to client | `response.output_audio.done` | Marks the end of response audio | | Server to client | `error` | Reports an invalid event, model, or audio payload | A minimal client sends events in this order: @@ -63,7 +63,7 @@ engine is consuming the stream. ## Audio Handling - Input is mono PCM16 at 16 kHz for the Qwen3-Omni example. -- `response.audio.delta.audio` contains base64-encoded PCM16 bytes. +- `response.output_audio.delta.audio` contains base64-encoded PCM16 bytes. - Read `sample_rate_hz` from each audio event instead of assuming an output rate. Qwen3-Omni output is typically 24 kHz. - Concatenate audio deltas in receive order to construct the output waveform. diff --git a/docs/serving/video_stream_api.md b/docs/serving/video_stream_api.md index e1f5839915f..78575e0fc75 100644 --- a/docs/serving/video_stream_api.md +++ b/docs/serving/video_stream_api.md @@ -45,8 +45,8 @@ WebSocket /v1/video/chat/stream | Server -> Client | `response.start` | none | Query generation started. | | Server -> Client | `response.text.delta` | `delta` | Incremental text output. | | Server -> Client | `response.text.done` | `text` | Final text output for the query. | -| Server -> Client | `response.audio.delta` | `data`, `format` | Incremental generated audio, base64 WAV. | -| Server -> Client | `response.audio.done` | none | Audio output finished. | +| Server -> Client | `response.output_audio.delta` | `data`, `format` | Incremental generated audio, base64 WAV. | +| Server -> Client | `response.output_audio.done` | none | Audio output finished. | | Server -> Client | `session.done` | none | Session closed. | | Server -> Client | `error` | `message` | Recoverable protocol or generation error. | diff --git a/examples/online_serving/qwen3_omni/README.md b/examples/online_serving/qwen3_omni/README.md index ef236741202..107e1f09815 100644 --- a/examples/online_serving/qwen3_omni/README.md +++ b/examples/online_serving/qwen3_omni/README.md @@ -1,5 +1,7 @@ # Qwen3-Omni: Online serving + + ## 🛠️ Installation Please refer to [README.md](../../../README.md) @@ -19,6 +21,7 @@ Additionally, NPU, ROCm, and XPU per-platform configuration deltas are determini `platforms`: section of the corresponding YAML. To explicitly utilize a custom deployment YAML, mandate the configuration path accordingly: + ```bash vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \ --deploy-config /path/to/your_deploy_config.yaml @@ -53,7 +56,7 @@ Use the stage-based CLI when you want to run one stage per process. The example below pins Stage 0 to GPU 0 and Stage 1/2 to GPU 1 via `CUDA_VISIBLE_DEVICES`. -**1. Stage 0 (Thinker + API server)** +#### 1. Stage 0 (Thinker + API server) ```bash CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ @@ -63,7 +66,7 @@ CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ --omni-master-port 26000 ``` -**2. Stage 1 (Talker)** +#### 2. Stage 1 (Talker) ```bash CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ @@ -73,7 +76,7 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ --omni-master-port 26000 ``` -**3. Stage 2 (Code2Wav)** +#### 3. Stage 2 (Code2Wav) ```bash CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ @@ -252,11 +255,12 @@ pattern works for Qwen2.5-Omni — replace `base_config:` with the path to ### Send Multi-modal Request Get into the example folder + ```bash cd examples/online_serving/qwen3_omni ``` -#### Send request via python +#### Send request via python ```bash python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py --model Qwen/Qwen3-Omni-30B-A3B-Instruct --query-type use_image --port 8091 --host "localhost" @@ -264,7 +268,7 @@ python examples/online_serving/openai_chat_completion_client_for_multimodal_gene #### Realtime WebSocket client (`openai_realtime_client.py`) -[`openai_realtime_client.py`](./openai_realtime_client.py) connects to **`ws://:/v1/realtime`**, streams a local WAV as **PCM16 mono @ 16 kHz** in fixed-size chunks, and receives **`response.audio.*`**, **`response.output_text.*`**, and transcript events. By default it uses explicit `input_audio_buffer.commit`; pass **`--server-vad`** to let the server detect speech endpoints and commit the turn automatically. The client concatenates audio deltas and writes **`--output-wav`** (model output is typically **24 kHz**). Optional **`--delta-dump-dir`** saves each delta as `delta_000001.wav`, … for debugging. +[`openai_realtime_client.py`](./openai_realtime_client.py) connects to **`ws://:/v1/realtime`**, streams a local WAV as **PCM16 mono @ 16 kHz** in fixed-size chunks, and receives **`response.output_audio.*`**, **`response.output_text.*`**, and transcript events. By default it uses explicit `input_audio_buffer.commit`; pass **`--server-vad`** to let the server detect speech endpoints and commit the turn automatically. The client concatenates audio deltas and writes **`--output-wav`** (model output is typically **24 kHz**). Optional **`--delta-dump-dir`** saves each delta as `delta_000001.wav`, … for debugging. Streaming input works well for translation-style use cases; if the Thinker runs while input is still incomplete, consider limiting **`max_tokens`** in your session / server defaults to avoid over-generation. @@ -289,7 +293,7 @@ python openai_realtime_client.py \ **Arguments:** | Flag | Default | Description | -|------|---------|-------------| +| ------ | --------- | ------------- | | `--url` | `ws://localhost:8091/v1/realtime` | Full WebSocket URL including path | | `--model` | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | Must match the served model (sent in `session.update`) | | `--input-wav` | *(required)* | Input WAV: mono, 16-bit PCM, **16 kHz** | @@ -297,7 +301,7 @@ python openai_realtime_client.py \ | `--output-text` | *(optional)* | If set, write final transcription text to this path | | `--chunk-ms` | `200` | Size of each uploaded audio chunk (milliseconds of audio) | | `--send-delay-ms` | `0` | Delay between chunk sends (simulate realtime upload) | -| `--delta-dump-dir` | *(optional)* | Directory to write per-`response.audio.delta` WAV files | +| `--delta-dump-dir` | *(optional)* | Directory to write per-`response.output_audio.delta` WAV files | | `--server-vad` | disabled | Enable server-side Silero VAD instead of sending explicit commits | | `--num-requests` | `1` | Number of sequential sessions (see `--concurrency`) | | `--concurrency` | `1` | Max concurrent WebSocket sessions when `--num-requests` > 1 | @@ -334,7 +338,6 @@ The Python client supports the following command-line arguments: - `--prompt` (or `-p`): Custom text prompt/question. If not provided, uses default prompt for the selected query type. Example: `--prompt "What are the main activities shown in this video?"` - `--speaker`: TTS speaker/voice for audio output when requesting audio (e.g. `ethan`, `chelsie`, `aiden`). Omit to use the model default. Example: `--speaker "chelsie"` - For example, to use a local video file with custom prompt: ```bash @@ -345,22 +348,22 @@ python examples/online_serving/openai_chat_completion_client_for_multimodal_gene --prompt "What are the main activities shown in this video?" ``` -#### Send request via curl +#### Send request via curl ```bash bash run_curl_multimodal_generation.sh use_image ``` - ### FAQ ## Modality control + You can control output modalities to specify which types of output the model should generate. This is useful when you only need text output and want to skip audio generation stages for better performance. ### Supported modalities | Modalities | Output | -|------------|--------| +| ------------ | -------- | | `["text"]` | Text only | | `["audio"]` | Audio only | | `["text", "audio"]` | Text + Audio | @@ -496,7 +499,9 @@ print(response.choices[0].message.audio) Supported speaker names depend on the model (e.g. `Ethan`, `Chelsie`, `Aiden`). Omit `speaker` to use the default. ## Streaming Output + If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. We support both text streaming output and audio streaming output. Other modalities can output normally. + ```bash python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \ --query-type use_image \ @@ -521,12 +526,14 @@ The convenience script launches both the vLLM server and Gradio demo together: ``` This script will: + 1. Start the vLLM server in the background 2. Wait for the server to be ready 3. Launch the Gradio demo 4. Handle cleanup when you press Ctrl+C The script supports the following arguments: + - `--model`: Model name/path (default: Qwen/Qwen3-Omni-30B-A3B-Instruct) - `--server-port`: Port for vLLM server (default: 8091) - `--gradio-port`: Port for Gradio demo (default: 7861) @@ -537,18 +544,19 @@ The script supports the following arguments: #### Option 2: Manual Launch (Two-Step Process) -**Step 1: Launch the vLLM API server** +##### Step 1: Launch the vLLM API server ```bash vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` If you have a custom deploy config file: + ```bash vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --deploy-config /path/to/deploy_config_file ``` -**Step 2: Run the Gradio demo** +##### Step 2: Run the Gradio demo In a separate terminal: @@ -561,7 +569,7 @@ Then open `http://localhost:7861/` on your local browser to interact with the we The gradio script supports the following arguments: - `--model`: Model name/path (should match the server model) -- `--api-base`: Base URL for the vLLM API server (default: http://localhost:8091/v1) +- `--api-base`: Base URL for the vLLM API server (default: ) - `--ip`: Host/IP for Gradio server (default: 127.0.0.1) - `--port`: Port for Gradio server (default: 7861) - `--share`: Share the Gradio demo publicly (creates a public link) diff --git a/examples/online_serving/qwen3_omni/openai_realtime_client.py b/examples/online_serving/qwen3_omni/openai_realtime_client.py index 8a687158077..0a98d69a0ac 100644 --- a/examples/online_serving/qwen3_omni/openai_realtime_client.py +++ b/examples/online_serving/qwen3_omni/openai_realtime_client.py @@ -6,14 +6,14 @@ This client: 1) Reads a local WAV file (must be mono, 16-bit PCM, 16kHz), 2) Streams PCM16 chunks to /v1/realtime with OpenAI-style events, -3) Receives response.audio.*, response.output_text.*, and transcript events, +3) Receives response.output_audio.*, response.output_text.*, and transcript events, 4) Saves synthesized audio to an output WAV file and optional text file. -By default each ``response.audio.delta`` is treated as an **incremental PCM** +By default each ``response.output_audio.delta`` is treated as an **incremental PCM** chunk and all chunks are concatenated into the final ``--output-wav``. Optional debugging: pass ``--delta-dump-dir DIR`` to write every -``response.audio.delta`` payload as ``delta_000001.wav``, ``delta_000002.wav``, … +``response.output_audio.delta`` payload as ``delta_000001.wav``, ``delta_000002.wav``, … Usage: python openai_realtime_client.py \ @@ -187,7 +187,7 @@ async def run_client( if event_type == "session.created": continue - if event_type == "response.audio.delta": + if event_type == "response.output_audio.delta": sr = event.get("sample_rate_hz") if isinstance(sr, int) and sr > 0: output_sample_rate = sr @@ -223,11 +223,11 @@ async def run_client( print(f"{log_prefix}text usage: {usage}") continue - if event_type == "response.audio_transcript.done": + if event_type == "response.output_audio_transcript.done": final_audio_transcript = event.get("transcript", "") continue - if event_type == "response.audio.done": + if event_type == "response.output_audio.done": if not server_vad: break continue @@ -339,7 +339,7 @@ def main() -> None: "--delta-dump-dir", type=Path, default=None, - help="If set, each response.audio.delta is saved as delta_NNNNNN.wav under this directory", + help="If set, each response.output_audio.delta is saved as delta_NNNNNN.wav under this directory", ) parser.add_argument( "--server-vad", diff --git a/tests/entrypoints/openai_api/test_duplex_handler.py b/tests/entrypoints/openai_api/test_duplex_handler.py index 0f969845792..206a38917ab 100644 --- a/tests/entrypoints/openai_api/test_duplex_handler.py +++ b/tests/entrypoints/openai_api/test_duplex_handler.py @@ -2197,9 +2197,7 @@ def test_native_realtime_protocol_audio_delta_preserves_sample_rate_hz(): ) audio_events = [ - payload - for payload in payloads - if payload["type"] in {"response.output_audio.delta", "response.output_audio.delta"} + payload for payload in payloads if payload["type"] in {"response.output_audio.delta", "response.audio.delta"} ] assert {payload["type"] for payload in audio_events} == {"response.output_audio.delta"} assert {payload["format"] for payload in audio_events} == {"pcm16"} diff --git a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py index 406668a4ce9..be83ea9dbde 100644 --- a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py +++ b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py @@ -175,7 +175,7 @@ async def _run_realtime_audio_roundtrip( if event_type == "session.created": continue - if event_type == "response.audio.delta": + if event_type == "response.output_audio.delta": delta_events += 1 sr = event.get("sample_rate_hz") if isinstance(sr, int) and sr > 0: @@ -195,7 +195,7 @@ async def _run_realtime_audio_roundtrip( final_text = event.get("text", "") or "".join(text_chunks) continue - if event_type == "response.audio.done": + if event_type == "response.output_audio.done": break if event_type == "error": @@ -315,7 +315,7 @@ def _server_vad_pcm16_input() -> bytes: def _assert_realtime_smoke(result: dict) -> None: out_pcm = result["output_pcm"] assert result["delta_events"] >= 1 - assert out_pcm, "No output PCM from response.audio.delta" + assert out_pcm, "No output PCM from response.output_audio.delta" assert len(out_pcm) % 2 == 0 assert len(out_pcm) >= 4096, "Output audio unexpectedly small" assert result["output_sample_rate"] > 0 @@ -388,7 +388,7 @@ def test_streaming_audio_input_pcm_output_async_chunk(self, omni_server) -> None @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", realtime_async_chunk_server_params, indirect=True) def test_long_audio_response_completes_async_chunk(self, omni_server) -> None: - """Regression for #6474: long async-chunk audio must emit response.audio.done.""" + """Regression for #6474: long async-chunk audio must emit response.output_audio.done.""" issue_6474_pcm16 = _synthetic_pcm16_input( phrase_text=ISSUE_6474_SYNTH_PHRASE_TEXT, ) @@ -450,8 +450,8 @@ def test_server_vad_multi_turn_without_client_commit( "input_audio_buffer.speech_stopped", "input_audio_buffer.committed", "response.created", - "response.audio.delta", - "response.audio.done", + "response.output_audio.delta", + "response.output_audio.done", "response.done", ] input_item_ids: list[str] = [] @@ -459,7 +459,7 @@ def test_server_vad_multi_turn_without_client_commit( for events in turns: event_types = [event["type"] for event in events] for event_type in required_sequence: - if event_type != "response.audio.delta": + if event_type != "response.output_audio.delta": assert event_types.count(event_type) == 1, event_types positions = [event_types.index(event_type) for event_type in required_sequence] assert positions == sorted(positions) @@ -486,7 +486,7 @@ def test_server_vad_multi_turn_without_client_commit( output_pcm = b"".join( base64.b64decode(event["delta"]) for event in events - if event["type"] == "response.audio.delta" and event.get("delta") + if event["type"] == "response.output_audio.delta" and event.get("delta") ) assert output_pcm assert created["id"] == done["id"] diff --git a/tests/entrypoints/openai_api/test_serving_video_stream.py b/tests/entrypoints/openai_api/test_serving_video_stream.py index b534d6c6706..58ea2b57f9d 100644 --- a/tests/entrypoints/openai_api/test_serving_video_stream.py +++ b/tests/entrypoints/openai_api/test_serving_video_stream.py @@ -14,6 +14,7 @@ from typing import Any import pytest +import torch from PIL import Image from vllm_omni.entrypoints.openai import video_stream_base, video_stream_envs @@ -457,6 +458,54 @@ async def _preprocess_to_engine_prompt(self, request): assert not any(m.get("type") == "response.text.delta" for m in ws_off.sent) +@pytest.mark.asyncio +@pytest.mark.parametrize("async_chunk_mode", ["on", "off"]) +async def test_audio_events_use_output_audio_names(monkeypatch, async_chunk_mode): + class AudioEngine: + def generate(self, **_kwargs): + async def _gen(): + yield _audio_result([torch.zeros(1)]) + + return _gen() + + class AudioHandler(QwenOmniStreamingVideoHandler): + async def _preprocess_to_engine_prompt(self, request): + return {"prompt": "x"} + + monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", async_chunk_mode) + monkeypatch.setattr( + OmniStreamingVideoHandler, + "_encode_audio_wav_b64", + staticmethod(lambda _audio: "audio-b64"), + ) + + ws = MockWebSocket() + handler = AudioHandler(chat_service=object(), engine_client=AudioEngine()) + config = StreamingVideoSessionConfig(model="test", modalities=["text", "audio"]) + + await handler._process_query_engine( + ws, + config, + [_b64(_make_jpeg())], + bytearray(), + [], + "describe", + "req-audio-events", + asyncio.Event(), + {}, + ) + + audio_events = [message for message in ws.sent if "audio" in message.get("type", "")] + assert audio_events == [ + { + "type": "response.output_audio.delta", + "data": "audio-b64", + "format": "wav", + }, + {"type": "response.output_audio.done"}, + ] + + @pytest.mark.asyncio async def test_query_without_engine_client_sends_error(): ws = MockWebSocket() diff --git a/vllm_omni/entrypoints/openai/realtime_connection.py b/vllm_omni/entrypoints/openai/realtime_connection.py index 86a8ba0319d..3d37fbe37fc 100644 --- a/vllm_omni/entrypoints/openai/realtime_connection.py +++ b/vllm_omni/entrypoints/openai/realtime_connection.py @@ -176,7 +176,7 @@ async def _run_generation( sent_audio = True await self.send_json( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "audio": self._pcm16_b64(chunk), "format": "pcm16", "sample_rate_hz": sample_rate, @@ -194,7 +194,7 @@ async def _run_generation( await self.send(TranscriptionDone(text=full_text, usage=usage)) if sent_audio: - await self.send_json({"type": "response.audio.done", "has_audio": True}) + await self.send_json({"type": "response.output_audio.done", "has_audio": True}) audio_done_sent = True except Exception as e: logger.exception("Error in generation: %s", e) @@ -213,9 +213,9 @@ async def _run_generation( # Always send terminal event so clients don't hang forever. if self._is_connected and not audio_done_sent: try: - await self.send_json({"type": "response.audio.done", "has_audio": sent_audio}) + await self.send_json({"type": "response.output_audio.done", "has_audio": sent_audio}) except Exception: - logger.exception("Failed to send response.audio.done") + logger.exception("Failed to send response.output_audio.done") while not self.audio_queue.empty(): self.audio_queue.get_nowait() diff --git a/vllm_omni/entrypoints/openai/serving_video_stream.py b/vllm_omni/entrypoints/openai/serving_video_stream.py index e05a40eeffa..5b30ff0faf9 100644 --- a/vllm_omni/entrypoints/openai/serving_video_stream.py +++ b/vllm_omni/entrypoints/openai/serving_video_stream.py @@ -18,8 +18,8 @@ {"type": "response.start"} {"type": "response.text.delta", "delta": "..."} {"type": "response.text.done", "text": "..."} - {"type": "response.audio.delta", "data": "...", "format": "wav"} - {"type": "response.audio.done"} + {"type": "response.output_audio.delta", "data": "...", "format": "wav"} + {"type": "response.output_audio.done"} {"type": "session.done"} {"type": "error", "message": "..."} """ diff --git a/vllm_omni/entrypoints/openai/video_stream_base.py b/vllm_omni/entrypoints/openai/video_stream_base.py index ab5d51c33a5..977a0b06526 100644 --- a/vllm_omni/entrypoints/openai/video_stream_base.py +++ b/vllm_omni/entrypoints/openai/video_stream_base.py @@ -21,8 +21,8 @@ {"type": "response.start"} {"type": "response.text.delta", "delta": "..."} {"type": "response.text.done", "text": "..."} - {"type": "response.audio.delta", "data": "...", "format": "wav"} - {"type": "response.audio.done"} + {"type": "response.output_audio.delta", "data": "...", "format": "wav"} + {"type": "response.output_audio.done"} {"type": "session.done"} {"type": "error", "message": "..."} """ @@ -763,7 +763,7 @@ async def _process_query_engine( if b64: await websocket.send_json( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "data": b64, "format": "wav", } @@ -807,7 +807,7 @@ async def _process_query_engine( if b64: await websocket.send_json( { - "type": "response.audio.delta", + "type": "response.output_audio.delta", "data": b64, "format": "wav", } @@ -816,7 +816,7 @@ async def _process_query_engine( logger.exception("Failed to coalesce off-path audio") if audio_chunk_count > 0: - await websocket.send_json({"type": "response.audio.done"}) + await websocket.send_json({"type": "response.output_audio.done"}) response_text = "".join(text_parts) self.on_turn_complete(message_history, user_message, response_text) From c111f3a2d6048848cd239b9cbce37cd55001d57c Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Fri, 11 Sep 2026 13:09:58 -0700 Subject: [PATCH 80/93] [Model] Add Tencent AuK speech generation and editing (encoder + diffusion pipeline) (#7385) Signed-off-by: Yueqian Lin Co-authored-by: Sy03 <1370724210@qq.com> --- docs/models/supported_models.md | 1 + .../text_to_speech/README.md | 1 + recipes/README.md | 1 + recipes/Tencent/AuK-H100.md | 167 +++++ .../models/auk/test_auk_transformer_parity.py | 439 ++++++++++++ .../models/auk/test_auk_vae_parity.py | 254 +++++++ .../diffusion/models/auk/test_pipeline_auk.py | 514 +++++++++++++ tests/e2e/offline_inference/test_auk.py | 121 ++++ .../stage_input_processors/test_auk.py | 188 +++++ tests/model_extras/test_auk_requests.py | 157 ++++ tests/model_tests/diffusion/test_alignment.py | 4 + tools/auk_parity_reference.py | 227 ++++++ tools/prepare_auk_checkpoint.py | 369 ++++++++++ vllm_omni/config/pipeline_registry.py | 2 + vllm_omni/deploy/auk.yaml | 51 ++ vllm_omni/diffusion/models/auk/__init__.py | 16 + .../diffusion/models/auk/auk_transformer.py | 677 ++++++++++++++++++ vllm_omni/diffusion/models/auk/auk_vae.py | 505 +++++++++++++ .../diffusion/models/auk/pipeline_auk.py | 486 +++++++++++++ vllm_omni/diffusion/registry.py | 6 + .../model_executor/models/auk/__init__.py | 5 + vllm_omni/model_executor/models/auk/auk.py | 304 ++++++++ .../model_executor/models/auk/pipeline.py | 60 ++ vllm_omni/model_executor/models/registry.py | 5 + .../stage_input_processors/auk.py | 148 ++++ vllm_omni/model_extras/auk.py | 145 ++++ .../transformers_utils/configs/__init__.py | 3 + vllm_omni/transformers_utils/configs/auk.py | 130 ++++ 28 files changed, 4986 insertions(+) create mode 100644 recipes/Tencent/AuK-H100.md create mode 100644 tests/diffusion/models/auk/test_auk_transformer_parity.py create mode 100644 tests/diffusion/models/auk/test_auk_vae_parity.py create mode 100644 tests/diffusion/models/auk/test_pipeline_auk.py create mode 100644 tests/e2e/offline_inference/test_auk.py create mode 100644 tests/model_executor/stage_input_processors/test_auk.py create mode 100644 tests/model_extras/test_auk_requests.py create mode 100644 tools/auk_parity_reference.py create mode 100755 tools/prepare_auk_checkpoint.py create mode 100644 vllm_omni/deploy/auk.yaml create mode 100644 vllm_omni/diffusion/models/auk/__init__.py create mode 100644 vllm_omni/diffusion/models/auk/auk_transformer.py create mode 100644 vllm_omni/diffusion/models/auk/auk_vae.py create mode 100644 vllm_omni/diffusion/models/auk/pipeline_auk.py create mode 100644 vllm_omni/model_executor/models/auk/__init__.py create mode 100644 vllm_omni/model_executor/models/auk/auk.py create mode 100644 vllm_omni/model_executor/models/auk/pipeline.py create mode 100644 vllm_omni/model_executor/stage_input_processors/auk.py create mode 100644 vllm_omni/model_extras/auk.py create mode 100644 vllm_omni/transformers_utils/configs/auk.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 6b0952eac47..1f4093828ea 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -68,6 +68,7 @@ th { | `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | — | | `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — | +| `AuKForConditionalGeneration` | AuK / AuK-Flash (instruction-driven speech generation and editing) | `tencent/AuK`, `tencent/AuK-Flash` (with `Qwen/Qwen2.5-Omni-3B`; assemble with `tools/prepare_auk_checkpoint.py`) | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Tencent/AuK-H100.md) | | `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | ✅︎ | — | | `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | ✅︎ | — | | `DotsTTSForConditionalGeneration` | dots.tts | `dots-studio/dots.tts-soar` | ✅︎ | | | | — | diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index a8a273026a5..89ff2dbce5e 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -14,6 +14,7 @@ list of supported architectures across all modalities, see | Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate | |---|---|---|---|---|---|---| +| AuK | `tencent/AuK`, `tencent/AuK-Flash` (assembled with `tools/prepare_auk_checkpoint.py`) | 2 (thinker encoder + DiT) | ✓ | — | instruction-driven editing, enhancement, separation; no example script (use `vllm_omni.model_extras.auk`) | 24 kHz | | Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | dual-AR | ✓ | ✓ | 11 languages | 44.1 kHz | | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz | | Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz | diff --git a/recipes/README.md b/recipes/README.md index ea45762014c..d9f4b68222c 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -84,6 +84,7 @@ recipes/ | [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x or 2x H200 / 1x AMD MI300X | | [`SenseNova/SenseNova-U1.5.md`](./SenseNova/SenseNova-U1.5.md) | Unified image generation and understanding, with an 8-step distilled LoRA | 1x A800 80GB | | [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB / 1x AMD MI300X | +| [`Tencent/AuK-H100.md`](./Tencent/AuK-H100.md) | Offline instruction-driven TTS, speech editing, enhancement and separation (AuK base and Flash) | 1x H100 80GB | | [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB | | [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 | | [`Wan-AI/Wan2.2-T2V.md`](./Wan-AI/Wan2.2-T2V.md) | Text-to-video serving (Wan2.2 14B) with Skip-Softmax sparse attention | 1x datacenter Blackwell (B200/B300) | diff --git a/recipes/Tencent/AuK-H100.md b/recipes/Tencent/AuK-H100.md new file mode 100644 index 00000000000..5acd01cb05e --- /dev/null +++ b/recipes/Tencent/AuK-H100.md @@ -0,0 +1,167 @@ +# AuK speech generation and editing — H100 + +> Instruction-driven zero-shot TTS, instruct TTS, speech editing, enhancement +> and separation with Tencent's AuK, served as a two-stage vLLM-Omni pipeline +> (Qwen2.5-Omni thinker encoder, then a rectified-flow DiT with a BigVGAN-flow +> VAE) on one H100. + +## Summary + +- Vendor: Tencent Hunyuan +- Model: `tencent/AuK` (32-step base) and `tencent/AuK-Flash` (distilled 4-step), + both with `Qwen/Qwen2.5-Omni-3B` as the frozen encoder +- Task: zero-shot TTS, instruct TTS, content and acoustic editing, + paralinguistic editing, speech enhancement, speaker and music separation +- Mode: offline inference (online chat serving is not yet qualified) +- Hardware: 1x H100 80GB +- Maintainer: Community + +## When to use this recipe + +Use this recipe to run AuK through vLLM-Omni for any of its instruction-driven +tasks with a single shared request shape: an instruction, an optional source or +reference clip, and a target duration. The task is chosen by the instruction +text, exactly as in the upstream repository's cookbook. + +## Supported model contract + +| Task family | Instruction | Audio input | Duration | +| --- | --- | --- | --- | +| Zero-shot TTS | "Say the following with the same voice: '...'" | reference clip | `gen_seconds` required | +| Instruct TTS | `Generate speech based on the following description: "{voice description}". The content to speak is: "{text}".` (upstream template; other orderings can make the model speak the description) | none | `gen_seconds` required | +| Content / lyric editing | "Replace '...' with '...'" and similar | source clip | defaults to the source length | +| Acoustic / paralinguistic editing | pitch, speed, volume, emotion, timbre, de-accent, nonverbal, whisper | source clip | defaults to the source length (speed edits scale it) | +| Enhancement / separation | "Keep pure speech voice, remove noise and reverberation." and similar | source clip | defaults to the source length | + +| Property | Value | +| --- | --- | +| Input audio | mono or stereo (averaged), any sample rate (resampled to 24 kHz for the VAE and to 16 kHz for the encoder), one clip per request | +| Output | mono float32 waveform at 24 kHz, `ceil(gen_seconds * 50)` latent frames | +| Sequence limit | reference plus target latents up to 65536 frames (about 21 minutes) | +| Sampling knobs | `num_inference_steps` (default 32), `guidance_scale` (default 2.0), `seed`; Flash pins 4 steps, CFG 0 and its own time grid | +| Request knobs | `gen_seconds`, `sway` (default -1.0), `t_grid`, `vae_sample` via `additional_information["auk"]` | +| Concurrency | one request per DiT forward; the encoder stage runs with `max_num_seqs: 1`, prefix caching and chunked prefill off | +| Streaming | none: the ODE runs over the whole target | + +## References + +- Upstream repository and cookbook: +- Technical report: +- Weights: , , + +- Pipeline topology: `vllm_omni/model_executor/models/auk/pipeline.py`; + deploy defaults: `vllm_omni/deploy/auk.yaml` +- Request helpers: `vllm_omni/model_extras/auk.py` +- Test: `tests/e2e/offline_inference/test_auk.py` + +## Hardware + +- Accelerator model and per-device memory: NVIDIA H100 80GB +- Number of devices: 1 (both stages share the device) +- Device interconnect: not applicable +- Host memory: 64 GB is ample; the assembled directory is symlinks only +- Qualification scope: offline `Omni` inference for the four cookbook tasks + below, base and Flash checkpoints, seed-reproducible, four concurrent + requests bit-identical to sequential runs + +## Software environment + +- OS: Linux +- Python: 3.12 +- Driver / runtime: CUDA 13 driver with the vLLM 0.29.0 wheel (torch 2.13) +- vLLM version: 0.29.0 +- vLLM-Omni version or commit: branch `auk-support` on top of `aff7d6494` + +## Command + +Assemble the checkpoint directory once (symlinks by default): + +```bash +hf download tencent/AuK --local-dir ckpts/AuK +hf download Qwen/Qwen2.5-Omni-3B --local-dir ckpts/Qwen2.5-Omni-3B +python tools/prepare_auk_checkpoint.py --auk-dir ckpts/AuK --qwen-dir ckpts/Qwen2.5-Omni-3B --out ckpts/auk-omni +# For AuK-Flash: --auk-dir ckpts/AuK-Flash --out ckpts/auk-omni-flash +``` + +Offline inference (any task; the instruction selects it): + +```python +import soundfile as sf +from vllm.multimodal.media.audio import load_audio + +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.model_extras.auk import auk_prompt, auk_sampling_params + +omni = Omni(model="ckpts/auk-omni") # loads vllm_omni/deploy/auk.yaml +wav, sr = load_audio("ref.wav", sr=None) +prompt = auk_prompt( + "Say the following with the same voice: 'Ladies and gentlemen, welcome.'", + (wav, sr), + gen_seconds=4.0, +) +out = omni.generate(prompt, auk_sampling_params(nfe=32, cfg=2.0, seed=0))[0] +sf.write("out.wav", out.multimodal_output["audio"], out.multimodal_output["audio_sample_rate"]) +``` + +Text-only instruct TTS passes `None` for the audio and must set `gen_seconds`. +Editing and enhancement requests pass the source clip and may omit +`gen_seconds` to keep the source length. + +## Verification + +```bash +VLLM_OMNI_AUK_MODEL_DIR=ckpts/auk-omni python -m pytest tests/e2e/offline_inference/test_auk.py -q +``` + +The parity suites compare the port against the upstream implementation and +need reference artifacts produced by the upstream package (install it with +`pip install -e /path/to/AuK` in its own environment; it pins its own torch): + +```bash +python tools/auk_parity_reference.py --auk-repo /path/to/AuK --ckpt-dir ckpts/AuK \ + --qwen-dir ckpts/Qwen2.5-Omni-3B --out auk-parity +AUK_CKPT_DIR=ckpts/AuK AUK_REF_WAV=/path/to/AuK/assets/demo-input-audio/zero-shot-tts/ref.wav \ + python -m pytest tests/diffusion/models/auk/test_auk_transformer_parity.py tests/diffusion/models/auk/test_auk_vae_parity.py -q +AUK_OMNI_CKPT_DIR=ckpts/auk-omni AUK_PARITY_REF=auk-parity/parity_ref/base \ + python -m pytest tests/diffusion/models/auk/test_pipeline_auk.py -q +``` + +Measured on one H100 against the upstream implementation with the same seeds +and mean VAE latents (four cookbook cases, base checkpoint): identical +transcripts, speaker similarity to the upstream output 0.993 to 0.999, log-mel +L1 0.11 to 0.47 (the upstream VAE-resampling noise floor is 0.39). Flash: +0.989 to 0.999 and 0.24 to 0.51. Wall per request on the second call: base +1.0 to 2.0 s for 3.5 to 10.9 s of audio, Flash 0.17 to 0.30 s; engine start +about 47 s with a warm page cache. + +## Notes + +- Memory usage: about 25 GB peak on the device for a 12 s generation with + both stages resident (deploy defaults: encoder 0.45, diffusion stage 0.35 + of device memory). +- Key flags: `enforce_eager` on both stages (the encoder walks the decoder + layers itself for the layer fusion). `enable_prefix_caching` must stay off + for the encoder: a cache hit skips prompt positions that the fused + condition needs. `enable_chunked_prefill` is off by default: forcing it + (128-token chunks, so two to three chunks per prompt) reproduces the + unchunked condition to bf16 rounding (per-token cosine 0.99999) and leaves + request walls unchanged, because the encoder is about 3% of a base request, + while concurrent runs stop being bit-identical to sequential ones. +- Known limitations: the encoder's audio tower runs without `flash_attn` in a + plain vLLM install, which shifts the encoder output on audio token positions + (per-token cosine 0.96 vs the upstream fp32 fusion; text positions 0.9999) + with waveform-level effects at the VAE noise floor; installing `flash-attn` + removes the gap. No streaming. Digits are read the way upstream reads them + (spell numbers out in the text). Instruct TTS without a reference can leak + the description into the speech, as upstream does. + +## Supported features + +| Feature | Status | Guide | +| --- | --- | --- | +| Offline `Omni.generate` | supported (base and Flash) | `docs/getting_started/quickstart.md` | +| Online `/v1/chat/completions` with audio | not yet qualified | `docs/serving/` | +| `/v1/audio/speech` | not supported (needs a TTS adapter) | `docs/contributing/model/adding_tts_model.md` | +| Streaming / async chunk | not supported | `docs/design/feature/async_chunk.md` | +| Batching across requests | one request per DiT forward | `docs/user_guide/diffusion/` | +| Tensor / sequence parallelism | not supported | `docs/configuration/composable_parallel.md` | diff --git a/tests/diffusion/models/auk/test_auk_transformer_parity.py b/tests/diffusion/models/auk/test_auk_transformer_parity.py new file mode 100644 index 00000000000..937369d054b --- /dev/null +++ b/tests/diffusion/models/auk/test_auk_transformer_parity.py @@ -0,0 +1,439 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""AuK transformer parity: the correctness oracle for the in-tree port. + +Checks that :class:`AuKTransformer` reproduces the upstream ``auk`` reference backbone +on the released ``auk_base`` weights. Both models are built from the same +state dict and fed the same inputs, under the reference's own inference +configuration: fp32 weights with ``torch.autocast(bfloat16)``. That +configuration matters. The reference keeps its rotary frequencies in a buffer, +so casting its weights to bf16 would round both the frequencies and the +integer positions and move the numbers for reasons unrelated to the port. + +Cases: + +* ``single`` - one conditional forward +* ``cfg`` - the doubled ``cfg_infer`` forward, both branches +* ``masked`` - batch of 2 with padded targets, references and text +* ``rope`` - the local rotary implementation against ``x_transformers``, and + its survival of a cast to bf16 +* ``sample`` - an 8-step Euler trajectory against the reference backbone + driven by a transcription of upstream's ODE loop +* ``generator`` - the sampler's two noise paths agreeing +* ``cache`` - the text projection cache and ``clear_cache()`` +* ``mask flag`` - ``attn_mask_enabled=False`` reproducing the reference with + its own flag off + +This directory deliberately has no ``__init__.py``, unlike its siblings. +pytest's prepend import mode puts the nearest non-package directory on +``sys.path``, and an ``auk`` package here would shadow the upstream ``auk`` +distribution this test compares against. + +Needs the released checkpoint and a GPU, so it runs locally rather than in +CI. Point ``AUK_CKPT_DIR`` at a directory holding ``auk_base.safetensors`` and +``config.yaml``; the test skips when that is unset or when the upstream +``auk`` package is not importable:: + + AUK_CKPT_DIR=/path/to/ckpts/AuK CUDA_VISIBLE_DEVICES=6 \ + python -m pytest tests/diffusion/models/auk/test_auk_transformer_parity.py -v -s +""" + +from __future__ import annotations + +import importlib.util +import math +import os +from pathlib import Path + +import pytest +import torch +import yaml + +from vllm_omni.diffusion.models.auk.auk_transformer import ( + AuKTransformer, + Rotary, + _apply_rope, + dit_state_dict, + sample_latents, +) + +CKPT_DIR = os.environ.get("AUK_CKPT_DIR") + +# Tolerances. The port evaluates rotary embeddings in fp32 and skips the +# reference's redundant uncond text projection, so the two graphs are not +# bit-identical under autocast; they agree well inside bf16 noise. +FORWARD_ATOL = 3e-2 +FORWARD_COS = 0.999 +SAMPLE_ATOL = 5e-2 +SAMPLE_COS = 0.999 +ROPE_ATOL = 1e-5 + +TARGET_FRAMES = 300 +REF_FRAMES = 150 +TEXT_TOKENS = 40 +SHORT_LENS = (220, 90, 25) # target, reference, text for the padded batch row +TIME = 0.37 +SAMPLE_STEPS = 8 +SWAY_COEF = -1.0 +SEED = 1234 +DEVICE = "cuda:0" +AUTOCAST_DTYPE = torch.bfloat16 + + +def _skip_reason() -> str | None: + if CKPT_DIR is None: + return "AUK_CKPT_DIR is not set" + if not Path(CKPT_DIR, "auk_base.safetensors").is_file(): + return f"no auk_base.safetensors under {CKPT_DIR}" + try: + # Resolving the submodule, not just "auk", also catches the shadow an + # __init__.py in this directory would create. + upstream = importlib.util.find_spec("auk.model") + except ModuleNotFoundError: + upstream = None + if upstream is None: + return "the upstream 'auk' package is not importable or is shadowed by this directory" + if not torch.cuda.is_available(): + return "no CUDA device" + return None + + +SKIP_REASON = _skip_reason() + +pytestmark = [ + pytest.mark.local_model, + pytest.mark.diffusion, + pytest.mark.gpu, + pytest.mark.cuda, + pytest.mark.skipif(SKIP_REASON is not None, reason=str(SKIP_REASON)), +] + + +def _arch() -> dict: + """Read the backbone geometry out of the checkpoint's config.""" + config = yaml.safe_load(Path(CKPT_DIR, "config.yaml").read_text())["model"] + return { + "dim": config["arch"]["dim"], + "heads": config["arch"]["heads"], + "dim_head": config["arch"].get("dim_head", 64), + "ff_mult": config["arch"]["ff_mult"], + "text_hidden_dim": config["arch"]["text_hidden_dim"], + "num_layers": config["arch"]["num_layers"], + "num_single_layers": config["arch"]["num_single_layers"], + "latent_dim": config["vae"]["latent_dim"], + } + + +def _stats(port: torch.Tensor, ref: torch.Tensor) -> tuple[float, float, float]: + """Return max abs diff, cosine similarity and the reference's RMS.""" + a = port.flatten().float() + b = ref.flatten().float() + return ( + (a - b).abs().max().item(), + torch.nn.functional.cosine_similarity(a, b, dim=0).item(), + b.pow(2).mean().sqrt().item(), + ) + + +def _assert_close(name: str, port: torch.Tensor, ref: torch.Tensor, atol: float, min_cos: float) -> None: + max_abs, cos, rms = _stats(port, ref) + print(f"[{name}] max|Δ|={max_abs:.4e} cos={cos:.8f} ref_rms={rms:.4f} rel={max_abs / rms:.3e}") + assert max_abs < atol, f"{name}: max|Δ|={max_abs:.4e} >= {atol:.1e} (reference rms {rms:.4f})" + assert cos > min_cos, f"{name}: cos={cos:.8f} <= {min_cos}" + + +def _autocast(): + return torch.autocast("cuda", dtype=AUTOCAST_DTYPE) + + +@pytest.fixture(scope="module") +def models() -> tuple[torch.nn.Module, AuKTransformer]: + """The reference and the port, same weights, fp32 on the accelerator.""" + from auk.model.flux2_edit import Flux2Edit + from safetensors.torch import load_file + + arch = _arch() + device = torch.device(DEVICE) + state = dit_state_dict(load_file(str(Path(CKPT_DIR, "auk_base.safetensors")), device="cpu")) + + reference = Flux2Edit(dropout=0.0, attn_backend="torch", attn_mask_enabled=True, **arch) + reference.load_state_dict(state, strict=True) + reference = reference.to(device).eval().requires_grad_(False) + + port = AuKTransformer(**arch) + # strict=True is the point of the port: every released tensor lands. + port.load_state_dict(state, strict=True) + port = port.to(device).eval().requires_grad_(False) + + del state + return reference, port + + +@pytest.fixture(scope="module") +def inputs() -> dict[str, torch.Tensor]: + """Fixed random inputs, as a batch of 2 whose second row is padded.""" + device = torch.device(DEVICE) + generator = torch.Generator(device="cpu").manual_seed(SEED) + arch = _arch() + + def randn(*shape: int) -> torch.Tensor: + return torch.randn(*shape, generator=generator, dtype=torch.float32).to(device) + + def pad_mask(width: int, short: int) -> torch.Tensor: + mask = torch.zeros(2, width, dtype=torch.bool, device=device) + mask[0, :] = True + mask[1, :short] = True + return mask + + mask2 = pad_mask(TARGET_FRAMES, SHORT_LENS[0]) + ref_mask2 = pad_mask(REF_FRAMES, SHORT_LENS[1]) + c_mask2 = pad_mask(TEXT_TOKENS, SHORT_LENS[2]) + + # Padding rows are all-zero in the text encoder's output, so keep the + # inputs consistent with the masks. + x2 = randn(2, TARGET_FRAMES, arch["latent_dim"]) * mask2.unsqueeze(-1) + ref2 = randn(2, REF_FRAMES, arch["latent_dim"]) * ref_mask2.unsqueeze(-1) + text2 = randn(2, TEXT_TOKENS, arch["text_hidden_dim"]) * c_mask2.unsqueeze(-1) + + return { + "x": x2[:1].contiguous(), + "ref": ref2[:1].contiguous(), + "text": text2[:1].contiguous(), + "c_mask": c_mask2[:1].contiguous(), + "x2": x2, + "ref2": ref2, + "text2": text2, + "mask2": mask2, + "ref_mask2": ref_mask2, + "c_mask2": c_mask2, + "time": torch.tensor(TIME, dtype=torch.float32, device=device), + } + + +def test_single_forward(models, inputs) -> None: + """One conditional forward: reference against port.""" + reference, port = models + kwargs = dict(text=inputs["text"], time=inputs["time"], c_mask=inputs["c_mask"], ref=inputs["ref"]) + with torch.no_grad(), _autocast(): + want = reference(x=inputs["x"], **kwargs) + got = port(inputs["x"], **kwargs) + + assert got.shape == want.shape == (1, TARGET_FRAMES, port.latent_dim) + _assert_close("single", got, want, FORWARD_ATOL, FORWARD_COS) + + +def test_cfg_forward(models, inputs) -> None: + """The doubled cfg_infer forward, cond branch first.""" + reference, port = models + kwargs = dict( + text=inputs["text"], + time=inputs["time"], + c_mask=inputs["c_mask"], + ref=inputs["ref"], + cfg_infer=True, + ) + with torch.no_grad(), _autocast(): + want = reference(x=inputs["x"], **kwargs) + got = port(inputs["x"], **kwargs) + + assert got.shape == want.shape == (2, TARGET_FRAMES, port.latent_dim) + _assert_close("cfg-cond", got[:1], want[:1], FORWARD_ATOL, FORWARD_COS) + _assert_close("cfg-uncond", got[1:], want[1:], FORWARD_ATOL, FORWARD_COS) + + +def test_masked_batch_forward(models, inputs) -> None: + """Batch of 2 with padded targets, references and text. + + Only valid target frames are compared. Padded frames are unconstrained: + the reference leaves them to whatever masked attention emits there. + """ + reference, port = models + kwargs = dict( + text=inputs["text2"], + time=inputs["time"], + mask=inputs["mask2"], + c_mask=inputs["c_mask2"], + ref=inputs["ref2"], + ref_mask=inputs["ref_mask2"], + ) + with torch.no_grad(), _autocast(): + want = reference(x=inputs["x2"], **kwargs) + got = port(inputs["x2"], **kwargs) + + assert got.shape == want.shape == (2, TARGET_FRAMES, port.latent_dim) + valid = inputs["mask2"].unsqueeze(-1).expand_as(got) + _assert_close("masked", got[valid], want[valid], FORWARD_ATOL, FORWARD_COS) + + +def _set_mask_flag(reference, port, enabled: bool) -> None: + """Flip attn_mask_enabled on both models in place, so no second copy is loaded.""" + for module in port.modules(): + if hasattr(module, "attn_mask_enabled"): + module.attn_mask_enabled = enabled + for blocks in (reference.transformer_blocks, reference.single_transformer_blocks): + for block in blocks: + block.attn.processor.attn_mask_enabled = enabled + + +def test_mask_flag_disabled(models, inputs) -> None: + """``attn_mask_enabled=False`` must track the reference with its own flag off. + + With the flag off both models run attention unmasked and still zero the + padded outputs, so the padded batch is compared in full. + """ + reference, port = models + kwargs = dict( + text=inputs["text2"], + time=inputs["time"], + mask=inputs["mask2"], + c_mask=inputs["c_mask2"], + ref=inputs["ref2"], + ref_mask=inputs["ref_mask2"], + ) + with torch.no_grad(), _autocast(): + masked = port(inputs["x2"], **kwargs) + _set_mask_flag(reference, port, False) + try: + want = reference(x=inputs["x2"], **kwargs) + got = port(inputs["x2"], **kwargs) + finally: + _set_mask_flag(reference, port, True) + + _assert_close("mask-flag-off", got, want, FORWARD_ATOL, FORWARD_COS) + # Guard against a vacuous pass: the flag has to actually change something. + assert (got.float() - masked.float()).abs().max().item() > 1e-3 + + +def test_sampler_generator_matches_seed(models, inputs) -> None: + """A generator seeded with 0 must draw exactly what ``seed=0`` draws.""" + _, port = models + common = dict( + gen_frames=TARGET_FRAMES, + nfe=2, + cfg_strength=0.0, + sway_sampling_coef=SWAY_COEF, + text=inputs["text"], + c_mask=inputs["c_mask"], + ref=inputs["ref"], + ref_mask=None, + ) + with torch.no_grad(), _autocast(): + by_seed = sample_latents(port, seed=0, **common) + by_generator = sample_latents(port, generator=torch.Generator(device=DEVICE).manual_seed(0), **common) + + print(f"[generator] max|Δ|={(by_generator.float() - by_seed.float()).abs().max().item():.4e}") + torch.testing.assert_close(by_generator, by_seed, rtol=0, atol=0) + + +def test_cache_needs_clearing(models, inputs) -> None: + """The cached text projection is reused until ``clear_cache()``. + + ``sample_latents`` clears on the way out, so a pipeline that reuses the + module across requests is safe without doing anything. A caller driving + ``forward(cache=True)`` itself has to clear when the text changes. + """ + _, port = models + other_text = inputs["text"] * 0.5 + kwargs = dict(time=inputs["time"], c_mask=inputs["c_mask"], ref=inputs["ref"], cfg_infer=True) + + port.clear_cache() + with torch.no_grad(), _autocast(): + want = port(inputs["x"], inputs["text"], **kwargs) + port(inputs["x"], other_text, cache=True, **kwargs) + stale = port(inputs["x"], inputs["text"], cache=True, **kwargs) + port.clear_cache() + fresh = port(inputs["x"], inputs["text"], cache=True, **kwargs) + port.clear_cache() + + # The cache is real: the second call reused the first call's projection. + assert (stale.float() - want.float()).abs().max().item() > 1e-3 + torch.testing.assert_close(fresh, want, rtol=0, atol=0) + assert port.text_cond is None and port.text_uncond is None + + +def test_rope_matches_x_transformers() -> None: + """The local rotary implementation against x_transformers, in fp32.""" + from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb + + device = torch.device(DEVICE) + dim_head, seq_len = 64, TEXT_TOKENS + REF_FRAMES + TARGET_FRAMES + q = torch.randn(2, 24, seq_len, dim_head, dtype=torch.float32, device=device) + + freqs = RotaryEmbedding(dim_head).to(device).forward_from_seq_len(seq_len)[0] + want = apply_rotary_pos_emb(q, freqs) + got = _apply_rope(q, Rotary(dim_head).to(device)(seq_len)) + + max_abs, cos, _ = _stats(got, want) + print(f"[rope] max|Δ|={max_abs:.4e} cos={cos:.10f}") + assert max_abs < ROPE_ATOL, f"rope: max|Δ|={max_abs:.4e} >= {ROPE_ATOL:.1e}" + + +def test_rope_survives_half_precision_cast() -> None: + """Casting the module to bf16 must not round the rotary frequencies. + + The reference stores them in a buffer, so ``.to(bfloat16)`` rounds both the + frequencies and the integer positions and silently changes every rotation. + The port recomputes them instead. + """ + dim_head, seq_len = 64, TEXT_TOKENS + REF_FRAMES + TARGET_FRAMES + exact = Rotary(dim_head)(seq_len) + got = Rotary(dim_head).to(torch.bfloat16)(seq_len) + + assert got.dtype == torch.float32 + torch.testing.assert_close(got, exact, rtol=0, atol=0) + + +def _reference_euler(reference, *, x, text, c_mask, ref, steps, cfg_strength, sway) -> torch.Tensor: + """Upstream's ODE loop, transcribed. + + ``CFMEdit.sample`` hands the backbone to ``torchdiffeq.odeint(method= + "euler")``, which on a fixed grid is ``y += (t1 - t0) * f(t0, y)``. + Constructing a real ``CFMEdit`` would drag in the Qwen-Omni text encoder, + so only the backbone here comes from upstream. + """ + t = torch.linspace(0, 1, steps + 1, device=x.device, dtype=torch.float32) + if sway is not None: + t = t + sway * (torch.cos(math.pi / 2 * t) - 1 + t) + + shared = dict(text=text, c_mask=c_mask, ref=ref, mask=None, ref_mask=None, cache=True) + for i in range(steps): + if cfg_strength < 1e-5: + v = reference(x=x, time=t[i], **shared) + else: + pred = reference(x=x, time=t[i], cfg_infer=True, **shared) + v_cond, v_uncond = torch.chunk(pred, 2, dim=0) + v = v_cond + (v_cond - v_uncond) * cfg_strength + x = x + (t[i + 1] - t[i]) * v + reference.clear_cache() + return x + + +@pytest.mark.parametrize("cfg_strength", [0.0, 2.0]) +def test_euler_sample(models, inputs, cfg_strength: float) -> None: + """An 8-step Euler trajectory, with and without guidance.""" + reference, port = models + shared = dict(text=inputs["text"], c_mask=inputs["c_mask"], ref=inputs["ref"]) + + with torch.no_grad(), _autocast(): + got = sample_latents( + port, + gen_frames=TARGET_FRAMES, + ref_mask=None, + nfe=SAMPLE_STEPS, + cfg_strength=cfg_strength, + sway_sampling_coef=SWAY_COEF, + seed=SEED, + **shared, + ) + # Re-seeding reproduces the noise sample_latents drew for itself. + torch.manual_seed(SEED) + noise = torch.randn( + TARGET_FRAMES, port.latent_dim, device=inputs["ref"].device, dtype=inputs["ref"].dtype + ).unsqueeze(0) + want = _reference_euler( + reference, x=noise, steps=SAMPLE_STEPS, cfg_strength=cfg_strength, sway=SWAY_COEF, **shared + ) + + assert got.shape == want.shape == (1, TARGET_FRAMES, port.latent_dim) + _assert_close(f"sample-cfg{cfg_strength:g}", got, want, SAMPLE_ATOL, SAMPLE_COS) diff --git a/tests/diffusion/models/auk/test_auk_vae_parity.py b/tests/diffusion/models/auk/test_auk_vae_parity.py new file mode 100644 index 00000000000..cafc99614fc --- /dev/null +++ b/tests/diffusion/models/auk/test_auk_vae_parity.py @@ -0,0 +1,254 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""AuK VAE parity: the numeric oracle for the vendored BigVGAN-flow codec. + +Checks ``AuKVAE`` against the AuK reference implementation on real speech, with +both models loaded from the same released ``vae.safetensors``: + + * encoder statistics (mean / log-std) agree, + * ``encode(sample=False)`` reproduces the reference mean path, + * ``encode(sample=True)`` reproduces ``encoding_and_normalization`` for a + matching RNG seed, + * ``decode`` reproduces ``denormalize`` + ``inference_from_latents``, in fp32 + and (reported, not gated tightly) under bf16 autocast, + * the waveform survives a mean-path round trip. + +Runs on GPU in fp32 and needs both the released checkpoint and the reference +``auk`` package importable, so it is a local_model test rather than CI. Point +``AUK_CKPT_DIR`` at the directory holding ``vae.safetensors`` and ``AUK_REF_WAV`` +at a speech clip (default ``$AUK_CKPT_DIR/ref.wav``):: + + AUK_CKPT_DIR=... AUK_REF_WAV=... python -m pytest \ + tests/diffusion/models/auk/test_auk_vae_parity.py -v -s + +Unlike its sibling directories this one carries no ``__init__.py`` on purpose: it would put +``tests/diffusion/models`` on ``sys.path`` ahead of site-packages, and this directory's own name +would then shadow the reference ``auk`` package the parity check imports. +""" + +from __future__ import annotations + +import importlib.util +import os +from pathlib import Path + +import pytest +import soundfile as sf +import torch +import torchaudio.functional as AF + +from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE + +pytestmark = [pytest.mark.local_model, pytest.mark.diffusion, pytest.mark.gpu, pytest.mark.cuda] + +# The released model.vae.model_init_kwargs, shared by both models under test. +VAE_CONFIG = { + "upsample_rates": [5, 4, 3, 2, 2, 2], + "upsample_kernel_sizes": [10, 8, 6, 4, 4, 4], + "upsample_initial_channel": 1536, + "resblock_kernel_sizes": [3, 7, 11], + "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + "downsample_rates": [2, 2, 2, 3, 4, 5], + "downsample_channels": [12, 24, 48, 96, 192, 384, 768], + "snake_logscale": True, + "latent_dim": 64, + "use_vae": True, + "causal": True, + "flow_hidden_channels": 256, + "act_causal": True, +} + +SAMPLE_RATE = 24000 +HOP_SIZE = 480 +LATENT_DIM = 64 +SEED = 1234 +# Keep the clip short: the parity run shares its GPU with other jobs. +MAX_SECONDS = 4.0 +TOLERANCE = 1e-5 +DEVICE = "cuda" + +_CKPT_DIR = os.environ.get("AUK_CKPT_DIR") +_VAE_PATH = Path(_CKPT_DIR) / "vae.safetensors" if _CKPT_DIR else None +_REF_WAV = Path(os.environ.get("AUK_REF_WAV") or (Path(_CKPT_DIR) / "ref.wav" if _CKPT_DIR else "unset")) + + +def _reference_importable() -> bool: + """Resolve ``auk.model``, not ``auk``. + + An ``__init__.py`` in this directory would make its own name resolve as ``auk`` and satisfy a + top-level probe, leaving the fixtures to fail at setup instead of skipping. The submodule + probe fails in that case, which is the honest answer. + """ + try: + return importlib.util.find_spec("auk.model") is not None + except (ImportError, ValueError): + return False + + +def _skip_reason() -> str | None: + if not _CKPT_DIR: + return "set AUK_CKPT_DIR to the directory holding vae.safetensors" + if _VAE_PATH is None or not _VAE_PATH.is_file(): + return f"no vae.safetensors under {_CKPT_DIR}" + if not _REF_WAV.is_file(): + return f"no reference clip at {_REF_WAV} (set AUK_REF_WAV)" + if not _reference_importable(): + return "cannot import the reference auk.model package (an __init__.py here would shadow it)" + if not torch.cuda.is_available(): + return "parity runs in fp32 on GPU" + return None + + +pytestmark.append(pytest.mark.skipif(_skip_reason() is not None, reason=_skip_reason() or "")) + + +def _max_abs(lhs: torch.Tensor, rhs: torch.Tensor) -> float: + return (lhs.float() - rhs.float()).abs().max().item() + + +def _snr_db(signal: torch.Tensor, estimate: torch.Tensor) -> float: + noise = (signal.float() - estimate.float()).pow(2).sum() + return (10.0 * torch.log10(signal.float().pow(2).sum() / noise)).item() + + +@pytest.fixture(scope="module") +def clip() -> torch.Tensor: + """Mono 24 kHz speech, trimmed to a whole number of latent frames.""" + audio, sample_rate = sf.read(str(_REF_WAV), dtype="float32", always_2d=True) # (T, C) + wav = torch.from_numpy(audio).mean(dim=1, keepdim=True).transpose(0, 1) + if sample_rate != SAMPLE_RATE: + wav = AF.resample(wav, sample_rate, SAMPLE_RATE) + samples = min(int(MAX_SECONDS * SAMPLE_RATE), wav.size(-1)) // HOP_SIZE * HOP_SIZE + wav = wav[:, :samples].contiguous().to(DEVICE) + print(f"\n[clip] {_REF_WAV.name}: {samples} samples, {samples / SAMPLE_RATE:.2f} s, {samples // HOP_SIZE} frames") + return wav + + +@pytest.fixture(scope="module") +def reference(clip: torch.Tensor) -> dict[str, torch.Tensor]: + """Reference tensors from the upstream model, which is then released so one model fits at a time.""" + from auk.model.vae import load_vae_model + from auk.model.vae.bigvgan_flow_vae import BigVGANFlowVAEConfig + + model = load_vae_model( + vae_name="BigVGANFlowVAE", + vae_cfg=BigVGANFlowVAEConfig.from_dict(VAE_CONFIG), + vae_ckpt=str(_VAE_PATH), + map_location="cpu", + ) + model = model.to(DEVICE).eval() + model.requires_grad_(False) + + with torch.no_grad(): + mean, log_std = model.audio_encoder(clip.unsqueeze(1)).chunk(2, 1) + scale = torch.sqrt(model.global_log_std.float()) + mean_latents = (mean.transpose(1, 2).float() - model.global_mean.float()) / scale + torch.manual_seed(SEED) + sampled_latents, _ = model.encoding_and_normalization(clip.unsqueeze(1)) + decoded = model.inference_from_latents(model.denormalize(mean_latents).permute(0, 2, 1)) + + tensors = { + "mean": mean.clone(), + "log_std": log_std.clone(), + "mean_latents": mean_latents.clone(), + "sampled_latents": sampled_latents.clone(), + "decoded": decoded.squeeze(1).clone(), + } + del model + torch.accelerator.empty_cache() + return tensors + + +@pytest.fixture(scope="module") +def port(reference: dict[str, torch.Tensor]) -> tuple[AuKVAE, list[str], list[str]]: + """The vendored module, loaded from the same checkpoint. Depends on ``reference`` for load ordering. + + Placed on the accelerator before loading so weight norm is folded there; see ``load_weights``. + """ + model = AuKVAE.from_config(VAE_CONFIG).to(DEVICE) + missing, unexpected = model.load_weights(_VAE_PATH) + model = model.eval() + model.requires_grad_(False) + return model, missing, unexpected + + +def test_checkpoint_covers_every_inference_tensor(port: tuple[AuKVAE, list[str], list[str]]) -> None: + model, missing, unexpected = port + print(f"[keys] missing={len(missing)} unexpected={len(unexpected)}") + print(f"[keys] unexpected prefixes={sorted({key.split('.')[0] for key in unexpected})}") + assert missing == [] + assert unexpected and all(key.startswith("flow.") for key in unexpected) + assert (model.hop_size, model.sample_rate, model.latent_dim) == (HOP_SIZE, SAMPLE_RATE, LATENT_DIM) + + +def test_encoder_statistics_match( + clip: torch.Tensor, reference: dict[str, torch.Tensor], port: tuple[AuKVAE, list[str], list[str]] +) -> None: + model = port[0] + with torch.no_grad(): + mean, log_std = model.audio_encoder(clip.unsqueeze(1)).chunk(2, 1) + mean_diff = _max_abs(mean, reference["mean"]) + log_std_diff = _max_abs(log_std, reference["log_std"]) + print(f"[encoder] max|dmean|={mean_diff:.3e} max|dlog_std|={log_std_diff:.3e}") + assert mean_diff < TOLERANCE + assert log_std_diff < TOLERANCE + # Folded and hook-computed weights are bit-identical on one device, so the whole encoder is. + assert (mean_diff, log_std_diff) == (0.0, 0.0) + + +def test_encode_mean_path_matches( + clip: torch.Tensor, reference: dict[str, torch.Tensor], port: tuple[AuKVAE, list[str], list[str]] +) -> None: + model = port[0] + with torch.no_grad(): + latents = model.encode(clip) + assert latents.shape == (1, clip.size(-1) // HOP_SIZE, LATENT_DIM) + diff = _max_abs(latents, reference["mean_latents"]) + print(f"[encode mean] shape={tuple(latents.shape)} max|d|={diff:.3e}") + assert diff == 0.0 + + +def test_encode_sampled_matches_seeded_reference( + clip: torch.Tensor, reference: dict[str, torch.Tensor], port: tuple[AuKVAE, list[str], list[str]] +) -> None: + model = port[0] + generator = torch.Generator(device=DEVICE) + generator.manual_seed(SEED) + with torch.no_grad(): + latents = model.encode(clip, sample=True, generator=generator) + diff = _max_abs(latents, reference["sampled_latents"]) + spread = _max_abs(latents, reference["mean_latents"]) + print(f"[encode sample] max|d vs seeded reference|={diff:.3e} max|d vs mean path|={spread:.3e}") + assert diff < TOLERANCE + assert diff == 0.0 # a fresh device Generator replays the default one after torch.manual_seed + assert spread > TOLERANCE # the draw actually moved the latents + + +def test_decode_matches(reference: dict[str, torch.Tensor], port: tuple[AuKVAE, list[str], list[str]]) -> None: + model = port[0] + latents = reference["mean_latents"] + with torch.no_grad(): + wav = model.decode(latents) + assert wav.shape == (1, latents.size(1) * HOP_SIZE) + diff = _max_abs(wav, reference["decoded"]) + print(f"[decode fp32] shape={tuple(wav.shape)} max|d|={diff:.3e}") + assert diff < TOLERANCE + assert diff == 0.0 + + with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16): + wav_bf16 = model.decode(latents) + bf16_diff = _max_abs(wav_bf16, reference["decoded"]) + print(f"[decode bf16 autocast] max|d| vs fp32 reference={bf16_diff:.3e} snr={_snr_db(wav, wav_bf16):.2f} dB") + assert torch.isfinite(wav_bf16).all() + assert bf16_diff < 0.5 # sanity only: bf16 is reported, not held to fp32 parity + + +def test_round_trip_snr(clip: torch.Tensor, port: tuple[AuKVAE, list[str], list[str]]) -> None: + model = port[0] + with torch.no_grad(): + wav = model.decode(model.encode(clip)) + assert wav.shape == clip.shape + snr = _snr_db(clip, wav) + print(f"[round trip] snr={snr:.2f} dB peak={wav.abs().max().item():.4f}") + assert snr > 10.0 diff --git a/tests/diffusion/models/auk/test_pipeline_auk.py b/tests/diffusion/models/auk/test_pipeline_auk.py new file mode 100644 index 00000000000..261c5240daa --- /dev/null +++ b/tests/diffusion/models/auk/test_pipeline_auk.py @@ -0,0 +1,514 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""AuK stage-1 pipeline: request parsing on CPU, generation parity on a GPU. + +The CPU tests stub the transformer, the codec and the ODE so they exercise only +what the pipeline itself owns: how a request turns into a target length, a +sampling schedule and a noise generator. + +The parity test replays the saved upstream reference for the zero-shot TTS case +and reports latent, per-frame and waveform distances rather than asserting bit +equality, because two differences live in the conditioning rather than in the +port: the saved text condition is an fp32 thinker forward while the reference +generated under bf16 autocast, and the reference's stochastic VAE draw came from +the process-global RNG seeded at 1234. + +It replays with ``vae_sample=False``. That is the deliberate choice: the +reference re-seeds the global RNG immediately before drawing the ODE noise, so a +per-request generator reproduces that noise exactly (verified bit-identical) as +long as nothing else has drawn from it. Sampling the VAE posterior consumes the +generator first, which moves the noise and produces a different, equally valid +realization; comparing latents then measures the noise, not the port. The +posterior mean costs a reference latent that is 0.045 MSE (0.969 cosine) away +from the draw the reference used, which is the smaller of the two errors. +Generate the reference artifacts once with the upstream package installed +(``tools/auk_parity_reference.py --auk-repo ... --ckpt-dir ... --qwen-dir ... +--out /path/to/auk-parity``; it writes ``parity_ref/`` and +``fusion_ref``), then run:: + + AUK_OMNI_CKPT_DIR=/path/to/auk-omni-base \ + AUK_PARITY_REF=/path/to/auk-parity/parity_ref/base \ + python -m pytest -s tests/diffusion/models/auk/test_pipeline_auk.py + +``AUK_OMNI_CKPT_DIR`` is the assembled directory this pipeline loads, which is +not the same thing as the released AuK snapshot that the transformer and codec +parity tests in this directory read from ``AUK_CKPT_DIR``. ``AUK_CKPT_DIR`` is +accepted as a fallback, and the test skips rather than fails when the directory +it resolves to is not an assembled one. +""" + +from __future__ import annotations + +import json +import math +import os +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +import soundfile as sf +import torch +import torch.nn.functional as F +import torchaudio +from torch import nn + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.models.auk import pipeline_auk +from vllm_omni.diffusion.models.auk.pipeline_auk import AuKPipeline +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +LATENT_DIM = 64 +HOP = 480 +SAMPLE_RATE = 24000 +TEXT_HIDDEN_DIM = 2048 + +# The distilled student's grid, exact values from the released checkpoint. +FLASH_T_GRID = [0.0, 0.07612049579620361, 0.2928932309150696, 0.6173166036605835, 1.0] + +CKPT_DIR = os.environ.get("AUK_OMNI_CKPT_DIR") or os.environ.get("AUK_CKPT_DIR") +PARITY_REF = os.environ.get("AUK_PARITY_REF") +# An assembled directory is the only thing this pipeline can load. +IS_ASSEMBLED = bool(CKPT_DIR) and (Path(CKPT_DIR).expanduser() / "config.json").is_file() + +# Measured on the released base checkpoint: mean frame cosine 0.961 to 0.968, +# relative latent MSE 0.051 to 0.063, log-mel distance 0.34. That spread is +# bf16 against fp32 DiT weights plus where the codec folds its weight norm, +# neither of which is worth gating on. A port that has lost the conditioning +# sits at 0.110 cosine, 1.82 relative MSE and 2.88 log-mel, so these gates +# still separate the two by an order of magnitude. +PARITY_MIN_FRAME_COSINE = 0.93 +PARITY_MAX_RELATIVE_MSE = 0.15 +PARITY_MAX_MEL_DISTANCE = 0.6 + + +class _StubTransformer(nn.Module): + """Records its construction arguments; owns one weight so strict loading runs.""" + + def __init__( + self, + *, + dim: int, + heads: int, + dim_head: int, + ff_mult: float, + latent_dim: int, + text_hidden_dim: int, + num_layers: int, + num_single_layers: int, + attn_mask_enabled: bool = True, + ) -> None: + super().__init__() + self.init_kwargs = { + "dim": dim, + "heads": heads, + "dim_head": dim_head, + "ff_mult": ff_mult, + "latent_dim": latent_dim, + "text_hidden_dim": text_hidden_dim, + "num_layers": num_layers, + "num_single_layers": num_single_layers, + "attn_mask_enabled": attn_mask_enabled, + } + self.latent_dim = latent_dim + self.proj = nn.Linear(latent_dim, latent_dim, bias=False) + + +class _StubVAE(nn.Module): + """Deterministic stand-in for the codec, one latent frame per hop.""" + + hop_size = HOP + sample_rate = SAMPLE_RATE + latent_dim = LATENT_DIM + + def __init__(self) -> None: + super().__init__() + self.encode_calls: list[dict[str, Any]] = [] + self.decode_calls = 0 + self.weights_path: str | None = None + self.config: dict[str, Any] = {} + + @classmethod + def from_config(cls, cfg: dict[str, Any]) -> _StubVAE: + vae = cls() + vae.config = dict(cfg) + return vae + + def load_weights(self, path: str, **_: Any) -> tuple[list[str], list[str]]: + self.weights_path = str(path) + return [], [] + + def encode( + self, + wav: torch.Tensor, + *, + sample: bool = False, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + frames = wav.shape[-1] // HOP + self.encode_calls.append({"samples": int(wav.shape[-1]), "sample": sample, "generator": generator}) + if sample: + return torch.randn(1, frames, LATENT_DIM, generator=generator, device=wav.device) + return torch.zeros(1, frames, LATENT_DIM, device=wav.device) + + def decode(self, latents: torch.Tensor) -> torch.Tensor: + self.decode_calls += 1 + return torch.zeros(1, latents.shape[1] * HOP, device=latents.device) + + +def _stub_sampler(calls: list[dict[str, Any]]): + """Record the ODE arguments and return noise drawn from the request generator.""" + + def sample_latents(dit: nn.Module, **kwargs: Any) -> torch.Tensor: + calls.append(kwargs) + shape = (1, kwargs["gen_frames"], kwargs["latent_dim"]) + generator = kwargs.get("generator") + if generator is None: + return torch.zeros(*shape, device=kwargs["device"], dtype=kwargs["dtype"]) + return torch.randn(*shape, generator=generator, device=kwargs["device"], dtype=kwargs["dtype"]) + + return sample_latents + + +def _write_checkpoint(root: Path, variant: str) -> Path: + """Write a tiny assembled checkpoint the stubs can load.""" + + from safetensors.torch import save_file + + model_dir = root / f"auk-{variant}" + model_dir.mkdir(parents=True, exist_ok=True) + config = { + "model_type": "auk", + "architectures": ["AuKForConditionalGeneration"], + "dit": { + "dim": 16, + "heads": 2, + "dim_head": 8, + "ff_mult": 2, + "text_hidden_dim": TEXT_HIDDEN_DIM, + "num_layers": 1, + "num_single_layers": 1, + "attn_mask_enabled": True, + }, + "vae": { + "latent_dim": LATENT_DIM, + "downsample_rate": HOP, + "target_sample_rate": SAMPLE_RATE, + "model_init_kwargs": {"latent_dim": LATENT_DIM}, + }, + "variant": variant, + "flash_t_grid": FLASH_T_GRID, + "defaults": {"nfe": 32, "cfg": 2.0, "sway": -1.0}, + } + (model_dir / "config.json").write_text(json.dumps(config)) + save_file( + { + "transformer.proj.weight": torch.zeros(LATENT_DIM, LATENT_DIM), + "layer_weights": torch.zeros(36), + "layer_scale": torch.ones(1), + }, + str(model_dir / "auk.safetensors"), + ) + save_file({"global_mean": torch.zeros(LATENT_DIM)}, str(model_dir / "vae.safetensors")) + return model_dir + + +@pytest.fixture +def build_pipeline(tmp_path, monkeypatch): + """Build an AuKPipeline whose transformer, codec and ODE are stubs.""" + + monkeypatch.setattr(pipeline_auk, "get_local_device", lambda: torch.device("cpu")) + monkeypatch.setattr(pipeline_auk, "AuKTransformer", _StubTransformer) + monkeypatch.setattr(pipeline_auk, "AuKVAE", _StubVAE) + + def _build(variant: str = "base") -> tuple[AuKPipeline, list[dict[str, Any]]]: + calls: list[dict[str, Any]] = [] + monkeypatch.setattr(pipeline_auk, "sample_latents", _stub_sampler(calls)) + od_config = OmniDiffusionConfig( + model=str(_write_checkpoint(tmp_path, variant)), + dtype=torch.float32, + model_class_name="AuKPipeline", + ) + return AuKPipeline(od_config=od_config), calls + + return _build + + +def _prompt( + *, + tokens: int = 8, + audio: Any = None, + knobs: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Build the diffusion prompt the stage input processor would emit.""" + + prompt: dict[str, Any] = { + "prompt": "", + "prompt_embeds": torch.zeros(tokens, TEXT_HIDDEN_DIM), + "additional_information": {"auk": dict(knobs or {})}, + } + if audio is not None: + prompt["multi_modal_data"] = {"audio": audio} + return prompt + + +def _batch(prompt: dict[str, Any], **sampling: Any) -> DiffusionRequestBatch: + request = OmniDiffusionRequest( + prompt=prompt, + sampling_params=OmniDiffusionSamplingParams(**sampling), + request_id="auk-test-0", + ) + return DiffusionRequestBatch(requests=[request]) + + +def _silence(seconds: float, sample_rate: int = SAMPLE_RATE) -> tuple[np.ndarray, int]: + return np.zeros(int(seconds * sample_rate), dtype=np.float32), sample_rate + + +@pytest.mark.core_model +@pytest.mark.cpu +class TestRequestParsing: + def test_pipeline_declares_audio_output(self, build_pipeline): + pipeline, _ = build_pipeline() + + assert pipeline.support_audio_output is True + assert pipeline.audio_sample_rate == SAMPLE_RATE + assert pipeline.supports_request_batch is False + # The warmup request cannot carry an encoder-stage text condition. + assert pipeline.dummy_run_num_frames == 0 + + def test_config_reaches_the_transformer(self, build_pipeline): + pipeline, _ = build_pipeline() + + assert pipeline.dit.init_kwargs["latent_dim"] == LATENT_DIM + assert pipeline.dit.init_kwargs["text_hidden_dim"] == TEXT_HIDDEN_DIM + assert pipeline.dit.init_kwargs["attn_mask_enabled"] is True + + def test_gen_seconds_sets_the_target_length(self, build_pipeline): + pipeline, calls = build_pipeline() + + outputs = pipeline.forward(_batch(_prompt(audio=_silence(2.0), knobs={"gen_seconds": 6.0}), seed=1)) + + assert calls[0]["gen_frames"] == math.ceil(6.0 * SAMPLE_RATE / HOP) + assert calls[0]["ref"].shape == (1, 100, LATENT_DIM) + assert calls[0]["ref_mask"].shape == (1, 100) + assert calls[0]["c_mask"].shape == (1, 8) + assert outputs[0].output.shape == (300 * HOP,) + assert outputs[0].output.dtype is torch.float32 + + def test_target_length_falls_back_to_the_source_clip(self, build_pipeline): + pipeline, calls = build_pipeline() + + pipeline.forward(_batch(_prompt(audio=_silence(2.0), knobs={"gen_seconds": None}), seed=1)) + + assert calls[0]["gen_frames"] == 100 + + def test_text_only_request_without_a_duration_is_rejected(self, build_pipeline): + pipeline, calls = build_pipeline() + + with pytest.raises(ValueError, match="gen_seconds"): + pipeline.forward(_batch(_prompt(knobs={"gen_seconds": None}), seed=1)) + assert calls == [] + + def test_text_only_request_gets_an_empty_reference(self, build_pipeline): + pipeline, calls = build_pipeline() + + pipeline.forward(_batch(_prompt(knobs={"gen_seconds": 3.5}), seed=1)) + + assert calls[0]["ref"].shape == (1, 0, LATENT_DIM) + assert calls[0]["gen_frames"] == 175 + + def test_source_clip_is_resampled_to_the_codec_rate(self, build_pipeline): + pipeline, _ = build_pipeline() + + pipeline.forward(_batch(_prompt(audio=_silence(1.0, 16000), knobs={"gen_seconds": 1.0}), seed=1)) + + assert pipeline.vae.encode_calls[0]["samples"] == SAMPLE_RATE + + def test_stereo_source_is_mixed_to_mono(self, build_pipeline): + pipeline, _ = build_pipeline() + stereo = np.zeros((2, SAMPLE_RATE), dtype=np.float32) + + pipeline.forward(_batch(_prompt(audio=(stereo, SAMPLE_RATE), knobs={"gen_seconds": 1.0}), seed=1)) + + assert pipeline.vae.encode_calls[0]["samples"] == SAMPLE_RATE + + def test_missing_text_condition_is_rejected(self, build_pipeline): + pipeline, _ = build_pipeline() + prompt = _prompt(knobs={"gen_seconds": 1.0}) + del prompt["prompt_embeds"] + + with pytest.raises(ValueError, match="prompt_embeds"): + pipeline.forward(_batch(prompt, seed=1)) + + def test_base_variant_uses_the_requested_schedule(self, build_pipeline): + pipeline, calls = build_pipeline() + + pipeline.forward( + _batch( + _prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "sway": -1.0}), + seed=1, + num_inference_steps=16, + guidance_scale=3.0, + ) + ) + + assert calls[0]["nfe"] == 16 + assert calls[0]["cfg_strength"] == 3.0 + assert calls[0]["sway_sampling_coef"] == -1.0 + assert calls[0]["t_grid"] is None + + def test_schedule_defaults_come_from_the_checkpoint(self, build_pipeline): + pipeline, calls = build_pipeline() + + pipeline.forward(_batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "sway": None}), seed=1)) + + assert calls[0]["nfe"] == 32 + assert calls[0]["cfg_strength"] == 2.0 + assert calls[0]["sway_sampling_coef"] == -1.0 + + def test_flash_variant_locks_the_distilled_recipe(self, build_pipeline): + pipeline, calls = build_pipeline("flash") + + pipeline.forward( + _batch( + _prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "sway": -1.0}), + seed=1, + num_inference_steps=32, + guidance_scale=2.0, + ) + ) + + assert calls[0]["nfe"] == 4 + assert calls[0]["cfg_strength"] == 0.0 + assert calls[0]["sway_sampling_coef"] is None + assert calls[0]["t_grid"] == FLASH_T_GRID + + def test_noise_comes_from_a_seeded_request_generator(self, build_pipeline): + pipeline, calls = build_pipeline() + + outputs = pipeline.forward( + _batch(_prompt(knobs={"gen_seconds": 1.0}), seed=0, output_type="latent"), + ) + + expected = torch.randn(1, 50, LATENT_DIM, generator=torch.Generator().manual_seed(0)) + assert torch.equal(outputs[0].output, expected) + # The global RNG is never seeded on the pipeline's behalf. + assert calls[0]["seed"] is None + assert isinstance(calls[0]["generator"], torch.Generator) + + def test_latent_output_type_skips_the_decoder(self, build_pipeline): + pipeline, _ = build_pipeline() + + outputs = pipeline.forward(_batch(_prompt(knobs={"gen_seconds": 1.0}), seed=1, output_type="latent")) + + assert outputs[0].output.shape == (1, 50, LATENT_DIM) + assert pipeline.vae.decode_calls == 0 + + def test_vae_sample_draws_from_the_request_generator(self, build_pipeline): + pipeline, _ = build_pipeline() + + pipeline.forward(_batch(_prompt(audio=_silence(1.0), knobs={"gen_seconds": 1.0, "vae_sample": True}), seed=1)) + + call = pipeline.vae.encode_calls[0] + assert call["sample"] is True + assert isinstance(call["generator"], torch.Generator) + + def test_one_request_per_forward(self, build_pipeline): + pipeline, _ = build_pipeline() + batch = _batch(_prompt(knobs={"gen_seconds": 1.0}), seed=1) + batch.requests.append(batch.requests[0]) + + with pytest.raises(AssertionError, match="one request per forward"): + pipeline.forward(batch) + + +def _reference_audio_path(messages: list[dict[str, Any]]) -> str: + for message in messages: + for item in message.get("content") or (): + if isinstance(item, dict) and item.get("type") == "audio": + path = item.get("audio") or item.get("audio_url") + if path: + return str(path) + raise AssertionError("the saved reference case carries no source clip") + + +def _log_mel_distance(left: torch.Tensor, right: torch.Tensor, sample_rate: int) -> float: + """Mean absolute log-mel difference between two mono waveforms.""" + + mel = torchaudio.transforms.MelSpectrogram(sample_rate=sample_rate, n_fft=1024, hop_length=256, n_mels=80) + length = min(left.shape[-1], right.shape[-1]) + spectra = [torch.log(mel(wav[..., :length].float()) + 1e-5) for wav in (left, right)] + return float((spectra[0] - spectra[1]).abs().mean()) + + +@pytest.mark.local_model +@pytest.mark.diffusion +@pytest.mark.skipif( + not (IS_ASSEMBLED and PARITY_REF), + reason="needs AUK_PARITY_REF and AUK_OMNI_CKPT_DIR pointing at an assembled checkpoint directory", +) +def test_zero_shot_tts_tracks_the_upstream_reference(): + """Replay the saved zero-shot TTS case and report the distances.""" + + parity_dir = Path(PARITY_REF) + fusion_dir = Path(os.environ.get("AUK_FUSION_REF") or parity_dir.parent.parent / "fusion_ref") + reference = torch.load(parity_dir / "zs_tts_en.pt", map_location="cpu", weights_only=False) + fusion = torch.load(fusion_dir / "zs_tts_audio.pt", map_location="cpu", weights_only=False) + + # soundfile rather than torchaudio.load: the repo's test convention, and + # torchaudio's decoder needs an ffmpeg the port environment does not have. + samples, source_rate = sf.read(_reference_audio_path(reference["messages"]), dtype="float32", always_2d=True) + source = torch.from_numpy(samples).transpose(0, 1) + knobs = {"gen_seconds": reference["gen_seconds"], "sway": -1.0, "t_grid": None, "vae_sample": False} + prompt = { + "prompt": "", + # The encoder stage emits bf16; the saved fusion is an fp32 forward. + "prompt_embeds": fusion["fused"].to(torch.bfloat16), + "multi_modal_data": {"audio": (source, int(source_rate))}, + "additional_information": {"auk": knobs}, + } + od_config = OmniDiffusionConfig(model=CKPT_DIR, dtype=torch.bfloat16, model_class_name="AuKPipeline") + pipeline = AuKPipeline(od_config=od_config) + + def run(output_type: str) -> torch.Tensor: + batch = _batch( + dict(prompt), + seed=reference["seed"], + num_inference_steps=32, + guidance_scale=2.0, + output_type=output_type, + ) + return pipeline.forward(batch)[0].output + + # Two deterministic runs: the ODE is seeded, so the latents belong to the + # waveform even though the decode happens in the second call. + latents = run("latent") + waveform = run("pt") + + ref_frames = int(reference["ref_latent_lens"][0]) + expected = reference["generated"][:, ref_frames:, :].float() + assert latents.shape == expected.shape + assert torch.isfinite(latents).all() + assert torch.isfinite(waveform).all() + assert waveform.shape[-1] == expected.shape[1] * HOP + + cosine = F.cosine_similarity(latents[0], expected[0], dim=-1) + metrics = { + "gen_frames": int(expected.shape[1]), + "ref_frames": ref_frames, + "latent_mse": float(torch.mean((latents - expected) ** 2)), + "latent_relative_mse": float(torch.mean((latents - expected) ** 2) / expected.var()), + "frame_cosine_mean": float(cosine.mean()), + "frame_cosine_min": float(cosine.min()), + "log_mel_distance": _log_mel_distance(waveform, reference["audio"].reshape(-1), int(reference["sr"])), + "rms_port": float(waveform.pow(2).mean().sqrt()), + "rms_reference": float(reference["audio"].reshape(-1).pow(2).mean().sqrt()), + } + print("\nAuK zero-shot TTS parity:", json.dumps(metrics, indent=1)) + + assert metrics["frame_cosine_mean"] > PARITY_MIN_FRAME_COSINE + assert metrics["latent_relative_mse"] < PARITY_MAX_RELATIVE_MSE + assert metrics["log_mel_distance"] < PARITY_MAX_MEL_DISTANCE diff --git a/tests/e2e/offline_inference/test_auk.py b/tests/e2e/offline_inference/test_auk.py new file mode 100644 index 00000000000..746af7bd24f --- /dev/null +++ b/tests/e2e/offline_inference/test_auk.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Offline E2E smoke test for AuK (encoder stage -> diffusion stage). + +AuK ships as ``config.yaml`` plus loose safetensors and needs the Qwen2.5-Omni-3B +snapshot next to it, so the test expects an assembled directory produced by +``tools/prepare_auk_checkpoint.py`` and skips when it is not available: + + VLLM_OMNI_AUK_MODEL_DIR=/path/to/auk-omni python -m pytest tests/e2e/offline_inference/test_auk.py + +Two requests are exercised: zero-shot TTS from a reference clip and text-only +instruct TTS. The checks are structural (duration, sample rate, level, +finiteness); fidelity against the upstream implementation is covered by the +parity suites under ``tests/diffusion/models/auk``. +""" + +from __future__ import annotations + +import functools +import math +import os +from pathlib import Path + +import numpy as np +import pytest +import soundfile as sf + +from tests.helpers.mark import hardware_test +from tests.helpers.media import get_asset_path +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.model_extras.auk import HOP, SAMPLE_RATE, auk_prompt, auk_sampling_params + +MODEL_DIR_ENV = "VLLM_OMNI_AUK_MODEL_DIR" + +# Any clean speech clip works as the cloning reference; reuse the vendored +# CosyVoice3 prompt so no new binary asset is needed. +REFERENCE_WAV_PATH = get_asset_path("cosyvoice3/zero_shot_prompt.wav") +ZERO_SHOT_TEXT = "The quick brown fox jumps over the lazy dog." +ZERO_SHOT_SECONDS = 3.0 +# The upstream cookbook template: description first, then the content. +INSTRUCT_TEXT = ( + 'Generate speech based on the following description: "A calm young woman speaking warmly and slowly.". ' + 'The content to speak is: "Welcome back, how was your day?".' +) +INSTRUCT_SECONDS = 2.5 +# 8 Euler steps keep the test short; the base recipe is 32. +TEST_NFE = 8 + +_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"), + pytest.mark.parametrize( + "omni_runner", + [pytest.param((str(Path(_model_dir or ".").resolve()), get_deploy_config_path("auk.yaml"), {}), id="auk")], + indirect=True, + ), +] + + +@functools.lru_cache(maxsize=1) +def _load_reference_wav() -> tuple[np.ndarray, int]: + audio, sr = sf.read(str(REFERENCE_WAV_PATH), dtype="float32", always_2d=False) + if isinstance(audio, np.ndarray) and audio.ndim > 1: + audio = np.mean(audio, axis=-1) + return np.asarray(audio, dtype=np.float32), int(sr) + + +def _audio_from_output(output) -> tuple[np.ndarray, int]: + mm = getattr(output, "multimodal_output", None) or {} + assert "audio" in mm, f"no audio in the final output; keys={sorted(mm)}" + audio = mm["audio"] + if isinstance(audio, list): + audio = np.concatenate([np.asarray(a, dtype=np.float32).reshape(-1) for a in audio]) + if hasattr(audio, "detach"): + audio = audio.detach().cpu().float().numpy() + audio = np.asarray(audio, dtype=np.float32).reshape(-1) + sr = mm.get("audio_sample_rate", SAMPLE_RATE) + if isinstance(sr, list) and sr: + sr = sr[-1] + if hasattr(sr, "item"): + sr = sr.item() + return audio, int(sr) + + +def _check_waveform(audio: np.ndarray, sr: int, *, gen_seconds: float) -> None: + assert sr == SAMPLE_RATE + expected = math.ceil(gen_seconds * SAMPLE_RATE / HOP) * HOP + assert audio.size == expected, f"expected {expected} samples for {gen_seconds}s, got {audio.size}" + assert np.isfinite(audio).all(), "non-finite samples" + assert np.abs(audio).max() <= 1.0 + 1e-6, "audio exceeds full scale" + assert float(np.sqrt(np.mean(audio**2))) > 1e-3, "audio is silent" + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_auk_zero_shot_tts(omni_runner: OmniRunner) -> None: + """A reference clip plus an instruction yields speech of the requested length.""" + wav, sr = _load_reference_wav() + prompt = auk_prompt( + f"Say the following with the same voice: '{ZERO_SHOT_TEXT}'", + (wav, sr), + gen_seconds=ZERO_SHOT_SECONDS, + ) + outputs = omni_runner.omni.generate(prompt, auk_sampling_params(nfe=TEST_NFE, cfg=2.0, seed=0)) + assert outputs, "no outputs returned" + audio, out_sr = _audio_from_output(outputs[0]) + _check_waveform(audio, out_sr, gen_seconds=ZERO_SHOT_SECONDS) + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_auk_instruct_tts_text_only(omni_runner: OmniRunner) -> None: + """A text-only request (no reference audio) takes the no-prompt-audio path.""" + prompt = auk_prompt(INSTRUCT_TEXT, None, gen_seconds=INSTRUCT_SECONDS) + assert prompt["prompt"].rstrip().endswith("<|im_start|>assistant\n".rstrip()) + assert "||" in prompt["prompt"] + outputs = omni_runner.omni.generate(prompt, auk_sampling_params(nfe=TEST_NFE, cfg=2.0, seed=0)) + assert outputs, "no outputs returned" + audio, out_sr = _audio_from_output(outputs[0]) + _check_waveform(audio, out_sr, gen_seconds=INSTRUCT_SECONDS) diff --git a/tests/model_executor/stage_input_processors/test_auk.py b/tests/model_executor/stage_input_processors/test_auk.py new file mode 100644 index 00000000000..b2904e1f7b0 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_auk.py @@ -0,0 +1,188 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for the AuK encoder -> diffusion stage input processor.""" + +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest +import torch + +from vllm_omni.model_executor.stage_input_processors.auk import encoder2dit + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +HIDDEN = 2048 + + +def _request_output( + mm_output: Any, + completion_mm_output: Any = None, + request_id: str | None = None, +) -> SimpleNamespace: + """Build a minimal encoder-stage output mock.""" + return SimpleNamespace( + outputs=[SimpleNamespace(token_ids=[0], multimodal_output=completion_mm_output)], + multimodal_output=mm_output, + request_id=request_id, + ) + + +def _payload(text_cond: Any) -> dict[str, Any]: + """The slot the encoder actually emits: hidden_states.output.""" + return {"hidden_states": {"output": text_cond}} + + +def _prompt( + *, + audio: bool = False, + gen_seconds: float | None = None, + has_audio: bool | None = None, +) -> dict[str, Any]: + knobs: dict[str, Any] = { + "gen_seconds": gen_seconds, + "sway": -1.0, + "t_grid": None, + "vae_sample": False, + } + if has_audio is not None: + knobs["has_audio"] = has_audio + prompt: dict[str, Any] = { + "prompt": "<|im_start|>system\n", + "additional_information": {"auk": knobs}, + } + if audio: + prompt["multi_modal_data"] = {"audio": (np.zeros(24000, dtype=np.float32), 24000)} + return prompt + + +class TestTextCondExtraction: + def test_tensor_output(self): + cond = torch.zeros(7, HIDDEN, dtype=torch.bfloat16) + out = encoder2dit([_request_output(_payload(cond))], _prompt()) + assert out is not None + assert out["prompt"] == "" + assert out["prompt_embeds"].shape == (7, HIDDEN) + assert out["prompt_embeds"].dtype == torch.bfloat16 + assert not out["prompt_embeds"].is_cuda + + def test_list_of_tensors_is_concatenated(self): + parts = [torch.zeros(3, HIDDEN), torch.ones(4, HIDDEN)] + out = encoder2dit([_request_output(_payload(parts))], _prompt()) + assert out["prompt_embeds"].shape == (7, HIDDEN) + assert torch.equal(out["prompt_embeds"][3:], torch.ones(4, HIDDEN)) + + def test_single_element_list(self): + out = encoder2dit([_request_output(_payload([torch.zeros(5, HIDDEN)]))], _prompt()) + assert out["prompt_embeds"].shape == (5, HIDDEN) + + def test_numpy_output(self): + cond = np.zeros((4, HIDDEN), dtype=np.float32) + out = encoder2dit([_request_output(_payload(cond))], _prompt()) + assert isinstance(out["prompt_embeds"], torch.Tensor) + assert out["prompt_embeds"].shape == (4, HIDDEN) + + def test_completion_output_fallback(self): + cond = torch.zeros(6, HIDDEN) + out = encoder2dit( + [_request_output(None, completion_mm_output=_payload(cond))], + _prompt(), + ) + assert out["prompt_embeds"].shape == (6, HIDDEN) + + def test_hidden_states_as_a_struct(self): + """A deserialized payload arrives as a struct, not a mapping.""" + cond = torch.zeros(9, HIDDEN) + mm_output = {"hidden_states": SimpleNamespace(output=cond, layers=None)} + out = encoder2dit([_request_output(mm_output)], _prompt()) + assert out["prompt_embeds"].shape == (9, HIDDEN) + + def test_flattened_payload_key(self): + cond = torch.zeros(8, HIDDEN) + out = encoder2dit([_request_output({"hidden_states.output": cond})], _prompt()) + assert out["prompt_embeds"].shape == (8, HIDDEN) + + def test_direct_text_cond_key(self): + cond = torch.zeros(2, HIDDEN) + out = encoder2dit([_request_output({"text_cond": cond})], _prompt()) + assert out["prompt_embeds"].shape == (2, HIDDEN) + + def test_no_source_outputs_returns_none(self): + assert encoder2dit([], _prompt()) is None + + def test_missing_payload_raises(self): + with pytest.raises(ValueError, match="hidden_states.output"): + encoder2dit([_request_output({})], _prompt()) + + def test_empty_hidden_states_slot_raises(self): + with pytest.raises(ValueError, match="hidden_states.output"): + encoder2dit([_request_output({"hidden_states": {}})], _prompt()) + + def test_empty_multimodal_output_raises(self): + with pytest.raises(ValueError, match="hidden_states.output"): + encoder2dit([_request_output(None)], _prompt()) + + def test_empty_list_raises(self): + with pytest.raises(ValueError, match="empty text condition"): + encoder2dit([_request_output(_payload([]))], _prompt()) + + def test_wrong_rank_raises(self): + with pytest.raises(ValueError, match=r"\[tokens, hidden\]"): + encoder2dit([_request_output(_payload(torch.zeros(2, 7, HIDDEN)))], _prompt()) + + +class TestKnobPropagation: + def test_knobs_are_forwarded(self): + prompt = _prompt(audio=True, gen_seconds=6.0) + prompt["additional_information"]["auk"]["t_grid"] = [0.0, 0.5, 1.0] + prompt["additional_information"]["auk"]["vae_sample"] = True + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], prompt) + knobs = out["additional_information"]["auk"] + assert knobs["gen_seconds"] == 6.0 + assert knobs["sway"] == -1.0 + assert knobs["t_grid"] == [0.0, 0.5, 1.0] + assert knobs["vae_sample"] is True + + def test_sampling_knobs_are_not_copied(self): + """nfe/cfg/seed ride the stage-1 sampling params, not the prompt.""" + prompt = _prompt() + prompt["additional_information"]["auk"].update({"nfe": 8, "cfg": 1.0, "seed": 3}) + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], prompt) + knobs = out["additional_information"]["auk"] + assert set(knobs) == {"gen_seconds", "sway", "t_grid", "vae_sample", "has_audio"} + + def test_multi_modal_data_is_forwarded(self): + prompt = _prompt(audio=True) + out = encoder2dit( + [_request_output(_payload(torch.zeros(3, HIDDEN)))], + prompt, + requires_multimodal_data=True, + ) + assert out["multi_modal_data"] is prompt["multi_modal_data"] + + def test_text_only_request_has_no_multi_modal_data(self): + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], _prompt()) + assert out["multi_modal_data"] is None + assert out["additional_information"]["auk"]["has_audio"] is False + + def test_has_audio_declared_by_the_prompt_wins(self): + prompt = _prompt(audio=True, has_audio=True) + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], prompt) + assert out["additional_information"]["auk"]["has_audio"] is True + + def test_has_audio_inferred_from_multi_modal_data(self): + prompt = _prompt(audio=True) + prompt["additional_information"]["auk"].pop("has_audio", None) + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], prompt) + assert out["additional_information"]["auk"]["has_audio"] is True + + def test_prompt_as_list(self): + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], [_prompt(gen_seconds=2.0)]) + assert out["additional_information"]["auk"]["gen_seconds"] == 2.0 + + def test_missing_prompt_yields_empty_knobs(self): + out = encoder2dit([_request_output(_payload(torch.zeros(3, HIDDEN)))], None) + knobs = out["additional_information"]["auk"] + assert knobs["gen_seconds"] is None + assert knobs["has_audio"] is False diff --git a/tests/model_extras/test_auk_requests.py b/tests/model_extras/test_auk_requests.py new file mode 100644 index 00000000000..1f319d09221 --- /dev/null +++ b/tests/model_extras/test_auk_requests.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for the AuK request helpers.""" + +from __future__ import annotations + +import numpy as np +import pytest +from vllm import SamplingParams + +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_extras.auk import ( + NO_PROMPT_AUDIO, + auk_prompt, + auk_sampling_params, + resolve_gen_frames, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +INSTRUCTION = "Say the following with the same voice: 'Ladies and gentlemen'" + +EXPECTED_WITH_AUDIO = ( + "<|im_start|>system\n" + "You are a helpful assistant.<|im_end|>\n" + "<|im_start|>user\n" + f"{INSTRUCTION}<|audio_bos|><|AUDIO|><|audio_eos|><|im_end|>\n" + "<|im_start|>assistant\n" +) + +EXPECTED_TEXT_ONLY = ( + "<|im_start|>system\n" + "You are a helpful assistant.<|im_end|>\n" + "<|im_start|>user\n" + f"{INSTRUCTION}{NO_PROMPT_AUDIO}<|im_end|>\n" + "<|im_start|>assistant\n" +) + + +def _audio() -> tuple[np.ndarray, int]: + return np.zeros(24000, dtype=np.float32), 24000 + + +class TestPromptRendering: + def test_with_audio(self): + audio = _audio() + prompt = auk_prompt(INSTRUCTION, audio, gen_seconds=6.0) + assert prompt["prompt"] == EXPECTED_WITH_AUDIO + assert prompt["multi_modal_data"] == {"audio": audio} + assert prompt["additional_information"]["auk"]["has_audio"] is True + + def test_text_only_appends_the_marker(self): + prompt = auk_prompt(INSTRUCTION, gen_seconds=6.0) + assert prompt["prompt"] == EXPECTED_TEXT_ONLY + assert "multi_modal_data" not in prompt + assert prompt["additional_information"]["auk"]["has_audio"] is False + + def test_marker_is_not_appended_twice(self): + once = auk_prompt(INSTRUCTION, gen_seconds=6.0) + twice = auk_prompt(INSTRUCTION + NO_PROMPT_AUDIO, gen_seconds=6.0) + assert once["prompt"] == twice["prompt"] + assert twice["prompt"].count(NO_PROMPT_AUDIO) == 1 + + def test_marker_is_not_added_when_audio_is_present(self): + prompt = auk_prompt(INSTRUCTION, _audio()) + assert NO_PROMPT_AUDIO not in prompt["prompt"] + + def test_knob_defaults(self): + knobs = auk_prompt(INSTRUCTION, _audio())["additional_information"]["auk"] + assert knobs == { + "gen_seconds": None, + "sway": -1.0, + "t_grid": None, + "vae_sample": False, + "has_audio": True, + } + + def test_knobs_are_carried(self): + grid = [0.0, 0.25, 1.0] + knobs = auk_prompt( + INSTRUCTION, + _audio(), + gen_seconds=3.5, + sway=0.0, + t_grid=grid, + vae_sample=True, + )["additional_information"]["auk"] + assert knobs["gen_seconds"] == 3.5 + assert knobs["sway"] == 0.0 + assert knobs["vae_sample"] is True + assert knobs["t_grid"] == grid + # A copy, so a caller mutating its list cannot reach the request. + assert knobs["t_grid"] is not grid + + def test_sampling_knobs_are_absent_from_the_prompt(self): + knobs = auk_prompt(INSTRUCTION, _audio())["additional_information"]["auk"] + assert not {"nfe", "cfg", "seed"} & set(knobs) + + +class TestSamplingParams: + def test_two_stages_in_order(self): + params = auk_sampling_params() + assert len(params) == 2 + assert isinstance(params[0], SamplingParams) + assert isinstance(params[1], OmniDiffusionSamplingParams) + + def test_encoder_stage_is_prefill_only(self): + encoder = auk_sampling_params()[0] + assert encoder.max_tokens == 1 + assert encoder.temperature == 0.0 + assert encoder.detokenize is False + + def test_defaults_are_the_base_recipe(self): + diffusion = auk_sampling_params()[1] + assert diffusion.num_inference_steps == 32 + assert diffusion.guidance_scale == 2.0 + assert diffusion.seed == 0 + + def test_overrides(self): + encoder, diffusion = auk_sampling_params(nfe=4, cfg=0.0, seed=None) + assert diffusion.num_inference_steps == 4 + assert diffusion.guidance_scale == 0.0 + assert diffusion.seed is None + assert encoder.seed is None + + +class TestDurationPolicy: + @pytest.mark.parametrize( + ("gen_seconds", "expected"), + [ + (6.0, 300), + (0.02, 1), + (0.001, 1), + (5.59, 280), + (1.0, 50), + ], + ) + def test_explicit_duration_rounds_up(self, gen_seconds: float, expected: int): + assert resolve_gen_frames(gen_seconds, 0) == expected + + def test_explicit_duration_wins_over_the_source(self): + assert resolve_gen_frames(2.0, 500) == 100 + + def test_source_length_is_the_default(self): + assert resolve_gen_frames(None, 137) == 137 + + def test_text_only_without_a_duration_raises(self): + with pytest.raises(ValueError, match="gen_seconds is required"): + resolve_gen_frames(None, 0) + + @pytest.mark.parametrize("gen_seconds", [0.0, -1.0]) + def test_non_positive_duration_raises(self, gen_seconds: float): + with pytest.raises(ValueError, match="must be positive"): + resolve_gen_frames(gen_seconds, 100) + + def test_custom_rate(self): + assert resolve_gen_frames(1.0, 0, sample_rate=16000, hop=320) == 50 diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index 183ef865944..f0a1285bc77 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -83,6 +83,10 @@ "Gr00tN1d7Pipeline", "Pi0Pipeline", "SanaWmPipeline", + # Audio pipeline conditioned by a separate Qwen2.5-Omni encoder stage; + # covered by tests/diffusion/models/auk (CPU pipeline tests, parity) and + # the real-checkpoint E2E test instead of a tiny model. + "AuKPipeline", ] diff --git a/tools/auk_parity_reference.py b/tools/auk_parity_reference.py new file mode 100644 index 00000000000..33b66423277 --- /dev/null +++ b/tools/auk_parity_reference.py @@ -0,0 +1,227 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Generate the upstream AuK reference artifacts consumed by the parity tests. + +This runs the *upstream* implementation, so it must be executed in an +environment where the upstream package is installed +(``pip install -e /path/to/AuK``; it pins its own torch/transformers line and +does not need vLLM-Omni). It writes: + + /parity_ref//.pt, .wav, meta.json + Upstream generations for four cookbook cases (zero-shot TTS, instruct + TTS, content edit, denoise) with the reference latents, the generated + latents and the waveform. Point ``AUK_PARITY_REF`` at + ``/parity_ref/`` for tests/diffusion/models/auk/test_pipeline_auk.py. + /fusion_ref/zs_tts_audio.pt, instruct_text_only.pt + The fp32 layer-fused text condition from the HF thinker for the + zero-shot and text-only prompts (``AUK_FUSION_REF``). + +By default the VAE encode uses the posterior mean, matching the port's default +and making the artifacts deterministic; ``--vae-sample`` reproduces the +upstream stochastic draw (seeded through ``--global-seed``). The ODE noise is +drawn from ``--seed`` on both sides. + +Usage: + python tools/auk_parity_reference.py \\ + --auk-repo /path/to/AuK --ckpt-dir /path/to/AuK \\ + --qwen-dir /path/to/Qwen2.5-Omni-3B --out /path/to/auk-parity \\ + [--variant base|flash] [--vae-sample] [--fusion-device cpu|cuda] [--skip-fusion] +""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path +from typing import Any + +import torch +import torch.nn.functional as F + +WEIGHTS = {"base": "auk_base.safetensors", "flash": "auk_flash.safetensors"} + +# Cookbook cases; audio paths are relative to /assets/demo-input-audio. +CASES: dict[str, dict[str, Any]] = { + "zs_tts_en": { + "text": ( + "Say the following with the same voice: 'Ladies and gentlemen, it's an honor to have the " + "opportunity to address such a distinguished audience'" + ), + "audio": "zero-shot-tts/ref.wav", + "gen_seconds": 6.0, + }, + "instruct_en": { + "text": ( + 'Generate speech based on the following description: "A calm young woman speaking warmly and slowly.". ' + 'The content to speak is: "Welcome back, how was your day at work?".' + ), + "audio": None, + "gen_seconds": 3.5, + }, + "content_edit": { + "text": "Replace 'but accepting what we cannot have' with 'and living well with dreams unmet'.", + "audio": "content-edit/content.wav", + "gen_seconds": 7.0, + }, + "denoise": { + "text": "Keep pure speech voice, remove noise and reverberation.", + "audio": "vocal-extraction/vocal-1-input.wav", + "gen_seconds": None, + }, +} +# The two prompts the fusion reference covers (they must match the cases above). +FUSION_CASES = {"zs_tts_audio": "zs_tts_en", "instruct_text_only": "instruct_en"} +NO_PROMPT_AUDIO = "||" + + +def _messages(case: dict[str, Any], assets: Path) -> list[dict[str, Any]]: + content: list[dict[str, Any]] = [{"type": "text", "text": case["text"]}] + if case["audio"]: + content.append({"type": "audio", "audio": str(assets / case["audio"])}) + return [{"role": "user", "content": content}] + + +def _mean_encode(vae): + """Posterior-mean version of the upstream ``encoding_and_normalization``.""" + + def encode(sample, sample_lengths=None): + latent_stats = vae.audio_encoder(sample) + if sample_lengths is None: + sample_lengths = torch.LongTensor([sample.size(-1)] * sample.size(0)).to(sample.device) + latent_lens = sample_lengths // vae.hop_size + mean, _log_std = latent_stats.chunk(2, 1) + latents = mean.transpose(1, 2).float() + latents = (latents - vae.global_mean.float()) / torch.sqrt(vae.global_log_std.float()) + return latents, torch.clamp(latent_lens, max=latents.size(1)) + + return encode + + +def generate_references(args: argparse.Namespace) -> None: + from auk.infer.infer_auk import AukInfer, save_audio + + ckpt = Path(args.ckpt_dir) + out = Path(args.out) / "parity_ref" / args.variant + out.mkdir(parents=True, exist_ok=True) + assets = Path(args.auk_repo) / "assets" / "demo-input-audio" + + engine = AukInfer(str(ckpt / "config.yaml"), str(ckpt / WEIGHTS[args.variant]), qwen_path=args.qwen_dir) + encode = engine.vae_model.encoding_and_normalization if args.vae_sample else _mean_encode(engine.vae_model) + recorded: dict[str, torch.Tensor] = {} + + def encode_hook(sample, sample_lengths=None): + latents, lens = encode(sample, sample_lengths) + recorded["ref_latents"] = latents.detach().float().cpu() + recorded["ref_latent_lens"] = lens.detach().cpu() + return latents, lens + + engine.vae_model.encoding_and_normalization = encode_hook + upstream_sample = engine.model.sample + + def sample_hook(*a, **k): + latents, trajectory = upstream_sample(*a, **k) + recorded["generated"] = latents.detach().float().cpu() + return latents, trajectory + + engine.model.sample = sample_hook + + meta: dict[str, Any] = {} + for name, case in CASES.items(): + messages = _messages(case, assets) + recorded.clear() + torch.manual_seed(args.global_seed) + started = time.perf_counter() + audio, sr = engine.generate(messages, gen_seconds=case["gen_seconds"], seed=args.seed) + wall = time.perf_counter() - started + save_audio(audio, sr, str(out / f"{name}.wav")) + torch.save( + { + "messages": messages, + "gen_seconds": case["gen_seconds"], + "seed": args.seed, + "global_seed": args.global_seed, + "vae_sample": bool(args.vae_sample), + "audio": audio.float().cpu(), + "sr": sr, + **recorded, + }, + out / f"{name}.pt", + ) + ref_frames = int(recorded["ref_latent_lens"][0]) if "ref_latent_lens" in recorded else 0 + meta[name] = { + "wall_s": wall, + "out_s": audio.shape[-1] / sr, + "sr": sr, + "ref_frames": ref_frames, + "gen_frames": int(recorded["generated"].shape[1]) - ref_frames, + } + print(name, meta[name], flush=True) + (out / "meta.json").write_text(json.dumps(meta, indent=1) + "\n") + + +def generate_fusion(args: argparse.Namespace) -> None: + from auk.model.cfm_edit import CFMEdit + from safetensors import safe_open + from transformers import Qwen2_5OmniProcessor, Qwen2_5OmniThinkerForConditionalGeneration + + out = Path(args.out) / "fusion_ref" + out.mkdir(parents=True, exist_ok=True) + assets = Path(args.auk_repo) / "assets" / "demo-input-audio" + with safe_open(str(Path(args.ckpt_dir) / WEIGHTS[args.variant]), "pt") as f: + layer_weights = f.get_tensor("layer_weights").float() + layer_scale = f.get_tensor("layer_scale").float() + + processor = Qwen2_5OmniProcessor.from_pretrained(args.qwen_dir) + thinker = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(args.qwen_dir, torch_dtype=torch.float32) + thinker.visual = None + thinker = thinker.to(args.fusion_device).eval() + + for out_name, case_name in FUSION_CASES.items(): + case = dict(CASES[case_name]) + if not case["audio"] and not case["text"].endswith(NO_PROMPT_AUDIO): + case["text"] = case["text"] + NO_PROMPT_AUDIO + cond = CFMEdit.build_cond_inputs([_messages(case, assets)], processor).to(args.fusion_device) + with torch.no_grad(): + hidden_states = thinker(**cond, output_hidden_states=True).hidden_states + width = hidden_states[0].shape[-1] + stacked = torch.stack([F.layer_norm(h.float(), [width]) for h in hidden_states[1:]], 0) + weights = F.softmax(layer_weights.to(stacked.device), 0) + fused = (stacked * weights[:, None, None, None]).sum(0) * layer_scale.to(stacked.device) + torch.save( + { + "input_ids": cond["input_ids"].cpu(), + "attention_mask": cond["attention_mask"].cpu(), + "fused": fused[0].float().cpu(), + "final_normed": hidden_states[-1][0].float().cpu(), + "text": case["text"], + }, + out / f"{out_name}.pt", + ) + print(out_name, {"tokens": int(cond["input_ids"].shape[1]), "fused_norm": float(fused.norm())}, flush=True) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--auk-repo", required=True, help="checkout of the upstream AuK repository (for demo clips)") + parser.add_argument("--ckpt-dir", required=True, help="tencent/AuK or tencent/AuK-Flash snapshot") + parser.add_argument("--qwen-dir", required=True, help="Qwen/Qwen2.5-Omni-3B snapshot") + parser.add_argument("--out", required=True, help="output root; writes parity_ref/ and fusion_ref") + parser.add_argument("--variant", choices=sorted(WEIGHTS), default="base") + parser.add_argument("--seed", type=int, default=0, help="ODE noise seed passed to the upstream sampler") + parser.add_argument("--global-seed", type=int, default=1234, help="global RNG seed set before each generate()") + parser.add_argument("--vae-sample", action="store_true", help="draw the VAE posterior sample as upstream does") + parser.add_argument("--fusion-device", default="cpu", help="device for the fp32 HF thinker forward") + parser.add_argument("--skip-fusion", action="store_true", help="only write the generation references") + args = parser.parse_args() + + generate_references(args) + if not args.skip_fusion: + generate_fusion(args) + print(f"reference artifacts written under {args.out}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/prepare_auk_checkpoint.py b/tools/prepare_auk_checkpoint.py new file mode 100755 index 00000000000..7f55dab8a70 --- /dev/null +++ b/tools/prepare_auk_checkpoint.py @@ -0,0 +1,369 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Assemble an AuK checkpoint directory that ``--model `` can load. + +The released AuK repos ship ``config.yaml`` plus loose safetensors and depend +on a separate Qwen2.5-Omni-3B snapshot for the frozen text encoder. This tool +writes the ``config.json`` the vLLM-Omni pipeline expects and links every +artifact both stages need into one flat directory: + + / + config.json the Qwen2.5-Omni-3B config with model_type "auk", + AuK's architecture, and the dit / vae / variant / + flash_t_grid / defaults sections added at the top + level (see transformers_utils/configs/auk.py) + auk.safetensors -> auk_base.safetensors or auk_flash.safetensors + vae.safetensors -> the AuK VAE weights + thinker/ -> the Qwen2.5-Omni-3B snapshot, for reference only + model-0000*.safetensors, model.safetensors.index.json, + tokenizer*, vocab.json, merges.txt, added_tokens.json, + special_tokens_map.json, chat_template.json, + preprocessor_config.json, generation_config.json, spk_dict.pt + -> the same snapshot's entries + +The Qwen entries are linked at the top level rather than left under +``thinker/`` because both stages load from this one directory: the encoder +reads the thinker shards, the tokenizer and the audio processor from here, +and takes the layer-fusion tensors out of ``auk.safetensors``. The weight +index keeps ``auk.safetensors`` and ``vae.safetensors`` out of the thinker's +own weight load, and the DiT and VAE weights are read by model code directly. + +Usage: + python tools/prepare_auk_checkpoint.py \ + --auk-dir /path/to/AuK \ + --qwen-dir /path/to/Qwen2.5-Omni-3B \ + --out /path/to/auk-omni-base + +Re-running against an existing directory is safe: config.json is rewritten +and the links are repointed. Pass ``--copy`` on filesystems without symlinks. +""" + +from __future__ import annotations + +import argparse +import ast +import json +import re +import shutil +import sys +from pathlib import Path +from typing import Any + +# The distilled student's 4-step grid. Duplicated from +# vllm_omni/transformers_utils/configs/auk.py so this tool stays runnable +# without importing the package (and its torch dependency). +FLASH_T_GRID = [ + 0.0, + 0.07612049579620361, + 0.2928932309150696, + 0.6173166036605835, + 1.0, +] + +# Names the assembled directory owns; everything else in it is a Qwen entry. +WEIGHTS_LINK = "auk.safetensors" +VAE_LINK = "vae.safetensors" +THINKER_LINK = "thinker" +CONFIG_NAME = "config.json" + +# Snapshot entries that are not reused: the config is rewritten, and the repo +# paperwork would misdescribe the assembled directory. +_SKIP_QWEN_ENTRIES = frozenset({CONFIG_NAME, "README.md", "LICENSE"}) + +DIT_KEYS = ("dim", "heads", "ff_mult", "text_hidden_dim", "num_layers", "num_single_layers") +VAE_KEYS = ("latent_dim", "downsample_rate", "target_sample_rate", "model_init_kwargs") + +_INTERPOLATION = re.compile(r"^\$\{([A-Za-z0-9_.]+)\}$") + + +def _strip_comment(line: str) -> str: + """Drop a trailing ``#`` comment that is not inside a quoted scalar.""" + out: list[str] = [] + quote: str | None = None + for ch in line: + if quote is not None: + if ch == quote: + quote = None + out.append(ch) + continue + if ch in "\"'": + quote = ch + out.append(ch) + continue + if ch == "#": + break + out.append(ch) + return "".join(out).rstrip() + + +def _parse_scalar(text: str) -> Any: + """Parse one YAML scalar or inline sequence.""" + text = text.strip() + if not text: + return None + lowered = text.lower() + if lowered in ("true", "yes"): + return True + if lowered in ("false", "no"): + return False + if lowered in ("null", "~"): + return None + if text[0] in "[{": + normalized = re.sub(r"\b(True|true|yes)\b", "True", text) + normalized = re.sub(r"\b(False|false|no)\b", "False", normalized) + normalized = re.sub(r"\b(null|~)\b", "None", normalized) + try: + return ast.literal_eval(normalized) + except (SyntaxError, ValueError): + return text + try: + return ast.literal_eval(text) + except (SyntaxError, ValueError): + return text.strip("\"'") + + +def _parse_yaml_minimal(text: str) -> dict[str, Any]: + """Parse the small subset of YAML the AuK configs use. + + Nested mappings by indentation, scalars, and inline sequences. Used only + when PyYAML is unavailable. + """ + root: dict[str, Any] = {} + stack: list[tuple[int, dict[str, Any]]] = [(-1, root)] + for raw in text.splitlines(): + line = _strip_comment(raw) + if not line.strip(): + continue + indent = len(line) - len(line.lstrip(" ")) + body = line.strip() + if ":" not in body: + raise ValueError(f"unsupported YAML line (no key): {raw!r}") + key, _, rest = body.partition(":") + key = key.strip() + while stack and indent <= stack[-1][0]: + stack.pop() + if not stack: + raise ValueError(f"unsupported YAML indentation at: {raw!r}") + parent = stack[-1][1] + if rest.strip(): + parent[key] = _parse_scalar(rest) + else: + child: dict[str, Any] = {} + parent[key] = child + stack.append((indent, child)) + return root + + +def _resolve_interpolations(node: Any, root: dict[str, Any]) -> Any: + """Replace ``${a.b.c}`` scalars with the value at that dotted path.""" + if isinstance(node, dict): + return {k: _resolve_interpolations(v, root) for k, v in node.items()} + if isinstance(node, list): + return [_resolve_interpolations(v, root) for v in node] + if isinstance(node, str): + match = _INTERPOLATION.match(node.strip()) + if match is None: + return node + target: Any = root + for part in match.group(1).split("."): + if not isinstance(target, dict) or part not in target: + raise ValueError(f"cannot resolve interpolation ${{{match.group(1)}}}") + target = target[part] + return target + return node + + +def load_auk_yaml(path: Path) -> dict[str, Any]: + """Load the AuK ``config.yaml`` with interpolations resolved.""" + text = path.read_text() + try: + import yaml + + parsed = yaml.safe_load(text) + except ImportError: + parsed = _parse_yaml_minimal(text) + if not isinstance(parsed, dict) or "model" not in parsed: + raise ValueError(f"{path} has no top-level 'model' section") + return _resolve_interpolations(parsed, parsed) + + +def detect_variant(auk_dir: Path, model_section: dict[str, Any]) -> str: + """Infer base vs flash from the weight filename, then the model name.""" + if (auk_dir / "auk_flash.safetensors").exists(): + return "flash" + if (auk_dir / "auk_base.safetensors").exists(): + return "base" + name = str(model_section.get("name", "")).lower() + return "flash" if "flash" in name else "base" + + +def weights_name(variant: str) -> str: + return "auk_flash.safetensors" if variant == "flash" else "auk_base.safetensors" + + +def build_config( + auk_yaml: dict[str, Any], + qwen_config: dict[str, Any], + variant: str, +) -> dict[str, Any]: + """Assemble the ``config.json`` body for the output directory. + + The Qwen config is kept whole, including ``talker_config`` and + ``token2wav_config``: ``AuKConfig`` subclasses ``Qwen2_5OmniConfig`` and + the encoder stage reuses the in-tree thinker processor, which asserts that + type. AuK's own sections are added at the top level. + """ + model = auk_yaml["model"] + arch = model.get("arch", {}) + vae = model.get("vae", {}) + + missing = [key for key in DIT_KEYS if key not in arch] + if missing: + raise ValueError(f"AuK config.yaml arch section is missing {missing}") + missing = [key for key in VAE_KEYS if key not in vae] + if missing: + raise ValueError(f"AuK config.yaml vae section is missing {missing}") + if not isinstance(qwen_config.get("thinker_config"), dict): + raise ValueError("Qwen snapshot config.json has no 'thinker_config' mapping") + + dit = {key: arch[key] for key in DIT_KEYS} + # dim_head is implicit upstream (dim // heads) but the port takes it + # explicitly, so it is recorded rather than recomputed at load time. + dit["dim_head"] = int(arch["dim"]) // int(arch["heads"]) + # Carried from the released config: the DiT builds an attention mask over + # the padded [audio | text] sequence and parity depends on it. + dit["attn_mask_enabled"] = bool(arch.get("attn_mask_enabled", True)) + + config = dict(qwen_config) + config["model_type"] = "auk" + config["architectures"] = ["AuKForConditionalGeneration"] + config["dit"] = dit + config["vae"] = {key: vae[key] for key in VAE_KEYS} + config["variant"] = variant + config["flash_t_grid"] = FLASH_T_GRID + config["defaults"] = {"nfe": 32, "cfg": 2.0, "sway": -1.0} + return config + + +def _replace_path(target: Path) -> None: + """Remove an existing file, directory, or link at ``target``.""" + if target.is_symlink() or target.is_file(): + target.unlink() + elif target.is_dir(): + shutil.rmtree(target) + + +def link_or_copy(source: Path, target: Path, *, copy: bool) -> None: + """Point ``target`` at ``source``, replacing whatever is already there.""" + if not source.exists(): + raise FileNotFoundError(f"missing source artifact: {source}") + if target.is_symlink() and not copy and target.readlink() == source.resolve(): + return + _replace_path(target) + if copy: + if source.is_dir(): + shutil.copytree(source, target) + else: + shutil.copy2(source, target) + return + target.symlink_to(source.resolve(), target_is_directory=source.is_dir()) + + +MANIFEST_NAME = ".auk-assembled.json" + + +def _reject_overlap(out: Path, *inputs: Path) -> None: + """Refuse an output directory that is, contains, or lives inside an input.""" + out_r = out.resolve() + for src in inputs: + src_r = src.resolve() + if out_r == src_r or out_r in src_r.parents or src_r in out_r.parents: + raise ValueError(f"--out {out} overlaps input directory {src}; choose a separate directory") + + +def assemble(auk_dir: Path, qwen_dir: Path, out: Path, variant: str, *, copy: bool) -> None: + """Link the AuK weights and the Qwen snapshot entries into ``out``.""" + _reject_overlap(out, auk_dir, qwen_dir) + out.mkdir(parents=True, exist_ok=True) + manifest_path = out / MANIFEST_NAME + previous: set[str] = set() + if manifest_path.is_file(): + previous = set(json.loads(manifest_path.read_text()).get("entries", [])) + weights = auk_dir / weights_name(variant) + + entries: dict[str, Path] = { + WEIGHTS_LINK: weights, + VAE_LINK: auk_dir / VAE_LINK, + THINKER_LINK: qwen_dir, + } + # The assembled config replaces the Qwen one; every other entry is reused. + for entry in qwen_dir.iterdir(): + if entry.name.startswith(".") or entry.name in _SKIP_QWEN_ENTRIES: + continue + entries[entry.name] = entry + + for name, source in entries.items(): + link_or_copy(source, out / name, copy=copy) + + # Drop only entries this tool created in an earlier run and no longer needs; + # anything else in the directory belongs to the user. + for name in previous - set(entries) - {CONFIG_NAME}: + stale = out / name + if stale.exists() or stale.is_symlink(): + _replace_path(stale) + manifest_path.write_text(json.dumps({"entries": sorted(entries)}, indent=1) + "\n") + + +def print_tree(out: Path) -> None: + """Print the assembled directory, one entry per line.""" + print(f"\n{out}/") + for entry in sorted(out.iterdir()): + if entry.is_symlink(): + print(f" {entry.name} -> {entry.readlink()}") + elif entry.is_dir(): + print(f" {entry.name}/") + else: + print(f" {entry.name} ({entry.stat().st_size / 1e9:.2f} GB)") + + +def main() -> int: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--auk-dir", required=True, help="HF snapshot of the AuK or AuK-Flash repo") + parser.add_argument("--qwen-dir", required=True, help="HF snapshot of Qwen2.5-Omni-3B") + parser.add_argument("--out", required=True, help="Directory to assemble (created if absent)") + parser.add_argument( + "--variant", + choices=("base", "flash"), + help="Override the variant; detected from the AuK snapshot by default", + ) + parser.add_argument( + "--copy", + action="store_true", + help="Copy the weights and the snapshot entries instead of symlinking", + ) + args = parser.parse_args() + + auk_dir = Path(args.auk_dir).expanduser().resolve() + qwen_dir = Path(args.qwen_dir).expanduser().resolve() + out = Path(args.out).expanduser().resolve() + _reject_overlap(out, auk_dir, qwen_dir) + + auk_yaml = load_auk_yaml(auk_dir / "config.yaml") + variant = args.variant or detect_variant(auk_dir, auk_yaml["model"]) + qwen_config = json.loads((qwen_dir / CONFIG_NAME).read_text()) + + config = build_config(auk_yaml, qwen_config, variant) + assemble(auk_dir, qwen_dir, out, variant, copy=args.copy) + (out / CONFIG_NAME).write_text(json.dumps(config, indent=2) + "\n") + + print(f"variant: {variant} (weights {weights_name(variant)})") + print_tree(out) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index a69fb08c65b..2378ce5b9ae 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -40,6 +40,7 @@ AUDEX_TTS_PIPELINE, ) from vllm_omni.model_executor.models.audio8_tts.pipeline import AUDIO8_TTS_PIPELINE +from vllm_omni.model_executor.models.auk.pipeline import AUK_PIPELINE from vllm_omni.model_executor.models.aura_omni.pipeline import AURA_OMNI_PIPELINE from vllm_omni.model_executor.models.bagel.pipeline import ( BAGEL_PIPELINE, @@ -164,6 +165,7 @@ "wan2_2_ti2v": WAN2_2_TI2V_PIPELINE, "voxcpm2": VOXCPM2_PIPELINE, "dots_tts": DOTS_TTS_PIPELINE, + "auk": AUK_PIPELINE, "cosyvoice3": COSYVOICE3_PIPELINE, "audex_tts": AUDEX_TTS_PIPELINE, "audex_tta": AUDEX_TTA_PIPELINE, diff --git a/vllm_omni/deploy/auk.yaml b/vllm_omni/deploy/auk.yaml new file mode 100644 index 00000000000..bbe544919ab --- /dev/null +++ b/vllm_omni/deploy/auk.yaml @@ -0,0 +1,51 @@ +# AuK deploy: thinker encoder (stage 0) + rectified-flow DiT (stage 1). +# Topology declared in vllm_omni/model_executor/models/auk/pipeline.py. +# +# Fields omitted from a stage fall back to StageDeployConfig dataclass +# defaults (see vllm_omni/config/stage_config.py). trust_remote_code, +# enable_prefix_caching and enable_chunked_prefill are pipeline-wide +# DeployConfig fields and apply to both stages. +# +# The ODE runs over the whole target segment, so there is nothing to stream +# between the stages: async_chunk stays off. +async_chunk: false +trust_remote_code: true +# Prefix caching must stay off: a cache hit skips the prefix positions, and +# the fused [tokens, 2048] condition needs every position of the prompt. +# Chunked prefill works (the runner concatenates the per-chunk hidden states) +# but measured no wall-time gain, since the encoder is a few percent of a +# request, and it costs bit-exact reproducibility across scheduling. +enable_prefix_caching: false +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. + - stage_id: 0 + devices: "0" + max_num_seqs: 1 + gpu_memory_utilization: 0.45 + enforce_eager: true + 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} + default_sampling_params: + temperature: 0.0 + max_tokens: 1 + detokenize: false + seed: 0 + + # Stage 1: AuKPipeline (DiT + BigVGAN-flow VAE encoder and decoder). + # Shares GPU 0 with the encoder; 32 Euler steps at CFG 2.0 is the Base + # recipe, and AuK-Flash overrides both with its distilled 4-step grid. + - stage_id: 1 + devices: "0" + max_num_seqs: 1 + gpu_memory_utilization: 0.35 + enforce_eager: true + default_sampling_params: + seed: 0 + num_inference_steps: 32 + guidance_scale: 2.0 diff --git a/vllm_omni/diffusion/models/auk/__init__.py b/vllm_omni/diffusion/models/auk/__init__.py new file mode 100644 index 00000000000..fd987d42ba2 --- /dev/null +++ b/vllm_omni/diffusion/models/auk/__init__.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""AuK diffusion-stage components: the flow transformer, the codec, the pipeline.""" + +from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, dit_state_dict, sample_latents +from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE +from vllm_omni.diffusion.models.auk.pipeline_auk import AuKPipeline, get_auk_post_process_func + +__all__ = [ + "AuKPipeline", + "AuKTransformer", + "AuKVAE", + "dit_state_dict", + "get_auk_post_process_func", + "sample_latents", +] diff --git a/vllm_omni/diffusion/models/auk/auk_transformer.py b/vllm_omni/diffusion/models/auk/auk_transformer.py new file mode 100644 index 00000000000..5aca8f39b82 --- /dev/null +++ b/vllm_omni/diffusion/models/auk/auk_transformer.py @@ -0,0 +1,677 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The AuK audio DiT as a self-contained inference module. + +Two phases. ``num_layers`` double-stream blocks attend jointly over text and +audio while keeping each stream's residual separate, then +``num_single_layers`` single-stream blocks run over the concatenated +``[text | ref audio | target audio]`` sequence. Only the target-audio slice is +projected back to latent space, so the module predicts a flow-matching +velocity for the target frames and :func:`sample_latents` integrates it. + +Parameter names and shapes match the AuK reference backbone, so a released +checkpoint loads with ``load_state_dict(strict=True)`` once the +``transformer.`` prefix is stripped (see :func:`dit_state_dict`). + +Deliberate differences from the reference, all inference-only: + +* ``x_transformers`` and ``torchdiffeq`` are gone. Rotary embeddings are + computed here (interleaved GPT-J pairs, no xpos scaling) and the ODE is an + explicit Euler loop. +* Dropout, activation checkpointing, the ``flash_attn`` backend switch and the + zero-init of the modulation layers are training concerns and are dropped. +* The timestep sinusoid is cast to the time MLP's weight dtype rather than to + the timestep's own dtype, so an fp32 timestep works against half-precision + weights outside autocast. +* ``attn_mask_enabled`` defaults to ``True``, the value the released config + sets, rather than the reference's ``False``. +""" + +import math +from collections.abc import Iterable, Mapping + +import torch +import torch.nn.functional as F +from torch import nn + +__all__ = ["AuKTransformer", "dit_state_dict", "sample_latents"] + + +def _sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor: + """The single attention call site, so the backend can be swapped in one place. + + ``q``/``k``/``v`` are ``[B, H, N, D]``. ``mask`` is a boolean key-padding + mask ``[B, K]`` in which ``True`` marks a valid key; it is broadcast to + ``[B, H, Q, K]`` only when given. + """ + attn_mask = None + if mask is not None: + attn_mask = mask[:, None, None, :].expand(q.shape[0], q.shape[1], q.shape[-2], k.shape[-2]) + return F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + """Rotate interleaved GPT-J pairs: ``(x1, x2) -> (-x2, x1)``.""" + pairs = x.reshape(*x.shape[:-1], -1, 2) + x1, x2 = pairs.unbind(dim=-1) + return torch.stack((-x2, x1), dim=-1).flatten(-2) + + +def _apply_rope(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """Apply rotary frequencies ``[1, 1, N, D]`` to ``x`` ``[B, H, N, D]`` in fp32.""" + out_dtype = x.dtype + x32 = x.float() + out = x32 * freqs.cos() + _rotate_half(x32) * freqs.sin() + return out.to(out_dtype) + + +class Rotary(nn.Module): + """Rotary position frequencies for one stream, evaluated in fp32. + + The frequencies live in a buffer so that a released checkpoint's own copy + of them loads with ``strict=True``. They are recomputed whenever the + module has been cast below fp32: rounding either the frequencies or the + integer positions to bf16 wrecks the phase at long positions. + """ + + def __init__(self, dim: int, base: float = 10000.0) -> None: + super().__init__() + if dim % 2 != 0: + raise ValueError(f"rotary dim must be even, got {dim}") + self.dim = dim + self.base = base + self.register_buffer("inv_freq", self._frequencies(torch.device("cpu"))) + + def _frequencies(self, device: torch.device) -> torch.Tensor: + exponents = torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim + return 1.0 / (self.base**exponents) + + def forward(self, seq_len: int) -> torch.Tensor: + """Return frequencies ``[1, 1, seq_len, dim]`` for positions ``0..seq_len-1``.""" + inv_freq = self.inv_freq + if inv_freq.dtype != torch.float32: + inv_freq = self._frequencies(inv_freq.device) + pos = torch.arange(seq_len, device=inv_freq.device, dtype=torch.float32) + freqs = torch.outer(pos, inv_freq) + return torch.stack((freqs, freqs), dim=-1).flatten(-2)[None, None] + + +class TimeEmbedding(nn.Module): + """Sinusoidal timestep features followed by a two-layer MLP.""" + + def __init__(self, dim: int, freq_embed_dim: int = 256) -> None: + super().__init__() + self.freq_embed_dim = freq_embed_dim + self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + + def forward(self, timestep: torch.Tensor, scale: float = 1000.0) -> torch.Tensor: + half = self.freq_embed_dim // 2 + decay = math.log(10000) / (half - 1) + freqs = torch.exp(torch.arange(half, device=timestep.device, dtype=torch.float32) * -decay) + angles = scale * timestep.unsqueeze(1) * freqs.unsqueeze(0) + hidden = torch.cat((angles.sin(), angles.cos()), dim=-1) + return self.time_mlp(hidden.to(self.time_mlp[0].weight.dtype)) + + +class ConvPosEmbedding(nn.Module): + """Depthwise-grouped conv stack that adds local position information.""" + + def __init__(self, dim: int, kernel_size: int = 31, groups: int = 16) -> None: + super().__init__() + if kernel_size % 2 == 0: + raise ValueError(f"kernel_size must be odd, got {kernel_size}") + padding = kernel_size // 2 + self.conv1d = nn.Sequential( + nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=padding), + nn.Mish(), + nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=padding), + nn.Mish(), + ) + + 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) + 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) + + +class AudioEmbedding(nn.Module): + """Project audio latents to model width and add conv position information.""" + + def __init__(self, latent_dim: int, dim: int) -> None: + super().__init__() + self.linear = nn.Linear(latent_dim, dim) + self.conv_pos_embed = ConvPosEmbedding(dim) + + def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor: + x = self.linear(x) + return self.conv_pos_embed(x, mask=mask) + x + + +class AdaLayerNorm(nn.Module): + """Timestep-conditioned modulation for a block: six chunks from one projection.""" + + def __init__(self, dim: int) -> None: + super().__init__() + self.silu = nn.SiLU() + self.linear = nn.Linear(dim, dim * 6) + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + def forward( + self, x: torch.Tensor, emb: torch.Tensor + ) -> 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) + x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] + return x, gate_msa, shift_mlp, scale_mlp, gate_mlp + + +class AdaLayerNormFinal(nn.Module): + """Timestep-conditioned modulation before the output projection.""" + + def __init__(self, dim: int) -> None: + super().__init__() + self.silu = nn.SiLU() + 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) + return self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] + + +class FeedForward(nn.Module): + """SwiGLU feed-forward with the gate projection fused into ``linear_in``.""" + + def __init__(self, dim: int, mult: float = 4.0) -> None: + super().__init__() + inner_dim = int(dim * mult) + self.linear_in = nn.Linear(dim, inner_dim * 2, bias=False) + self.linear_out = nn.Linear(inner_dim, dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate, up = self.linear_in(x).chunk(2, dim=-1) + return self.linear_out(F.silu(gate) * up) + + +class Attention(nn.Module): + """Self-attention with per-head QK RMSNorm and rotary positions. + + ``to_out`` stays a ``ModuleList`` because the released checkpoints name the + output projection ``to_out.0``; the reference's trailing dropout carries no + parameters and is dropped. + """ + + def __init__(self, dim: int, heads: int, dim_head: int, attn_mask_enabled: bool = True) -> None: + super().__init__() + self.heads = heads + self.dim_head = dim_head + self.attn_mask_enabled = attn_mask_enabled + inner_dim = heads * dim_head + self.to_qkv = nn.Linear(dim, 3 * inner_dim) + self.q_norm = nn.RMSNorm(dim_head, elementwise_affine=True) + self.k_norm = nn.RMSNorm(dim_head, elementwise_affine=True) + self.to_out = nn.ModuleList([nn.Linear(inner_dim, dim)]) + + def _merge_heads(self, x: torch.Tensor) -> torch.Tensor: + return x.transpose(1, 2).reshape(x.shape[0], -1, self.heads * self.dim_head) + + def _qkv( + self, + packed: torch.Tensor, + q_norm: nn.Module, + k_norm: nn.Module, + rope: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Split a packed QKV projection into heads, RMSNorm Q/K, then rotate.""" + q, k, v = (t.view(t.shape[0], -1, self.heads, self.dim_head).transpose(1, 2) for t in packed.chunk(3, dim=-1)) + q = q_norm(q) + k = k_norm(k) + if rope is not None: + q = _apply_rope(q, rope) + k = _apply_rope(k, rope) + return q, k, v + + def _project(self, x: torch.Tensor, rope: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return self._qkv(self.to_qkv(x), self.q_norm, self.k_norm, rope) + + def forward( + self, x: torch.Tensor, mask: torch.Tensor | None = None, rope: torch.Tensor | None = None + ) -> torch.Tensor: + q, k, v = self._project(x, rope) + attn_mask = mask if self.attn_mask_enabled else None + out = self._merge_heads(_sdpa(q, k, v, attn_mask).to(q.dtype)) + out = self.to_out[0](out) + if mask is not None: + out = out.masked_fill(~mask.unsqueeze(-1), 0.0) + return out + + +class JointAttention(Attention): + """Attention over the audio and text streams jointly, with separate projections.""" + + def __init__(self, dim: int, heads: int, dim_head: int, context_dim: int, attn_mask_enabled: bool = True) -> None: + super().__init__(dim, heads, dim_head, attn_mask_enabled) + inner_dim = heads * dim_head + self.to_qkv_c = nn.Linear(context_dim, 3 * inner_dim) + self.c_q_norm = nn.RMSNorm(dim_head, elementwise_affine=True) + self.c_k_norm = nn.RMSNorm(dim_head, elementwise_affine=True) + self.to_out_c = nn.Linear(inner_dim, context_dim) + + def _project_context(self, c: torch.Tensor, c_rope: torch.Tensor | None) -> tuple[torch.Tensor, ...]: + return self._qkv(self.to_qkv_c(c), self.c_q_norm, self.c_k_norm, c_rope) + + def forward( # type: ignore[override] # joint attention takes the context stream too + self, + x: torch.Tensor, + c: torch.Tensor, + mask: torch.Tensor | None = None, + rope: torch.Tensor | None = None, + c_rope: torch.Tensor | None = None, + c_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + audio_len = x.shape[1] + q, k, v = self._project(x, rope) + c_q, c_k, c_v = self._project_context(c, c_rope) + + q = torch.cat([q, c_q], dim=2) + k = torch.cat([k, c_k], dim=2) + v = torch.cat([v, c_v], dim=2) + + joint_mask = None + if self.attn_mask_enabled and mask is not None: + if c_mask is not None: + joint_mask = torch.cat([mask, c_mask], dim=1) + else: + joint_mask = F.pad(mask, (0, c.shape[1]), value=True) + + out = self._merge_heads(_sdpa(q, k, v, joint_mask).to(q.dtype)) + x_out = self.to_out[0](out[:, :audio_len]) + c_out = self.to_out_c(out[:, audio_len:]) + + if mask is not None: + x_out = x_out.masked_fill(~mask.unsqueeze(-1), 0.0) + if c_mask is not None: + c_out = c_out.masked_fill(~c_mask.unsqueeze(-1), 0.0) + return x_out, c_out + + +class DoubleBlock(nn.Module): + """MM-DiT block: joint attention, separate modulation and feed-forward per stream.""" + + def __init__(self, dim: int, heads: int, dim_head: int, ff_mult: float, attn_mask_enabled: bool) -> None: + super().__init__() + self.attn_norm_c = AdaLayerNorm(dim) + self.attn_norm_x = AdaLayerNorm(dim) + self.attn = JointAttention( + dim=dim, heads=heads, dim_head=dim_head, context_dim=dim, attn_mask_enabled=attn_mask_enabled + ) + self.ff_norm_c = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff_c = FeedForward(dim=dim, mult=ff_mult) + self.ff_norm_x = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff_x = FeedForward(dim=dim, mult=ff_mult) + + def forward( + self, + x: torch.Tensor, + c: torch.Tensor, + t: torch.Tensor, + mask: torch.Tensor | None, + rope: torch.Tensor, + c_rope: torch.Tensor, + c_mask: torch.Tensor | 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) + + x_attn, c_attn = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope, c_mask=c_mask) + + c = c + c_gate_msa[:, None] * c_attn + norm_c = self.ff_norm_c(c) * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None] + c = c + c_gate_mlp[:, None] * self.ff_c(norm_c) + + x = x + x_gate_msa[:, None] * x_attn + norm_x = self.ff_norm_x(x) * (1 + x_scale_mlp[:, None]) + x_shift_mlp[:, None] + x = x + x_gate_mlp[:, None] * self.ff_x(norm_x) + return c, x + + +class SingleBlock(nn.Module): + """DiT block over the concatenated text and audio sequence.""" + + def __init__(self, dim: int, heads: int, dim_head: int, ff_mult: float, attn_mask_enabled: bool) -> None: + super().__init__() + self.attn_norm = AdaLayerNorm(dim) + self.attn = Attention(dim=dim, heads=heads, dim_head=dim_head, attn_mask_enabled=attn_mask_enabled) + self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + self.ff = FeedForward(dim=dim, mult=ff_mult) + + def forward(self, x: torch.Tensor, t: torch.Tensor, mask: torch.Tensor | None, rope: torch.Tensor) -> torch.Tensor: + norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, t) + x = x + gate_msa[:, None] * self.attn(x=norm, mask=mask, rope=rope) + norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + return x + gate_mlp[:, None] * self.ff(norm) + + +class AuKTransformer(nn.Module): + """Flow-matching velocity predictor for AuK audio latents. + + Args: + dim: Model width. + heads: Attention heads. + dim_head: Per-head width. + ff_mult: Feed-forward expansion factor. + latent_dim: Audio VAE latent channels, the input and output width. + text_hidden_dim: Width of the pre-encoded text hidden states. + num_layers: Double-stream (joint attention) block count. + num_single_layers: Single-stream block count. + attn_mask_enabled: Apply padding masks inside attention. When ``False`` + attention runs unmasked, though padded outputs are still zeroed, + which is what the reference does with the flag off. + """ + + def __init__( + self, + *, + dim: int, + heads: int, + dim_head: int, + ff_mult: float, + latent_dim: int, + text_hidden_dim: int, + num_layers: int, + num_single_layers: int, + attn_mask_enabled: bool = True, + ) -> None: + super().__init__() + self.dim = dim + self.latent_dim = latent_dim + self.text_cond: torch.Tensor | None = None + self.text_uncond: torch.Tensor | None = None + + self.time_embed = TimeEmbedding(dim) + self.txt_norm = nn.RMSNorm(dim, elementwise_affine=True) + self.txt_proj = nn.Linear(text_hidden_dim, dim) + self.audio_embed = AudioEmbedding(latent_dim, dim) + self.rotary_embed = Rotary(dim_head) + + self.transformer_blocks = nn.ModuleList( + [ + DoubleBlock( + dim=dim, heads=heads, dim_head=dim_head, ff_mult=ff_mult, attn_mask_enabled=attn_mask_enabled + ) + for _ in range(num_layers) + ] + ) + self.single_transformer_blocks = nn.ModuleList( + [ + SingleBlock( + dim=dim, heads=heads, dim_head=dim_head, ff_mult=ff_mult, attn_mask_enabled=attn_mask_enabled + ) + for _ in range(num_single_layers) + ] + ) + + self.norm_out = AdaLayerNormFinal(dim) + self.proj_out = nn.Linear(dim, latent_dim) + + def project_text(self, text: torch.Tensor) -> torch.Tensor: + """Project LLM hidden states ``[B, nt, text_hidden_dim]`` to model width.""" + return self.txt_norm(self.txt_proj(text)) + + def clear_cache(self) -> None: + """Drop the cached text projection. Call this when the conditioning changes.""" + self.text_cond = None + self.text_uncond = None + + def _embed_audio( + self, + x: torch.Tensor, + ref: torch.Tensor | None, + drop_audio_cond: bool, + mask: torch.Tensor | None, + ref_mask: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None, int]: + """Embed the target latents, prepending the reference prompt when present. + + Returns the ``[ref | target]`` sequence, its padding mask and the + reference length, which the caller uses to slice the target back out. + """ + if ref is not None and ref.shape[1] == 0: + ref = None + + target = self.audio_embed(x, mask=mask) + if ref is None: + return target, mask, 0 + + if drop_audio_cond: + ref = torch.zeros_like(ref) + prompt = self.audio_embed(ref, mask=ref_mask) + audio = torch.cat([prompt, target], dim=1) + prompt_len = prompt.shape[1] + + audio_mask = None + if mask is not None or ref_mask is not None: + if mask is None: + mask = x.new_ones(x.shape[:2], dtype=torch.bool) + if ref_mask is None: + ref_mask = ref.new_ones(ref.shape[:2], dtype=torch.bool) + audio_mask = torch.cat([ref_mask, mask], dim=1) + return audio, audio_mask, prompt_len + + def forward( + self, + x: torch.Tensor, + text: torch.Tensor, + time: torch.Tensor, + *, + mask: torch.Tensor | None = None, + c_mask: torch.Tensor | None = None, + ref: torch.Tensor | None = None, + ref_mask: torch.Tensor | None = None, + drop_audio_cond: bool = False, + drop_text: bool = False, + cfg_infer: bool = False, + cache: bool = False, + ) -> torch.Tensor: + """Predict the flow-matching velocity for the target frames of ``x``. + + Args: + x: Noised target latents ``[B, n, latent_dim]``. + text: Pre-encoded text hidden states ``[B, nt, text_hidden_dim]``. + time: Flow-matching timestep, a scalar or ``[B]``. + mask: Target padding mask ``[B, n]``, ``True`` where valid. + c_mask: Text padding mask ``[B, nt]``. Inferred from all-zero text + rows when omitted. + ref: Reference prompt latents ``[B, np, latent_dim]``, prepended to + the audio sequence. A zero-length ``ref`` counts as absent. + ref_mask: Reference padding mask ``[B, np]``. + drop_audio_cond: Zero the reference prompt, for the CFG uncond branch. + drop_text: Zero the projected text, for the CFG uncond branch. + cfg_infer: Run the cond and uncond branches as one batch of ``2B``, + cond first. ``drop_audio_cond`` and ``drop_text`` are then + driven per branch and ignored. + cache: Reuse the projected text across calls, for an ODE loop over + fixed conditioning. Only the ``cfg_infer`` path caches, as in + the reference. Call :meth:`clear_cache` when the text changes. + + Returns: + Velocity ``[B, n, latent_dim]``, or ``[2B, n, latent_dim]`` under + ``cfg_infer``. + """ + batch = x.shape[0] + if time.ndim == 0: + time = time.repeat(batch) + t = self.time_embed(time) + + if c_mask is None: + # Padding rows are all-zero in the text encoder's output. + c_mask = text.abs().sum(-1) > 0 + + if cfg_infer: + if cache and self.text_cond is not None: + c_cond, c_uncond = self.text_cond, self.text_uncond + else: + c_cond = self.project_text(text) + # The uncond branch drops the text, which is exactly a zeroed projection. + c_uncond = torch.zeros_like(c_cond) + if cache: + self.text_cond, self.text_uncond = c_cond, c_uncond + x_cond, mask_cond, prompt_len = self._embed_audio(x, ref, False, mask, ref_mask) + x_uncond, mask_uncond, _ = self._embed_audio(x, ref, True, mask, ref_mask) + + x = torch.cat((x_cond, x_uncond), dim=0) + c = torch.cat((c_cond, c_uncond), dim=0) + t = torch.cat((t, t), dim=0) + c_mask = torch.cat((c_mask, c_mask), dim=0) + audio_mask = None + if mask_cond is not None and mask_uncond is not None: + audio_mask = torch.cat((mask_cond, mask_uncond), dim=0) + else: + c = self.project_text(text) + if drop_text: + c = torch.zeros_like(c) + x, audio_mask, prompt_len = self._embed_audio(x, ref, drop_audio_cond, mask, ref_mask) + + seq_len = x.shape[1] + text_len = c.shape[1] + rope_audio = self.rotary_embed(seq_len) + rope_text = self.rotary_embed(text_len) + + for block in self.transformer_blocks: + c, x = block(x, c, t, mask=audio_mask, rope=rope_audio, c_rope=rope_text, c_mask=c_mask) + + x = torch.cat([c, x], dim=1) + rope = self.rotary_embed(text_len + seq_len) + single_mask = None if audio_mask is None else torch.cat([c_mask, audio_mask], dim=1) + + for block in self.single_transformer_blocks: + x = block(x, t, mask=single_mask, rope=rope) + + x = x[:, text_len + prompt_len :] + return self.proj_out(self.norm_out(x, t)) + + +def dit_state_dict( + checkpoint: Mapping[str, torch.Tensor] | Iterable[tuple[str, torch.Tensor]], + prefix: str = "transformer.", +) -> dict[str, torch.Tensor]: + """Select the DiT tensors out of a full AuK checkpoint and strip ``prefix``. + + Released checkpoints wrap the backbone under ``transformer.`` alongside the + text-encoder layer-fusion parameters, which this module does not own. + """ + items = checkpoint.items() if isinstance(checkpoint, Mapping) else checkpoint + return {name[len(prefix) :]: tensor for name, tensor in items if name.startswith(prefix)} + + +@torch.no_grad() +def sample_latents( + dit: AuKTransformer, + *, + text: torch.Tensor, + c_mask: torch.Tensor | None, + ref: torch.Tensor, + ref_mask: torch.Tensor | None, + gen_frames: int, + nfe: int = 32, + cfg_strength: float = 1.0, + sway_sampling_coef: float | None = None, + t_grid: list[float] | None = None, + seed: int | None = None, + generator: torch.Generator | None = None, + latent_dim: int | None = None, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + """Integrate the flow from noise to audio latents with explicit Euler steps. + + Single-request only: the batch dimension carries the CFG branches, not + separate prompts, so no target padding mask is needed. + + Args: + dit: The velocity model. + text: Pre-encoded text hidden states ``[1, nt, text_hidden_dim]``. + c_mask: Text padding mask ``[1, nt]``. + ref: Reference prompt latents ``[1, np, latent_dim]``. + ref_mask: Reference padding mask ``[1, np]``. + gen_frames: Target latent frames to generate. + nfe: Euler steps, ignored when ``t_grid`` is given. + cfg_strength: Classifier-free guidance weight. Below ``1e-5`` the + uncond branch is skipped entirely. + sway_sampling_coef: Reshapes the uniform time grid towards ``t=0``. + t_grid: Explicit timesteps, overriding ``nfe`` and the sway reshape. + seed: Seeds the global RNG before drawing the noise. Ignored when + ``generator`` is given. + generator: Draws the initial noise from its own RNG, leaving the global + one untouched. + latent_dim: Latent channels; defaults to the model's. + device: Device for the noise; defaults to the reference's. + dtype: Dtype for the noise and the Euler accumulator; defaults to the + reference's, which is the reference implementation's rule. Pass + ``torch.float32`` explicitly when serving half-precision weights. + Drawing the noise at bf16 instead of fp32 moves a 32-step CFG + trajectory by 0.025 relative MSE, roughly 17 times what bf16 + weights themselves cost, because the draw consumes the generator + differently and starts the ODE somewhere else. Keeping it fp32 + costs a few tens of KB and no measurable time. + + Returns: + The target latents at ``t=1``, ``[1, gen_frames, latent_dim]``. + """ + device = torch.device(device) if device is not None else ref.device + dtype = dtype if dtype is not None else ref.dtype + latent_dim = latent_dim if latent_dim is not None else dit.latent_dim + + if generator is None: + if seed is not None: + torch.manual_seed(seed) + x = torch.randn(gen_frames, latent_dim, device=device, dtype=dtype).unsqueeze(0) + else: + x = ( + torch.randn(gen_frames, latent_dim, generator=generator, device=generator.device, dtype=dtype) + .to(device) + .unsqueeze(0) + ) + + if t_grid is not None: + t = torch.tensor(t_grid, device=device, dtype=torch.float32) + else: + t = torch.linspace(0, 1, nfe + 1, device=device, dtype=torch.float32) + if sway_sampling_coef is not None: + t = t + sway_sampling_coef * (torch.cos(math.pi / 2 * t) - 1 + t) + + guided = cfg_strength >= 1e-5 + if t.numel() < 2 or not bool(torch.all(t[1:] > t[:-1])): + raise ValueError( + f"AuK sampling needs a strictly increasing time grid with at least two points; got {t.tolist()}" + ) + try: + for i in range(t.shape[0] - 1): + if guided: + pred = dit( + x, + text, + t[i], + c_mask=c_mask, + ref=ref, + ref_mask=ref_mask, + cfg_infer=True, + cache=True, + ) + v_cond, v_uncond = pred.chunk(2, dim=0) + v = v_cond + (v_cond - v_uncond) * cfg_strength + else: + v = dit(x, text, t[i], c_mask=c_mask, ref=ref, ref_mask=ref_mask) + x = x + (t[i + 1] - t[i]) * v + finally: + # The cached text projections belong to this request only; a failed + # step must not leak them into the next one. + dit.clear_cache() + return x diff --git a/vllm_omni/diffusion/models/auk/auk_vae.py b/vllm_omni/diffusion/models/auk/auk_vae.py new file mode 100644 index 00000000000..1fabeb327ac --- /dev/null +++ b/vllm_omni/diffusion/models/auk/auk_vae.py @@ -0,0 +1,505 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# +# Adapted from the AuK reference release (MIT), which in turn adapts NVIDIA +# BigVGAN (MIT), jik876/hifi-gan (MIT), junjun3518/alias-free-torch +# (Apache-2.0), EdwardDixon/snake (MIT) and adefossez/julius (MIT). +"""BigVGAN-flow VAE for AuK: 24 kHz mono waveform <-> 50 Hz, 64-channel latents. + +Inference only. The reference module also carries a residual coupling flow, a KL +term and a discriminator; all three are training-only, so the released +``flow.*`` tensors are skipped (see :meth:`AuKVAE.load_weights`). + +Parameter and buffer names follow the reference release, so the published +``vae.safetensors`` loads without a key map. That is also why the convolutions +are built with the legacy ``weight_norm`` parametrization (``weight_g`` / +``weight_v``): :meth:`remove_weight_norm` folds them into plain weights once +the checkpoint is in. +""" + +from __future__ import annotations + +import inspect +import math +import warnings +from pathlib import Path + +import torch +import torch.nn.functional as F +from safetensors.torch import load_file +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 + +logger = init_logger(__name__) + +__all__ = ["AuKVAE"] + +# Checkpoint prefix owned by the training-only coupling flow. +TRAINING_ONLY_PREFIX = "flow." + + +def _weight_norm(module: nn.Module) -> nn.Module: + """Legacy ``weight_g``/``weight_v`` split, which is how the release stores these tensors.""" + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message=".*weight_norm.*deprecated.*", category=FutureWarning) + return _apply_weight_norm(module) + + +def _same_padding(kernel_size: int, dilation: int = 1) -> int: + return int((kernel_size * dilation - dilation) / 2) + + +def _kaiser_sinc_filter(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor: + """Kaiser-windowed sinc low-pass, returned as a ``[1, 1, kernel_size]`` conv kernel.""" + half_size = kernel_size // 2 + attenuation = 2.285 * (half_size - 1) * math.pi * (4 * half_width) + 7.95 + if attenuation > 50.0: + beta = 0.1102 * (attenuation - 8.7) + elif attenuation >= 21.0: + beta = 0.5842 * (attenuation - 21.0) ** 0.4 + 0.07886 * (attenuation - 21.0) + else: + beta = 0.0 + window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) + if kernel_size % 2 == 0: + time = torch.arange(-half_size, half_size) + 0.5 + else: + time = torch.arange(kernel_size) - half_size + if cutoff == 0: + return torch.zeros_like(time).view(1, 1, kernel_size) + taps = 2 * cutoff * window * torch.sinc(2 * cutoff * time) + # Normalize to unit sum so a constant input passes through unscaled. + return (taps / taps.sum()).view(1, 1, kernel_size) + + +class Conv(nn.Conv1d): + """1-D convolution over ``[B, C, T]``; ``causal`` swaps centred padding for a left pad.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int = 1, + stride: int = 1, + dilation: int = 1, + bias: bool = True, + causal: bool = False, + ) -> None: + super().__init__( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=0 if causal else _same_padding(kernel_size, dilation), + dilation=dilation, + bias=bias, + ) + self.causal = causal + self.left_padding = dilation * (kernel_size - 1) if causal else 0 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.causal: + x = F.pad(x, (self.left_padding, 0)) + return super().forward(x) + + +class ConvTranspose(nn.ConvTranspose1d): + """Transposed 1-D convolution; ``causal`` drops the trailing ``stride`` samples instead of padding.""" + + def __init__( + self, in_channels: int, out_channels: int, kernel_size: int, stride: int, causal: bool = False + ) -> None: + if causal and kernel_size != 2 * stride: + raise ValueError(f"causal ConvTranspose needs kernel_size == 2 * stride, got {kernel_size} and {stride}") + super().__init__( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=0 if causal else (kernel_size - stride) // 2, + ) + self.causal = causal + self.trim = stride + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = super().forward(x) + return x[:, :, : -self.trim] if self.causal else x + + +class NormConv(nn.Module): + """Weight-normalized conv held under ``layer``; the attribute name is load-bearing for the checkpoint.""" + + def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 1, stride: int = 1) -> None: + super().__init__() + self.layer = _weight_norm( + nn.Conv1d(in_channels, out_channels, kernel_size, stride=stride, padding=(kernel_size - 1) // 2) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.layer(x) + + +class ResStack(nn.Module): + """Stack of dilated residual conv pairs. The leaky slope is torch's 0.01 default, as upstream.""" + + def __init__(self, channels: int, kernel_size: int = 3, base: int = 3, nums: int = 4) -> None: + super().__init__() + self.layers = nn.ModuleList( + [ + nn.Sequential( + nn.LeakyReLU(), + _weight_norm(nn.Conv1d(channels, channels, kernel_size, dilation=base**i, padding=base**i)), + nn.LeakyReLU(), + _weight_norm(nn.Conv1d(channels, channels, kernel_size, dilation=1, padding=1)), + ) + for i in range(nums) + ] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for layer in self.layers: + x = x + layer(x) + return x + + +class Encoder(nn.Module): + """Strided conv encoder over ``[B, 1, T]``, emitting ``[B, 2 * latent_dim, T / hop]`` mean/log-std pairs.""" + + def __init__( + self, + latent_dim: int, + channels: tuple[int, ...], + down_sample_factors: tuple[int, ...], + in_channels: int = 1, + proj_kernel_size: int = 3, + stack_kernel_size: int = 3, + stack_dilation_base: int = 2, + stacks: int = 6, + ) -> None: + super().__init__() + layers: list[nn.Module] = [ + NormConv(in_channels, channels[0], kernel_size=proj_kernel_size), + nn.LeakyReLU(0.2, True), + ] + for in_c, out_c, factor in zip(channels[:-1], channels[1:], down_sample_factors, strict=True): + layers += [ + NormConv(in_c, out_c, kernel_size=factor * 2, stride=factor), + ResStack(out_c, stack_kernel_size, stack_dilation_base, stacks), + nn.LeakyReLU(0.2, True), + ] + layers.append(NormConv(channels[-1], latent_dim * 2, kernel_size=proj_kernel_size)) + self.generator = nn.Sequential(*layers) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.generator(x) + + +class SnakeBeta(nn.Module): + """``x + sin^2(alpha * x) / beta`` with per-channel learned alpha and beta (log-scale in AuK).""" + + def __init__(self, channels: int, alpha_logscale: bool = False) -> None: + super().__init__() + self.alpha_logscale = alpha_logscale + init = torch.zeros(channels) if alpha_logscale else torch.ones(channels) + self.alpha = nn.Parameter(init.clone()) + self.beta = nn.Parameter(init.clone()) + self.eps = 1e-9 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + 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 x + (1.0 / (beta + self.eps)) * torch.sin(x * alpha).pow(2) + + +class LowPass(nn.Module): + """Fixed FIR low-pass with replicate padding, applied per channel.""" + + def __init__( + self, + cutoff: float = 0.5, + half_width: float = 0.6, + stride: int = 1, + kernel_size: int = 12, + causal: bool = False, + ) -> None: + super().__init__() + if not 0.0 <= cutoff <= 0.5: + raise ValueError(f"cutoff must lie in [0, 0.5], got {cutoff}") + if causal: + self.pad_left, self.pad_right = kernel_size - 1, 0 + else: + even = kernel_size % 2 == 0 + self.pad_left = kernel_size // 2 - int(even) + self.pad_right = kernel_size // 2 + self.stride = stride + self.register_buffer("filter", _kaiser_sinc_filter(cutoff, half_width, kernel_size)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + channels = x.size(1) + x = F.pad(x, (self.pad_left, self.pad_right), mode="replicate") + return F.conv1d(x, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels) + + +class Upsample(nn.Module): + """Band-limited interpolation by ``ratio``. Always non-causal, matching how AuK builds it.""" + + def __init__(self, ratio: int = 2, kernel_size: int = 12) -> None: + super().__init__() + self.ratio = ratio + self.pad = kernel_size // ratio - 1 + self.pad_left = self.pad * ratio + (kernel_size - ratio) // 2 + self.pad_right = self.pad * ratio + (kernel_size - ratio + 1) // 2 + self.register_buffer("filter", _kaiser_sinc_filter(0.5 / ratio, 0.6 / ratio, kernel_size)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + channels = x.size(1) + x = F.pad(x, (self.pad, self.pad), mode="replicate") + x = self.ratio * F.conv_transpose1d(x, self.filter.expand(channels, -1, -1), stride=self.ratio, groups=channels) + return x[..., self.pad_left : -self.pad_right] + + +class Downsample(nn.Module): + """Decimation by ``ratio`` behind a low-pass; ``lowpass`` holds the filter buffer.""" + + def __init__(self, ratio: int = 2, kernel_size: int = 12, causal: bool = False) -> None: + super().__init__() + self.lowpass = LowPass( + cutoff=0.5 / ratio, half_width=0.6 / ratio, stride=ratio, kernel_size=kernel_size, causal=causal + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.lowpass(x) + + +class AliasFreeActivation(nn.Module): + """Oversample, apply the pointwise activation, decimate, so snake harmonics do not fold back.""" + + def __init__( + self, + activation: nn.Module, + up_ratio: int = 2, + down_ratio: int = 2, + up_kernel_size: int = 12, + down_kernel_size: int = 12, + causal: bool = False, + ) -> None: + super().__init__() + self.act = activation + self.upsample = Upsample(up_ratio, up_kernel_size) + self.downsample = Downsample(down_ratio, down_kernel_size, causal=causal) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.downsample(self.act(self.upsample(x))) + + +class AmpBlock(nn.Module): + """Anti-aliased multi-periodicity residual block: alternating alias-free snake and dilated convs.""" + + def __init__( + self, + channels: int, + kernel_size: int, + dilations: tuple[int, ...], + snake_logscale: bool = True, + causal: bool = True, + act_causal: bool = False, + ) -> None: + super().__init__() + self.convs1 = nn.ModuleList( + [_weight_norm(Conv(channels, channels, kernel_size, 1, dilation=d, causal=causal)) for d in dilations] + ) + self.convs2 = nn.ModuleList( + [_weight_norm(Conv(channels, channels, kernel_size, 1, dilation=1, causal=causal)) for _ in dilations] + ) + self.activations = nn.ModuleList( + [ + AliasFreeActivation(SnakeBeta(channels, alpha_logscale=snake_logscale), causal=act_causal) + for _ in range(len(self.convs1) + len(self.convs2)) + ] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + acts1, acts2 = self.activations[::2], self.activations[1::2] + for conv1, conv2, act1, act2 in zip(self.convs1, self.convs2, acts1, acts2, strict=True): + x = conv2(act2(conv1(act1(x)))) + x + return x + + +class AuKVAE(nn.Module): + """AuK's audio codec: :meth:`encode` waveform to normalized latents, :meth:`decode` back. + + Build from the released ``model.vae.model_init_kwargs`` with :meth:`from_config`, then + :meth:`load_weights`. Latents are normalized by the checkpoint's global statistics, which is + the space the rectified-flow transformer works in. + """ + + def __init__( + self, + *, + upsample_rates: tuple[int, ...] = (5, 4, 3, 2, 2, 2), + upsample_kernel_sizes: tuple[int, ...] = (10, 8, 6, 4, 4, 4), + upsample_initial_channel: int = 1536, + resblock_kernel_sizes: tuple[int, ...] = (3, 7, 11), + resblock_dilation_sizes: tuple[tuple[int, ...], ...] = ((1, 3, 5), (1, 3, 5), (1, 3, 5)), + downsample_rates: tuple[int, ...] = (2, 2, 2, 3, 4, 5), + downsample_channels: tuple[int, ...] = (12, 24, 48, 96, 192, 384, 768), + snake_logscale: bool = True, + latent_dim: int = 64, + use_vae: bool = True, + causal: bool = True, + act_causal: bool = True, + sample_rate: int = 24000, + ) -> None: + super().__init__() + if not use_vae: + raise ValueError("AuKVAE requires use_vae=True: the encoder head emits mean/log-std pairs") + self.latent_dim = latent_dim + self.sample_rate = sample_rate + self.hop_size = math.prod(downsample_rates) + self.num_kernels = len(resblock_kernel_sizes) + self.num_upsamples = len(upsample_rates) + self._norm_folded = False + + self.register_buffer("global_mean", torch.zeros(latent_dim, dtype=torch.float32)) + self.register_buffer("global_log_std", torch.ones(latent_dim, dtype=torch.float32)) + + self.audio_encoder = Encoder( + latent_dim=latent_dim, + channels=tuple(downsample_channels), + down_sample_factors=tuple(downsample_rates), + ) + + self.conv_pre = _weight_norm(Conv(latent_dim, upsample_initial_channel, 7, 1, causal=False)) + self.ups = nn.ModuleList() + for i, (rate, kernel) in enumerate(zip(upsample_rates, upsample_kernel_sizes, strict=True)): + # One-element inner list: the checkpoint keys are ``ups..0.*``. + self.ups.append( + nn.ModuleList( + [ + _weight_norm( + ConvTranspose( + upsample_initial_channel // (2**i), + upsample_initial_channel // (2 ** (i + 1)), + kernel, + rate, + causal=causal, + ) + ) + ] + ) + ) + + self.resblocks = nn.ModuleList() + channels = upsample_initial_channel + for i in range(self.num_upsamples): + channels = upsample_initial_channel // (2 ** (i + 1)) + for kernel, dilations in zip(resblock_kernel_sizes, resblock_dilation_sizes, strict=True): + self.resblocks.append( + AmpBlock( + channels, + kernel, + tuple(dilations), + snake_logscale=snake_logscale, + causal=causal, + act_causal=act_causal, + ) + ) + + self.activation_post = AliasFreeActivation( + SnakeBeta(channels, alpha_logscale=snake_logscale), causal=act_causal + ) + self.conv_post = _weight_norm(Conv(channels, 1, 7, 1, causal=causal, bias=False)) + + @classmethod + def from_config(cls, cfg: dict) -> AuKVAE: + """Build from the released ``model_init_kwargs`` mapping, ignoring training-only entries.""" + accepted = set(inspect.signature(cls.__init__).parameters) - {"self"} + return cls(**{key: value for key, value in cfg.items() if key in accepted}) + + def load_weights( + self, path: str | Path, *, device: str | torch.device | None = None, fold_weight_norm: bool = True + ) -> tuple[list[str], list[str]]: + """Load a released ``vae.safetensors``, then fold weight norm away. Returns (missing, unexpected). + + ``device`` defaults to where the module already lives, which is the placement that matters: + the fold reduces ``weight_v`` to a norm, and a CPU reduction lands up to one ULP away from the + accelerator's. The encoder amplifies its activations roughly fiftyfold, so that one ULP grows + into a ~4e-4 relative drift by the latent head. Move the module first, then load. + """ + target = self.global_mean.device if device is None else torch.device(device) + missing, unexpected = self.load_state_dict(load_file(str(path), device=str(target)), strict=False) + missing, unexpected = list(missing), list(unexpected) + stray = [key for key in missing + unexpected if not key.startswith(TRAINING_ONLY_PREFIX)] + if stray: + raise ValueError(f"AuK VAE weights at {path} do not match this module: {sorted(stray)[:8]}") + logger.info( + "Loaded AuK VAE from %s (skipped %d training-only %s* tensors)", path, len(unexpected), TRAINING_ONLY_PREFIX + ) + if fold_weight_norm: + self.remove_weight_norm() + return missing, unexpected + + def remove_weight_norm(self) -> None: + """Fold ``weight_g``/``weight_v`` into plain weights, on whichever device the module is on. + + Idempotent, and meant to run once the checkpoint is in. See :meth:`load_weights` for why the + device this runs on is a numeric decision rather than a convenience. + """ + if self._norm_folded: + return + folded = 0 + for module in self.modules(): + if hasattr(module, "weight_g"): + _fold_weight_norm(module) + folded += 1 + self._norm_folded = True + logger.debug("Folded weight norm on %d AuK VAE convolutions", folded) + + def encode( + self, wav: torch.Tensor, *, sample: bool = False, generator: torch.Generator | None = None + ) -> torch.Tensor: + """Encode a mono ``[1, T]`` waveform at :attr:`sample_rate` into ``[1, T / hop_size, latent_dim]``. + + ``sample=False`` takes the posterior mean, which is what a reproducible pipeline wants; + ``sample=True`` draws ``mean + randn * exp(log_std)`` from ``generator``. Latents come back + normalized by the checkpoint's global statistics. + """ + if wav.dim() != 2 or wav.size(0) != 1: + raise ValueError(f"encode expects a mono [1, T] waveform, got {tuple(wav.shape)}") + # Statistics stay in fp32 even under an ambient autocast, as in the reference. + with torch.autocast(device_type=wav.device.type, enabled=False): + mean, log_std = self.audio_encoder(wav.unsqueeze(1)).chunk(2, dim=1) + if sample: + latents = mean + self._noise_like(mean, generator) * torch.exp(log_std) + else: + latents = mean + latents = latents.transpose(1, 2).float() + return (latents - self.global_mean.float()) / torch.sqrt(self.global_log_std.float()) + + def decode(self, latents: torch.Tensor) -> torch.Tensor: + """Decode normalized ``[B, Np, latent_dim]`` latents into a ``[B, Np * hop_size]`` waveform in [-1, 1].""" + if latents.dim() != 3 or latents.size(-1) != self.latent_dim: + raise ValueError(f"decode expects [B, Np, {self.latent_dim}] latents, got {tuple(latents.shape)}") + x = latents.float() * torch.sqrt(self.global_log_std.float()) + self.global_mean.float() + x = self.conv_pre(x.permute(0, 2, 1)) + for i in range(self.num_upsamples): + for up in self.ups[i]: + x = up(x) + first = i * self.num_kernels + residual = self.resblocks[first](x) + for j in range(1, self.num_kernels): + residual = residual + self.resblocks[first + j](x) + x = residual / self.num_kernels + x = self.conv_post(self.activation_post(x)) + return torch.clamp(x, min=-1.0, max=1.0).squeeze(1) + + @staticmethod + def _noise_like(reference: torch.Tensor, generator: torch.Generator | None) -> torch.Tensor: + """Draw standard normal noise shaped like ``reference``; ``generator`` decides which RNG stream.""" + device = reference.device if generator is None else generator.device + noise = torch.randn(reference.shape, generator=generator, dtype=reference.dtype, device=device) + return noise.to(reference.device) diff --git a/vllm_omni/diffusion/models/auk/pipeline_auk.py b/vllm_omni/diffusion/models/auk/pipeline_auk.py new file mode 100644 index 00000000000..6ad12b02de7 --- /dev/null +++ b/vllm_omni/diffusion/models/auk/pipeline_auk.py @@ -0,0 +1,486 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""AuK audio-editing pipeline for vLLM-Omni. + +Turns a layer-fused text condition (produced by the encoder stage) plus an +optional source clip into a 24 kHz waveform: the source clip is encoded to VAE +latents, a rectified-flow DiT integrates the target latents conditioned on +both, and the BigVGAN-flow decoder renders the waveform. +""" + +from __future__ import annotations + +import json +import os +from collections.abc import Iterable +from typing import Any, ClassVar + +import numpy as np +import torch +import torchaudio +from safetensors import safe_open +from torch import nn +from vllm.logger import init_logger + +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_vae import AuKVAE +from vllm_omni.diffusion.models.interface import ( + SupportAudioInput, + SupportAudioOutput, + SupportsComponentDiscovery, +) +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.model_extras.auk import resolve_gen_frames + +logger = init_logger(__name__) + +# Longest ref + target latent sequence the released checkpoints accept. +MAX_LATENT_FRAMES = 65536 + +# Weight-file layout: the DiT lives under this prefix; the layer-fusion +# parameters sit next to it but belong to the encoder stage. +_DIT_PREFIX = "transformer." +_FUSION_KEYS = frozenset({"layer_weights", "layer_scale"}) + +# The distilled variant only reproduces its training recipe at these settings. +_FLASH_NFE = 4 +_FLASH_CFG = 0.0 +# Decorrelates the VAE posterior draw from the initial latent for the same request seed. +_VAE_SEED_OFFSET = 0x5EED_0A0C + + +def get_auk_post_process_func(od_config: OmniDiffusionConfig): + """Create the post-processing function for AuK audio output. + + Tensor output types pass through; anything else becomes a numpy waveform. + The sample rate is not attached here: the output formatter reads + ``AuKPipeline.audio_sample_rate`` for audio-output pipelines. + """ + + del od_config # The conversion does not depend on the config. + + def post_process_func( + audio: torch.Tensor, + output_type: str = "np", + ): + if output_type in ("pt", "latent"): + return audio + return audio.cpu().float().numpy() + + return post_process_func + + +def _prompt_mapping(prompt: Any) -> dict[str, Any]: + """Return the prompt dict, or an empty mapping for a bare text prompt.""" + + if isinstance(prompt, dict): + return prompt + return {} + + +def _unwrap_single(value: Any) -> Any: + """Unwrap the one-element lists the serving layer sometimes produces.""" + + if isinstance(value, list) and len(value) == 1: + return value[0] + return value + + +def _split_audio(audio: Any, default_sample_rate: int) -> tuple[Any, int]: + """Split a multimodal audio item into raw samples and a sample rate.""" + + if isinstance(audio, (tuple, list)) and len(audio) == 2 and isinstance(audio[1], (int, float)): + return audio[0], int(audio[1]) + return audio, default_sample_rate + + +class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComponentDiscovery): + """Instruction-driven audio generation and editing with AuK. + + One request per forward: the rectified-flow ODE runs over the whole target + span and the batch dimension carries the CFG branches, so there is nothing + to share between requests yet. + + Args: + od_config: OmniDiffusion configuration. ``od_config.model`` must be an + assembled AuK directory (``config.json``, ``auk.safetensors``, + ``vae.safetensors``), as produced by + ``tools/prepare_auk_checkpoint.py``. + prefix: Unused; kept for the pipeline construction contract. + """ + + supports_request_batch = False + + # Picked up by ``supports_audio_output`` in the diffusion engine so the + # default stage metadata reports ``final_output_type="audio"`` and the + # ``multimodal_output`` payload includes the sample rate. + support_audio_output: ClassVar[bool] = True + support_audio_input: ClassVar[bool] = True + audio_sample_rate: ClassVar[int] = 24000 + + # The DiT cannot run without an encoder-stage text condition, which the + # engine's synthetic warmup request has no way to produce. + dummy_run_num_frames: ClassVar[int] = 0 + + _dit_modules: ClassVar[list[str]] = ["dit"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = ["vae"] + + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ) -> None: + super().__init__() + del prefix # Weights are not namespaced: one checkpoint, one pipeline. + self.od_config = od_config + self.device = get_local_device() + self.dtype = getattr(od_config, "dtype", None) or torch.bfloat16 + + model_dir = od_config.model + if not model_dir or not os.path.isdir(model_dir): + raise ValueError( + f"AuK needs an assembled local checkpoint directory, got {model_dir!r}. " + "Build one with tools/prepare_auk_checkpoint.py." + ) + config = _read_config(model_dir) + + vae_config = dict(config["vae"]) + self.latent_dim = int(vae_config["latent_dim"]) + self.hop_size = int(vae_config["downsample_rate"]) + self.sample_rate = int(vae_config["target_sample_rate"]) + + self.variant = str(config.get("variant", "base")) + self.is_flash = self.variant == "flash" + self.flash_t_grid = [float(t) for t in config.get("flash_t_grid") or ()] + if self.is_flash and len(self.flash_t_grid) != _FLASH_NFE + 1: + raise ValueError(f"The flash variant needs a {_FLASH_NFE + 1}-point flash_t_grid, got {self.flash_t_grid}.") + defaults = dict(config.get("defaults") or {}) + self.default_nfe = int(defaults.get("nfe", 32)) + self.default_cfg = float(defaults.get("cfg", 2.0)) + self.default_sway = defaults.get("sway", -1.0) + self._flash_lock_logged = False + + # The codec is small and numerically sensitive: it stays in fp32. It + # must sit on the inference device BEFORE load_weights folds weight + # norm: a CPU fold lands one ULP off the accelerator's and the encoder + # amplifies that into a 1e-2 latent drift. + self.vae = AuKVAE.from_config(vae_config["model_init_kwargs"]).to(device=self.device, dtype=torch.float32) + self.vae.load_weights(os.path.join(model_dir, "vae.safetensors")) + self.vae = self.vae.eval() + self.vae.requires_grad_(False) + self._check_vae_geometry() + + # The transformer owns every key the checkpoint tool writes into `dit`, + # `attn_mask_enabled` included, so the section passes straight through. + self.dit = AuKTransformer(latent_dim=self.latent_dim, **config["dit"]) + self.dit = self.dit.to(dtype=self.dtype) + self.dit.load_state_dict(_read_dit_weights(model_dir, self.dtype), strict=True) + self.dit = self.dit.to(device=self.device).eval() + self.dit.requires_grad_(False) + + logger.info( + "AuK pipeline ready: variant=%s dtype=%s latent_dim=%d hop=%d sample_rate=%d", + self.variant, + self.dtype, + self.latent_dim, + self.hop_size, + self.sample_rate, + ) + + # The assembled checkpoint is not a diffusers layout: __init__ reads + # auk.safetensors and vae.safetensors directly, so the loader has no + # component sources to stream and load_weights only reports what is loaded. + weights_sources: ClassVar[tuple] = () + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stray = [name for name, _ in weights] + if stray: + logger.warning("AuKPipeline ignores %d loader-provided weights (e.g. %s)", len(stray), stray[:3]) + return {name for name, _ in self.named_parameters()} + + def _check_vae_geometry(self) -> None: + """Fail fast when config.json and the VAE disagree on latent geometry.""" + + declared = { + "latent_dim": (self.latent_dim, int(self.vae.latent_dim)), + "downsample_rate": (self.hop_size, int(self.vae.hop_size)), + "target_sample_rate": (self.sample_rate, int(self.vae.sample_rate)), + } + bad = {name: pair for name, pair in declared.items() if pair[0] != pair[1]} + if bad: + raise ValueError(f"config.json and the AuK VAE disagree (config, vae): {bad}") + if self.sample_rate != self.audio_sample_rate: + raise ValueError( + f"AuK advertises {self.audio_sample_rate} Hz output but the checkpoint is {self.sample_rate} Hz." + ) + + def _resolve_generator(self, sampling_params: Any) -> torch.Generator | None: + """Per-request generator; the process-global RNG is never seeded.""" + + generator = sampling_params.generator + if isinstance(generator, list): + if len(generator) > 1: + logger.warning( + "AuKPipeline runs one request per forward; using the first of %d generators", len(generator) + ) + generator = generator[0] if generator else None + if generator is None and sampling_params.seed is not None: + generator = torch.Generator(device=self.device).manual_seed(int(sampling_params.seed)) + return generator + + def _prepare_waveform(self, audio: Any) -> torch.Tensor: + """Return mono float32 samples ``[1, T]`` at the VAE sample rate. + + ``T`` is truncated to a whole number of latent frames. The encoder is + causal, so dropping a trailing partial frame leaves every emitted frame + unchanged. + """ + + data, sample_rate = _split_audio(audio, self.sample_rate) + if isinstance(data, np.ndarray): + wav = torch.from_numpy(np.ascontiguousarray(data)) + else: + wav = torch.as_tensor(data) + wav = wav.detach().to(device="cpu", dtype=torch.float32) + + if wav.ndim == 1: + wav = wav.unsqueeze(0) + elif wav.ndim == 2: + # Accept both [channels, samples] and soundfile's [samples, channels]; + # the channel axis is the short one (at most 8 wide). + if wav.shape[0] > wav.shape[1] and wav.shape[1] <= 8: + wav = wav.transpose(0, 1) + if wav.shape[0] > 1: + wav = wav.mean(dim=0, keepdim=True) + else: + raise ValueError(f"AuK audio input must be [samples] or [channels, samples], got {tuple(wav.shape)}.") + + if wav.shape[-1] == 0: + raise ValueError("AuK audio input is empty.") + if not torch.isfinite(wav).all(): + raise ValueError("AuK audio input contains NaN or Inf.") + + if sample_rate != self.sample_rate: + wav = torchaudio.functional.resample(wav, sample_rate, self.sample_rate) + + frames = wav.shape[-1] // self.hop_size + if frames < 1: + raise ValueError( + f"AuK audio input is shorter than one latent frame ({self.hop_size} samples at {self.sample_rate} Hz)." + ) + return wav[..., : frames * self.hop_size] + + def _encode_source( + self, + audio: Any, + *, + sample: bool, + generator: torch.Generator | None, + ) -> torch.Tensor: + """Encode the source clip to normalized latents ``[1, np, latent_dim]``.""" + + 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) + + def _target_frames(self, gen_seconds: Any, ref_frames: int) -> int: + """Target latent length, capped by the remaining context.""" + + frames = resolve_gen_frames( + None if gen_seconds is None else float(gen_seconds), + ref_frames, + sample_rate=self.sample_rate, + hop=self.hop_size, + ) + budget = MAX_LATENT_FRAMES - ref_frames + if budget < 1: + raise ValueError( + f"The source clip already fills the {MAX_LATENT_FRAMES}-frame context ({ref_frames} frames)." + ) + if frames > budget: + logger.warning("AuK target of %d frames exceeds the remaining context; clamping to %d.", frames, budget) + frames = budget + return frames + + def _resolve_schedule( + self, + sampling_params: Any, + knobs: dict[str, Any], + ) -> tuple[int, float, float | None, list[float] | None]: + """Resolve (nfe, cfg, sway, t_grid), locking the distilled recipe on Flash. + + A knob present but ``None`` counts as unspecified: the stage input + processor fills every key it knows about, whether the caller set it or + not. + """ + + steps = sampling_params.num_inference_steps + nfe = self.default_nfe if steps is None else int(steps) + if nfe < 1: + raise ValueError(f"num_inference_steps must be >= 1 for AuK; got {steps}") + cfg = float(sampling_params.guidance_scale) if sampling_params.guidance_scale_provided else self.default_cfg + sway = knobs.get("sway") + sway = self.default_sway if sway is None else float(sway) + raw_grid = knobs.get("t_grid") + t_grid = [float(t) for t in raw_grid] if raw_grid else None + + if not self.is_flash: + return nfe, cfg, sway, t_grid + + locked = (_FLASH_NFE, _FLASH_CFG, None, self.flash_t_grid) + if (nfe, cfg, sway, t_grid) != locked and not self._flash_lock_logged: + self._flash_lock_logged = True + logger.warning( + "AuK-Flash ignores per-request sampling knobs (asked nfe=%s cfg=%s sway=%s t_grid=%s); " + "using the distilled recipe nfe=%d cfg=%s with the checkpoint time grid.", + nfe, + cfg, + sway, + t_grid, + _FLASH_NFE, + _FLASH_CFG, + ) + return locked + + def _dit_autocast(self): + """Match the reference, which runs the DiT under bf16 autocast.""" + + return torch.autocast( + device_type=self.device.type, + dtype=self.dtype, + enabled=self.dtype in (torch.bfloat16, torch.float16), + ) + + def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: + """Generate one waveform. + + Args: + req: Request batch holding exactly one request. The prompt carries + ``prompt_embeds`` (the fused text condition ``[nt, 2048]``), an + optional ``multi_modal_data["audio"]`` source clip, and + ``additional_information["auk"]`` with ``gen_seconds``, + ``sway``, ``t_grid`` and ``vae_sample``. ``num_inference_steps``, + ``guidance_scale`` and ``seed``/``generator`` come from the + sampling params. + + Returns: + One ``DiffusionOutput`` whose ``output`` is a float32 mono waveform + ``[T]`` at 24 kHz, or the normalized target latents + ``[1, gen_frames, latent_dim]`` when ``output_type`` is ``latent``. + """ + + assert req.num_reqs == 1, f"AuKPipeline runs one request per forward, got {req.num_reqs}." + prompt = _prompt_mapping(req.prompts[0]) + sampling_params = req.sampling_params + knobs = dict((prompt.get("additional_information") or {}).get("auk") or {}) + + text = _unwrap_single(prompt.get("prompt_embeds")) + if text is None: + raise ValueError("AuK stage 1 needs `prompt_embeds`, the fused text condition from the encoder stage.") + text = torch.as_tensor(text).to(device=self.device, dtype=self.dtype) + if text.ndim == 3 and text.shape[0] == 1: + text = text[0] + if text.ndim != 2: + raise ValueError(f"AuK `prompt_embeds` must be [nt, text_hidden_dim], got {tuple(text.shape)}.") + text = text.unsqueeze(0) + c_mask = torch.ones(text.shape[:2], dtype=torch.bool, device=self.device) + + generator = self._resolve_generator(sampling_params) + audio = (prompt.get("multi_modal_data") or {}).get("audio") + output_type = sampling_params.output_type or "np" + + with torch.inference_mode(): + # The VAE posterior draw gets its own generator, seeded with a + # fixed offset from the request seed: the initial latent below then + # still equals the reference's fresh manual_seed draw, and the two + # noise streams are not copies of each other. + vae_generator = None + if knobs.get("vae_sample") and generator is not None: + vae_seed = (int(generator.initial_seed()) + _VAE_SEED_OFFSET) % (2**63) + vae_generator = torch.Generator(device=self.device).manual_seed(vae_seed) + ref = self._encode_source(audio, sample=bool(knobs.get("vae_sample")), generator=vae_generator) + ref_frames = ref.shape[1] + ref_mask = torch.ones(ref.shape[:2], dtype=torch.bool, device=self.device) + gen_frames = self._target_frames(knobs.get("gen_seconds"), ref_frames) + nfe, cfg, sway, t_grid = self._resolve_schedule(sampling_params, knobs) + + with self._dit_autocast(): + latents = sample_latents( + self.dit, + text=text, + c_mask=c_mask, + ref=ref, + ref_mask=ref_mask, + gen_frames=gen_frames, + nfe=nfe, + cfg_strength=cfg, + sway_sampling_coef=sway, + t_grid=t_grid, + seed=None, + latent_dim=self.latent_dim, + device=self.device, + dtype=torch.float32, + generator=generator, + ) + + latents = latents.float() + if not torch.isfinite(latents).all(): + raise RuntimeError("AuK generated latents contain NaN or Inf.") + if output_type == "latent": + return [DiffusionOutput(output=latents.detach().cpu())] + + wav = self.vae.decode(latents) + + # One mono waveform per request; the formatter expects [T]. + wav = wav.detach().to(device="cpu", dtype=torch.float32).reshape(-1) + if not torch.isfinite(wav).all(): + raise RuntimeError("AuK generated audio contains NaN or Inf.") + return [DiffusionOutput(output=wav)] + + +def _read_config(model_dir: str) -> dict[str, Any]: + """Read the assembled AuK ``config.json``.""" + + path = os.path.join(model_dir, "config.json") + try: + with open(path) as handle: + config = json.load(handle) + except FileNotFoundError as exc: + raise ValueError(f"AuK checkpoint is missing config.json: {path}") from exc + except json.JSONDecodeError as exc: + logger.error("AuK config.json is not valid JSON: %s", path) + raise ValueError(f"AuK config.json is not valid JSON: {path}") from exc + + missing = [key for key in ("dit", "vae") if key not in config] + if missing: + raise ValueError(f"AuK config.json is missing required sections {missing}: {path}") + return config + + +def _read_dit_weights(model_dir: str, dtype: torch.dtype) -> dict[str, torch.Tensor]: + """Read the DiT state dict from ``auk.safetensors``, cast to ``dtype``. + + Tensors are cast one at a time so the fp32 checkpoint is never fully + resident alongside the model. + """ + + path = os.path.join(model_dir, "auk.safetensors") + if not os.path.isfile(path): + raise ValueError(f"AuK checkpoint is missing auk.safetensors: {path}") + + with safe_open(path, framework="pt", device="cpu") as checkpoint: + keys = list(checkpoint.keys()) + state = dit_state_dict(((key, checkpoint.get_tensor(key).to(dtype)) for key in keys), _DIT_PREFIX) + + if not state: + raise ValueError(f"AuK checkpoint has no '{_DIT_PREFIX}*' weights: {path}") + unexpected = sorted(key for key in keys if key not in _FUSION_KEYS and not key.startswith(_DIT_PREFIX)) + if unexpected: + logger.warning("Ignoring %d unrecognized keys in %s, e.g. %s", len(unexpected), path, unexpected[:5]) + return state diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 39fff7df256..26b052a3ab2 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -116,6 +116,11 @@ "pipeline_minimax_h3", "MiniMaxH3Pipeline", ), + "AuKPipeline": ( + "auk", + "pipeline_auk", + "AuKPipeline", + ), "StableAudioPipeline": ( "stable_audio", "pipeline_stable_audio", @@ -567,6 +572,7 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) - "LTX2I2VDMD2Pipeline": "get_ltx2_post_process_func", "MiniMaxH3Pipeline": "get_minimax_h3_post_process_func", "MiniMaxH3ModularPipeline": "get_minimax_h3_post_process_func", + "AuKPipeline": "get_auk_post_process_func", "StableAudioPipeline": "get_stable_audio_post_process_func", "WanImageToVideoPipeline": "get_wan22_i2v_post_process_func", "WanS2VPipeline": "get_wan22_s2v_post_process_func", diff --git a/vllm_omni/model_executor/models/auk/__init__.py b/vllm_omni/model_executor/models/auk/__init__.py new file mode 100644 index 00000000000..d77d6a64daf --- /dev/null +++ b/vllm_omni/model_executor/models/auk/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +from vllm_omni.model_executor.models.auk.auk import AuKForConditionalGeneration + +__all__ = ["AuKForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/auk/auk.py b/vllm_omni/model_executor/models/auk/auk.py new file mode 100644 index 00000000000..ebb20711dfe --- /dev/null +++ b/vllm_omni/model_executor/models/auk/auk.py @@ -0,0 +1,304 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""AuK (Tencent): instruction-driven speech generation and editing. + +Stage 0 (``model_stage="encoder"``) runs the frozen Qwen2.5-Omni thinker as a +prefill-only encoder. The outputs of all decoder layers are fused with learned +softmax weights into one ``[tokens, hidden]`` text condition, emitted through +``multimodal_outputs["hidden_states"]["output"]`` for the diffusion stage +(``vllm_omni/diffusion/models/auk``). The stage samples exactly one token +(``max_tokens=1``, forced to EOS); that token is discarded downstream. +""" + +from __future__ import annotations + +import os +from collections.abc import Iterable, Mapping +from functools import cached_property +from itertools import islice + +import torch +import torch.nn as nn +import torch.nn.functional as F +from safetensors import safe_open +from transformers import PretrainedConfig +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP +from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalFeatureSpec +from vllm.sequence import IntermediateTensors +from vllm.v1.sample.sampler import Sampler + +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni import Qwen2_5OmniForConditionalGeneration +from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( + Qwen2_5OmniThinkerDummyInputsBuilder, + Qwen2_5OmniThinkerForConditionalGeneration, + Qwen2_5OmniThinkerMultiModalProcessor, + Qwen2_5OmniThinkerProcessingInfo, +) +from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights + +logger = init_logger(__name__) + +# File next to the thinker shards in the assembled checkpoint directory that +# carries the DiT (``transformer.*``) and the two layer-fusion tensors. +AUK_WEIGHTS_FILE = "auk.safetensors" +LAYER_WEIGHTS_KEY = "layer_weights" +LAYER_SCALE_KEY = "layer_scale" +# The fused text condition travels as ``multimodal_outputs["hidden_states"]["output"]`` +# (``HiddenStatesStruct.output``), the same stage-wire slot MiniMax H3 uses. +TEXT_COND_KEY = "output" +# ``F.layer_norm`` epsilon used by the upstream fusion (torch default). +_FUSION_LN_EPS = 1e-5 + + +class AuKProcessingInfo(Qwen2_5OmniThinkerProcessingInfo): + """Qwen2.5-Omni thinker processing restricted to AuK's single audio input. + + AuK conditions on text plus at most one reference or source clip; images + and video are not part of its contract, so they are not advertised (and the + thinker then builds no vision tower). + """ + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"audio": 1} + + +@MULTIMODAL_REGISTRY.register_processor( + Qwen2_5OmniThinkerMultiModalProcessor, + info=AuKProcessingInfo, + dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder, +) +class AuKForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP, SupportsMRoPE): + """Encoder stage of AuK: Qwen2.5-Omni thinker with learned layer fusion. + + The forward pass re-implements the thinker language model's layer loop so + the per-layer outputs can be fused on the fly: vLLM's decoder layers carry + the residual stream separately, and ``hidden + residual`` after layer ``i`` + is the HF ``output_hidden_states[i + 1]`` entry, while the last entry is + the final-normed state. Only one accumulator lives at a time. + """ + + have_multimodal_outputs = True + # The fused condition is the payload; the runner's default last-hidden + # copy and prefix-cached hidden states are not needed (MiniMax H3 pattern). + omni_pooler_payload_include_hidden = False + requires_full_prefix_cached_hidden_states = False + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.has_preprocess = False + self.vllm_config = vllm_config + config = vllm_config.model_config.hf_config + + self.model_stage = getattr(vllm_config.model_config, "model_stage", None) or "encoder" + if self.model_stage != "encoder": + raise ValueError( + f"AuKForConditionalGeneration only hosts the 'encoder' stage, got model_stage={self.model_stage!r}; " + "the DiT stage is the AuKPipeline diffusion model." + ) + + thinker_config = config.thinker_config + self.thinker_config = thinker_config + self.thinker = init_vllm_registered_model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "thinker"), + hf_config=thinker_config, + architectures=["Qwen2_5OmniThinkerModel"], + ) + # The audio-only mm limits keep the thinker from building its vision + # tower; drop it explicitly as well, as the upstream implementation does, + # so no vision weights are loaded or resident. + if getattr(self.thinker, "visual", None) is not None: + del self.thinker.visual + self.thinker.visual = None + self.model = self.thinker + self.make_empty_intermediate_tensors = self.thinker.make_empty_intermediate_tensors + + num_layers = thinker_config.text_config.num_hidden_layers + # Filled by load_weights() from AUK_WEIGHTS_FILE; buffers so the vLLM + # loader does not expect them among the thinker shards. + self.register_buffer("layer_weights", torch.zeros(num_layers, dtype=torch.float32)) + self.register_buffer("layer_scale", torch.ones(1, dtype=torch.float32)) + self._fusion_loaded = False + + # -------------------- delegations to the thinker -------------------- + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + return Qwen2_5OmniThinkerForConditionalGeneration.get_placeholder_str(modality, i) + + def get_language_model(self) -> nn.Module: + return self.thinker.get_language_model() + + @cached_property + def sampler(self): + if hasattr(self.thinker, "sampler"): + return self.thinker.sampler + return Sampler() + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings=None, + is_multimodal=None, + ) -> torch.Tensor: + return self.thinker.embed_input_ids( + input_ids=input_ids, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + def embed_multimodal(self, **kwargs): + return self.thinker.embed_multimodal(**kwargs) + + def get_mrope_input_positions( + self, + input_tokens: list[int], + mm_features: list[MultiModalFeatureSpec] | None = None, + *, + hf_config: PretrainedConfig, + image_grid_thw: list[list[int]] | torch.Tensor, + video_grid_thw: list[list[int]] | torch.Tensor, + second_per_grid_ts: list[float] | None = None, + context_len: int = 0, + seq_len: int | None = None, + audio_feature_lengths: torch.Tensor | None = None, + use_audio_in_video: bool = False, + ) -> tuple[torch.Tensor, int]: + # The omni runner filters kwargs by this signature (the extended + # Qwen2.5-Omni form); the unified Qwen2.5-Omni class implements the + # audio-aware position math without touching instance state. + return Qwen2_5OmniForConditionalGeneration.get_mrope_input_positions( + self, + input_tokens, + mm_features, + hf_config=hf_config, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + context_len=context_len, + seq_len=seq_len, + audio_feature_lengths=audio_feature_lengths, + use_audio_in_video=use_audio_in_video, + ) + + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, **kwargs: object) -> torch.Tensor | None: + """Emit EOS immediately; the useful result is the prefill text condition.""" + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + text_config = self.thinker_config.text_config + logits = hidden_states.new_full( + (hidden_states.shape[0], int(text_config.vocab_size)), + torch.finfo(hidden_states.dtype).min, + ) + logits[:, self._eos_token_id()] = 0 + return logits + + def _eos_token_id(self) -> int: + for cfg in (self.thinker_config.text_config, self.thinker_config, self.vllm_config.model_config.hf_config): + eos = getattr(cfg, "eos_token_id", None) + if isinstance(eos, (list, tuple)) and eos: + eos = eos[0] + if isinstance(eos, int): + return eos + # Qwen2.5-Omni tokenizers end turns with <|im_end|>. + return 151645 + + # -------------------- forward -------------------- + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> OmniOutput: + if intermediate_tensors is not None: + raise NotImplementedError("AuK encoder does not support pipeline parallelism") + language_model = self.thinker.language_model.model + num_layers = len(language_model.layers) + if language_model.start_layer != 0 or language_model.end_layer != num_layers: + raise NotImplementedError("AuK encoder needs all decoder layers on one rank") + + if inputs_embeds is None: + hidden_states = language_model.embed_input_ids(input_ids) + else: + hidden_states = inputs_embeds + residual: torch.Tensor | None = None + + weights = torch.softmax(self.layer_weights, dim=0) + fused: torch.Tensor | None = None + normed_shape = (hidden_states.shape[-1],) + layers = islice(language_model.layers, language_model.start_layer, language_model.end_layer) + for idx, layer in enumerate(layers): + hidden_states, residual = layer(positions, hidden_states, residual) + if idx + 1 < num_layers: + layer_output = hidden_states + residual + else: + # HF's last hidden_states entry is post final norm. + hidden_states, _ = language_model.norm(hidden_states, residual) + layer_output = hidden_states + normed = F.layer_norm(layer_output.float(), normed_shape, eps=_FUSION_LN_EPS) + contribution = normed * weights[idx] + fused = contribution if fused is None else fused + contribution + + assert fused is not None + text_cond = (fused * self.layer_scale).to(hidden_states.dtype) + # Stage-wire layout: [tokens, features] under hidden_states.output so + # the runner slices it per request (and concatenates across chunked + # prefill steps); the stage-1 processor reads the same path. + return OmniOutput( + text_hidden_states=hidden_states, + multimodal_outputs={"hidden_states": {"output": text_cond}}, + ) + + # -------------------- weights -------------------- + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + thinker_weights: list[tuple[str, torch.Tensor]] = [] + skipped = 0 + for name, tensor in weights: + if name.startswith("thinker."): + thinker_weights.append((name, tensor)) + elif name.startswith(("talker.", "token2wav.", "transformer.")): + skipped += 1 + else: + logger.warning("AuK encoder: ignoring unexpected weight %s", name) + if skipped: + logger.debug("AuK encoder: skipped %d non-thinker weights", skipped) + + loaded = self.thinker.load_weights(thinker_weights) + loaded = add_prefix_to_loaded_weights(loaded, "thinker") + self._load_fusion_params() + return loaded + + def _load_fusion_params(self) -> None: + model_dir = self.vllm_config.model_config.model + path = os.path.join(model_dir, AUK_WEIGHTS_FILE) + if not os.path.isfile(path): + raise FileNotFoundError( + f"AuK encoder: {AUK_WEIGHTS_FILE} not found in {model_dir!r}; assemble the checkpoint with " + "tools/prepare_auk_checkpoint.py" + ) + with safe_open(path, framework="pt", device="cpu") as f: + layer_weights = f.get_tensor(LAYER_WEIGHTS_KEY).float() + layer_scale = f.get_tensor(LAYER_SCALE_KEY).float().reshape(1) + if layer_weights.shape != self.layer_weights.shape: + raise ValueError( + f"AuK encoder: {LAYER_WEIGHTS_KEY} has {tuple(layer_weights.shape)} entries, " + f"thinker has {tuple(self.layer_weights.shape)} layers" + ) + self.layer_weights.copy_(layer_weights.to(self.layer_weights.device)) + self.layer_scale.copy_(layer_scale.to(self.layer_scale.device)) + self._fusion_loaded = True + logger.info( + "AuK encoder: loaded layer fusion (scale %.3f, top layer weight %.3f at layer %d)", + float(layer_scale[0]), + float(torch.softmax(layer_weights, 0).max()), + int(torch.softmax(layer_weights, 0).argmax()) + 1, + ) diff --git a/vllm_omni/model_executor/models/auk/pipeline.py b/vllm_omni/model_executor/models/auk/pipeline.py new file mode 100644 index 00000000000..8b9d6855bab --- /dev/null +++ b/vllm_omni/model_executor/models/auk/pipeline.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""AuK pipeline topology (frozen). + +Stage 0: Encoder — frozen Qwen2.5-Omni thinker, prefill only, emits the + learned layer fusion as ``multimodal_outputs["hidden_states"]["output"]`` +Stage 1: DiT — rectified-flow transformer plus the BigVGAN-flow VAE, + emits a 24 kHz waveform +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.auk" + +AUK_PIPELINE = PipelineConfig( + model_type="auk", + default_deploy_config_name="auk.yaml", + model_arch="AuKForConditionalGeneration", + hf_architectures=("AuKForConditionalGeneration",), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="encoder", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=False, + owns_tokenizer=True, + requires_multimodal_data=True, + hf_config_name="thinker_config", + engine_output_type="latent", + model_arch="AuKForConditionalGeneration", + # No model_subdir: tools/prepare_auk_checkpoint.py assembles one + # flat directory, so the encoder finds the thinker shards, the + # tokenizer, the audio processor and the layer-fusion tensors in + # auk.safetensors all under the pipeline root. + # One prefill step produces the text condition; the sampled token + # is discarded, and detokenizing it would only cost latency. + sampling_constraints={"max_tokens": 1, "temperature": 0.0, "detokenize": False}, + ), + StagePipelineConfig( + stage_id=1, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(0,), + requires_multimodal_data=True, + final_output=True, + final_output_type="audio", + model_arch="AuKPipeline", + custom_process_input_func=f"{_PROC}.encoder2dit", + omni_kv_config={"need_recv_cache": False}, + # Single replica, and the whole ODE runs inside one forward, so + # the stage stays in the orchestrator process. + inline_diffusion=True, + ), + ), +) diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index b9373e5f229..febda68afa4 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -87,6 +87,11 @@ "cosyvoice3", "CosyVoice3Model", ), + "AuKForConditionalGeneration": ( + "auk", + "auk", + "AuKForConditionalGeneration", + ), "NemotronDenseForCausalLM": ( "audex", "audex_thinker", diff --git a/vllm_omni/model_executor/stage_input_processors/auk.py b/vllm_omni/model_executor/stage_input_processors/auk.py new file mode 100644 index 00000000000..7c417aee171 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/auk.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Stage input processor for AuK: thinker encoder -> diffusion transition.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import torch +from vllm.inputs import TextPrompt +from vllm.logger import init_logger + +from vllm_omni.inputs.data import OmniTokensPrompt + +logger = init_logger(__name__) + +# The encoder emits the fused condition as +# ``multimodal_outputs["hidden_states"]["output"]`` (``HiddenStatesStruct.output``, +# the same stage-wire slot MiniMax H3 uses). The payload flattener may deliver +# it under the dotted spelling instead, and a plain ``text_cond`` key is +# accepted so an encoder that emits the condition directly still works. +HIDDEN_STATES_KEY = "hidden_states" +TEXT_COND_KEY = "output" +_FALLBACK_KEYS = ("hidden_states.output", "text_cond") + +# Per-request knobs the prompt carries for stage 1. nfe/cfg/seed are not here: +# they ride the stage-1 OmniDiffusionSamplingParams instead. +_AUK_KNOBS = ("gen_seconds", "sway", "t_grid", "vae_sample") + + +def _as_dict(prompt: Any) -> dict[str, Any]: + """Normalize the many prompt spellings the orchestrator may pass.""" + if isinstance(prompt, list): + prompt = prompt[0] if prompt else {} + if prompt is None: + return {} + if isinstance(prompt, dict): + return prompt + if hasattr(prompt, "_asdict"): + return prompt._asdict() + if hasattr(prompt, "__dict__"): + return vars(prompt) + return {} + + +def _has_audio(mm_data: Any) -> bool: + return isinstance(mm_data, Mapping) and mm_data.get("audio") is not None + + +def _from_hidden_states(hidden: Any) -> Any: + """Read the fused condition out of the payload's hidden-states slot.""" + if isinstance(hidden, Mapping): + return hidden.get(TEXT_COND_KEY) + # A deserialized payload arrives as a msgspec struct rather than a dict. + return getattr(hidden, TEXT_COND_KEY, None) + + +def _extract_text_cond(ar_output: Any) -> Any: + """Read the fused text condition from the encoder output. + + ``multimodal_output`` is attached to the RequestOutput; the CompletionOutput + is checked as a fallback for vLLM versions that attach it there. + """ + for holder in (ar_output, *(getattr(ar_output, "outputs", None) or ())): + mm_output = getattr(holder, "multimodal_output", None) + if not isinstance(mm_output, Mapping): + continue + text_cond = _from_hidden_states(mm_output.get(HIDDEN_STATES_KEY)) + if text_cond is not None: + return text_cond + for key in _FALLBACK_KEYS: + if mm_output.get(key) is not None: + return mm_output[key] + return None + + +def _to_cpu_tensor(text_cond: Any) -> torch.Tensor: + """Coerce the emitted condition to a single 2-D CPU tensor.""" + if isinstance(text_cond, list): + if not text_cond: + raise ValueError("AuK encoder emitted an empty text condition list") + parts = [_to_cpu_tensor(part) for part in text_cond] + return torch.cat(parts, dim=0) if len(parts) > 1 else parts[0] + if isinstance(text_cond, torch.Tensor): + tensor = text_cond + else: + # numpy array or any array-like the connector round-tripped. + tensor = torch.as_tensor(text_cond) + tensor = tensor.detach().cpu() + if tensor.ndim != 2: + raise ValueError(f"AuK text condition must be [tokens, hidden]; got shape {tuple(tensor.shape)}") + return tensor + + +def encoder2dit( + source_outputs: list[Any], + prompt: OmniTokensPrompt | TextPrompt | list | None = None, + requires_multimodal_data: bool = False, + streaming_context: Any | None = None, +) -> dict[str, Any] | None: + """Turn the AuK encoder stage's fused hidden states into a diffusion prompt. + + The source clip is not decoded here: the original ``multi_modal_data`` is + forwarded so stage 1 resamples and VAE-encodes it itself. It is forwarded + unconditionally because the editing and enhancement tasks cannot run + without it; ``requires_multimodal_data`` only records the stage's + declaration. + """ + del streaming_context, requires_multimodal_data + if not source_outputs: + return None + + ar_output = source_outputs[0] + text_cond = _extract_text_cond(ar_output) + if text_cond is None: + raise ValueError( + "AuK encoder stage produced no hidden_states.output multimodal payload; " + "stage 1 cannot run without the fused thinker hidden states" + ) + prompt_embeds = _to_cpu_tensor(text_cond) + + original = _as_dict(prompt) + original_knobs = original.get("additional_information") or {} + if isinstance(original_knobs, Mapping): + original_knobs = original_knobs.get("auk") or {} + if not isinstance(original_knobs, Mapping): + original_knobs = {} + + mm_data = original.get("multi_modal_data") + knobs: dict[str, Any] = {key: original_knobs.get(key) for key in _AUK_KNOBS} + declared_audio = original_knobs.get("has_audio") + knobs["has_audio"] = bool(declared_audio) if declared_audio is not None else _has_audio(mm_data) + + logger.debug( + "[encoder2dit] text_cond=%s dtype=%s has_audio=%s gen_seconds=%s", + tuple(prompt_embeds.shape), + prompt_embeds.dtype, + knobs["has_audio"], + knobs["gen_seconds"], + ) + + return { + "prompt": "", + "prompt_embeds": prompt_embeds, + "multi_modal_data": mm_data, + "additional_information": {"auk": knobs}, + } diff --git a/vllm_omni/model_extras/auk.py b/vllm_omni/model_extras/auk.py new file mode 100644 index 00000000000..16d6af52ac1 --- /dev/null +++ b/vllm_omni/model_extras/auk.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Request helpers for AuK: prompt rendering, sampling params, duration policy. + +AuK's contract is one instruction plus an optional reference clip in, audio +out; the task (zero-shot TTS, instruct TTS, content or acoustic editing, +enhancement, separation) is carried by the instruction text alone, so there is +no per-task request surface. +""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING, Any + +from vllm import SamplingParams + +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +if TYPE_CHECKING: + import numpy as np + import torch + +# The upstream ChatML render of a single user turn through +# ``Qwen2_5OmniProcessor.apply_chat_template(..., add_generation_prompt=True)`` +# with the template's default system turn. +_CHATML = ( + "<|im_start|>system\n" + "You are a helpful assistant.<|im_end|>\n" + "<|im_start|>user\n" + "{instruction}{audio_part}<|im_end|>\n" + "<|im_start|>assistant\n" +) +_AUDIO_PART = "<|audio_bos|><|AUDIO|><|audio_eos|>" + +# Literal marker the reference appends to the instruction when no reference +# clip is supplied. The model was trained with it, so it is not decorative. +NO_PROMPT_AUDIO = "||" + +# Latent frame rate: 24 kHz waveform, VAE downsample 480 -> 50 frames/s. +SAMPLE_RATE = 24000 +HOP = 480 + +DEFAULT_NFE = 32 +DEFAULT_CFG = 2.0 +DEFAULT_SWAY = -1.0 + + +def auk_prompt( + instruction: str, + audio: tuple[np.ndarray | torch.Tensor, int] | None = None, + *, + gen_seconds: float | None = None, + sway: float = DEFAULT_SWAY, + t_grid: list[float] | None = None, + vae_sample: bool = False, +) -> dict[str, Any]: + """Build the stage-0 prompt for one AuK request. + + Args: + instruction: The task instruction; for text-only requests the + no-reference marker is appended if it is not already there. + audio: Reference or source clip as ``(waveform, sample_rate)``. + ``None`` selects the text-only (instruct TTS) path. + gen_seconds: Target duration. ``None`` reuses the source clip's + length and is an error for text-only requests. + sway: Sway coefficient for the Euler time grid. + t_grid: Explicit time grid, overriding ``sway``. + vae_sample: Draw the reference latent from the VAE posterior instead + of taking its mean. + + Returns: + A prompt dict for ``Omni.generate``. + """ + has_audio = audio is not None and not (isinstance(audio, (tuple, list)) and (len(audio) == 0 or audio[0] is None)) + if not has_audio and not instruction.endswith(NO_PROMPT_AUDIO): + instruction = instruction + NO_PROMPT_AUDIO + + prompt: dict[str, Any] = { + "prompt": _CHATML.format( + instruction=instruction, + audio_part=_AUDIO_PART if has_audio else "", + ), + "additional_information": { + "auk": { + "gen_seconds": gen_seconds, + "sway": sway, + "t_grid": list(t_grid) if t_grid is not None else None, + "vae_sample": vae_sample, + "has_audio": has_audio, + } + }, + } + if has_audio: + prompt["multi_modal_data"] = {"audio": audio} + return prompt + + +def auk_sampling_params( + *, + nfe: int = DEFAULT_NFE, + cfg: float = DEFAULT_CFG, + seed: int | None = 0, +) -> list[SamplingParams | OmniDiffusionSamplingParams]: + """Build the per-stage sampling params for one AuK request. + + Stage 0 is a prefill-only encoder, so its vLLM sampling is a no-op that + only has to satisfy the scheduler. Stage 1 carries the real knobs: ``nfe`` + Euler steps and CFG strength. AuK-Flash ignores both and pins its + distilled recipe. + """ + return [ + SamplingParams(max_tokens=1, temperature=0.0, detokenize=False, seed=seed), + OmniDiffusionSamplingParams(num_inference_steps=nfe, guidance_scale=cfg, seed=seed), + ] + + +def resolve_gen_frames( + gen_seconds: float | None, + ref_frames: int, + *, + sample_rate: int = SAMPLE_RATE, + hop: int = HOP, +) -> int: + """Resolve the target latent length in frames. + + An explicit duration wins; otherwise the target is as long as the source + clip. A text-only request with no duration has nothing to generate, which + is an error rather than a silent one-frame output. + """ + if gen_seconds is not None: + if gen_seconds <= 0: + raise ValueError(f"gen_seconds must be positive; got {gen_seconds}") + return max(1, math.ceil(gen_seconds * sample_rate / hop)) + if ref_frames > 0: + return ref_frames + raise ValueError("gen_seconds is required when the request carries no source audio") + + +__all__ = [ + "NO_PROMPT_AUDIO", + "auk_prompt", + "auk_sampling_params", + "resolve_gen_frames", +] diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index 6beb9101866..fdbfb4b6017 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -34,11 +34,13 @@ "MingFlashOmniConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni", "Qwen3VLMoeVisionConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni", "WhisperEncoderConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni", + "AuKConfig": "vllm_omni.transformers_utils.configs.auk", "SenseNovaU1Config": "vllm_omni.transformers_utils.configs.sensenova_u1", } __all__ = [ "Audio8TTSConfig", + "AuKConfig", "Audio8TTSFastARConfig", "Audio8TTSSlowARConfig", "HiggsAudioV3Config", @@ -82,6 +84,7 @@ def __dir__(): # Eagerly import all config modules so their AutoConfig.register() side-effects # run as soon as `vllm_omni.transformers_utils.configs` is imported. from vllm_omni.transformers_utils.configs import audio8_tts as _audio8_tts # noqa: F401, E402 +from vllm_omni.transformers_utils.configs import auk as _auk # noqa: F401, E402 from vllm_omni.transformers_utils.configs import cosyvoice3 as _cosyvoice3 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import dots_tts as _dots_tts # noqa: F401, E402 from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech # noqa: F401, E402 diff --git a/vllm_omni/transformers_utils/configs/auk.py b/vllm_omni/transformers_utils/configs/auk.py new file mode 100644 index 00000000000..d1e8fa805b6 --- /dev/null +++ b/vllm_omni/transformers_utils/configs/auk.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Config for the AuK instruction-driven audio generation model. + +The upstream checkpoint ships ``config.yaml`` plus loose safetensors, so an +assembled directory is produced by ``tools/prepare_auk_checkpoint.py``: the +Qwen2.5-Omni-3B ``config.json`` with this ``model_type``, AuK's architecture, +and the five sections below added at the top level. + +Subclassing ``Qwen2_5OmniConfig`` is load-bearing, not stylistic. The encoder +stage reuses the in-tree Qwen2.5-Omni thinker processor, whose +``get_hf_config`` resolves ``self.ctx.get_hf_config(Qwen2_5OmniConfig)`` and +asserts the type, so a bare ``PretrainedConfig`` subclass would fail before +the first request. Inheriting also leaves the parent's nested +``thinker_config`` / ``talker_config`` / ``token2wav_config`` handling and its +``get_text_config`` in place, which is what the encoder stage reads. +""" + +from typing import Any + +from transformers import AutoConfig +from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import Qwen2_5OmniConfig + +# The distilled AuK-Flash checkpoint pins a 4-step time grid recorded in its +# metadata. The values are exact and must never be rounded: the student was +# trained on them, and a rounded grid changes the output amplitude. +AUK_FLASH_T_GRID: tuple[float, ...] = ( + 0.0, + 0.07612049579620361, + 0.2928932309150696, + 0.6173166036605835, + 1.0, +) + +DEFAULT_DIT: dict[str, Any] = { + "dim": 1536, + "heads": 24, + "dim_head": 64, + "ff_mult": 2, + "text_hidden_dim": 2048, + "num_layers": 10, + "num_single_layers": 20, + "attn_mask_enabled": True, +} + +DEFAULT_VAE: dict[str, Any] = { + "latent_dim": 64, + "downsample_rate": 480, + "target_sample_rate": 24000, + "model_init_kwargs": { + "upsample_rates": [5, 4, 3, 2, 2, 2], + "upsample_kernel_sizes": [10, 8, 6, 4, 4, 4], + "upsample_initial_channel": 1536, + "resblock_kernel_sizes": [3, 7, 11], + "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + "downsample_rates": [2, 2, 2, 3, 4, 5], + "downsample_channels": [12, 24, 48, 96, 192, 384, 768], + "snake_logscale": True, + "latent_dim": 64, + "use_vae": True, + "causal": True, + "flow_hidden_channels": 256, + "act_causal": True, + }, +} + +# Base AuK sampling recipe; AuK-Flash replaces nfe and cfg with its own. +DEFAULT_SAMPLING: dict[str, Any] = {"nfe": 32, "cfg": 2.0, "sway": -1.0} + +VARIANTS = ("base", "flash") + + +class AuKConfig(Qwen2_5OmniConfig): + """Assembled AuK checkpoint config: a Qwen2.5-Omni config plus DiT and VAE.""" + + model_type = "auk" + + def __init__( + self, + dit: dict[str, Any] | None = None, + vae: dict[str, Any] | None = None, + variant: str = "base", + flash_t_grid: list[float] | None = None, + defaults: dict[str, Any] | None = None, + **kwargs: Any, + ) -> None: + if variant not in VARIANTS: + raise ValueError(f"AuK variant must be one of {VARIANTS}; got {variant!r}") + + # Assigned before the parent's init so nothing it runs during + # validation sees a half-built config. + self.dit = {**DEFAULT_DIT, **(dit or {})} + self.vae = {**DEFAULT_VAE, **(vae or {})} + self.variant = variant + self.flash_t_grid = list(flash_t_grid) if flash_t_grid is not None else list(AUK_FLASH_T_GRID) + self.defaults = {**DEFAULT_SAMPLING, **(defaults or {})} + + kwargs.setdefault("architectures", ["AuKForConditionalGeneration"]) + # The parent owns thinker_config / talker_config / token2wav_config and + # get_text_config; nothing here overrides them. + super().__init__(**kwargs) + + @property + def is_flash(self) -> bool: + """Whether this checkpoint is the distilled 4-step student.""" + return getattr(self, "variant", "base") == "flash" + + @property + def latent_dim(self) -> int: + return int(self.vae["latent_dim"]) + + @property + def downsample_rate(self) -> int: + return int(self.vae["downsample_rate"]) + + @property + def target_sample_rate(self) -> int: + return int(self.vae["target_sample_rate"]) + + +AutoConfig.register("auk", AuKConfig) + +__all__ = [ + "AUK_FLASH_T_GRID", + "DEFAULT_DIT", + "DEFAULT_SAMPLING", + "DEFAULT_VAE", + "VARIANTS", + "AuKConfig", +] From 6fcab74aa9b048ab16b2a3815484649ec1ea32f6 Mon Sep 17 00:00:00 2001 From: Joshna-Medisetty Date: Fri, 11 Sep 2026 13:54:28 -0700 Subject: [PATCH 81/93] [XPU][Docker] Align XPU image and CI with vLLM v0.29.0 (#7441) Signed-off-by: Joshna-Medisetty --- .buildkite/intel/pipeline-intel.yml | 2 +- docker/Dockerfile.xpu | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.buildkite/intel/pipeline-intel.yml b/.buildkite/intel/pipeline-intel.yml index d6134b0280f..1e2f86af26a 100644 --- a/.buildkite/intel/pipeline-intel.yml +++ b/.buildkite/intel/pipeline-intel.yml @@ -10,7 +10,7 @@ steps: DOCKER_BUILDKIT: "1" # Buildkite will automatically replace this with the actual commit hash VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}" - VLLM_VERSION: "v0.28.0" + VLLM_VERSION: "v0.29.0" priority: 100 timeout_in_minutes: 60 soft_fail: false diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index acfa71588bd..018bc4316ca 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -1,4 +1,4 @@ -# Target vLLM version: v0.28.0. Keep VLLM_VERSION aligned with the target vLLM +# Target vLLM version: v0.29.0. Keep VLLM_VERSION aligned with the target vLLM # release on every rebase. # # This file layers vLLM-Omni on top of the published upstream XPU image, which @@ -18,7 +18,7 @@ # # .buildkite/intel/scripts/run-xpu-test.sh automates that pull-then-build fallback. ARG VLLM_BASE_IMAGE=vllm/vllm-openai-xpu -ARG VLLM_VERSION=v0.28.0 +ARG VLLM_VERSION=v0.29.0 ARG VLLM_BASE=${VLLM_BASE_IMAGE}:${VLLM_VERSION} FROM ${VLLM_BASE} AS vllm-omni From 13b85c562248831a7c5f527cc08531dc72c0fbad Mon Sep 17 00:00:00 2001 From: MaciejBalaNV Date: Sat, 12 Sep 2026 00:59:49 +0200 Subject: [PATCH 82/93] [Bugfix] Add explicit error when using CFGP with distilled Cosmos3 models (#7427) Signed-off-by: Maciej Bala --- .../diffusion/parallelism/cfg_parallel.md | 6 +++ .../models/cosmos3/test_cosmos3_pipeline.py | 45 ++++++++++++++++--- .../models/cosmos3/pipeline_cosmos3.py | 35 ++++++++------- 3 files changed, 65 insertions(+), 21 deletions(-) diff --git a/docs/user_guide/diffusion/parallelism/cfg_parallel.md b/docs/user_guide/diffusion/parallelism/cfg_parallel.md index 54475440313..1e3b998d8bd 100644 --- a/docs/user_guide/diffusion/parallelism/cfg_parallel.md +++ b/docs/user_guide/diffusion/parallelism/cfg_parallel.md @@ -19,6 +19,12 @@ CFG-Parallel accelerates diffusion models by distributing positive and negative See supported models list in [Supported Models](../../diffusion_features.md#supported-models). +Distilled Cosmos3 checkpoints, including `nvidia/Cosmos3-Super-Image2Video-4Step`, +run without CFG and reject `--cfg-parallel-size` greater than 1 at startup. +Use `--ulysses-degree 2` for two-GPU inference. These checkpoints always use +`guidance_scale=1.0`; other explicitly requested values produce a warning, and +`negative_prompt` does not affect generation. + --- ## Quick Start diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index 302409aebaa..9892a72d744 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -9,6 +9,7 @@ from dataclasses import dataclass from types import SimpleNamespace from typing import Any +from unittest.mock import Mock import numpy as np import pytest @@ -376,23 +377,31 @@ def _tokenize( @pytest.mark.parametrize( - ("provided", "value", "default", "is_distilled", "expected"), + ("provided", "value", "default", "is_distilled", "expected", "expected_warning"), [ - (False, 1.0, 7.0, False, 7.0), - (True, 1.0, 7.0, False, 1.0), - (True, 4.5, 7.0, False, 4.5), - (False, 1.0, 7.0, True, 1.0), - (True, 4.5, 7.0, True, 1.0), + (False, 1.0, 7.0, False, 7.0, False), + (True, 1.0, 7.0, False, 1.0, False), + (True, 4.5, 7.0, False, 4.5, False), + (False, 1.0, 7.0, True, 1.0, False), + (True, 1.0, 7.0, True, 1.0, False), + (True, 4.5, 7.0, True, 1.0, True), + (True, 0.0, 7.0, True, 1.0, True), ], ) def test_resolve_guidance_scale( make_cosmos3_pipeline, + monkeypatch: pytest.MonkeyPatch, provided: bool, value: float, default: float, is_distilled: bool, expected: float, + expected_warning: bool, ) -> None: + from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3 + + warning_once = Mock() + monkeypatch.setattr(pipeline_cosmos3.logger, "warning_once", warning_once) pipeline = make_cosmos3_pipeline() pipeline.is_distilled_model = is_distilled sp = make_sampling_params( @@ -401,6 +410,12 @@ def test_resolve_guidance_scale( ) assert pipeline._resolve_guidance_scale(sp, default) == expected + if expected_warning: + warning_once.assert_called_once() + assert "overridden to 1.0" in warning_once.call_args.args[0] + assert "negative_prompt does not affect generation" in warning_once.call_args.args[0] + else: + warning_once.assert_not_called() def test_distilled_generation_accepts_t2i_and_i2v(make_cosmos3_pipeline) -> None: @@ -735,6 +750,7 @@ def _make_od_config( custom_pipeline_args={}, model_config=model_config or {}, tf_model_config=tf_model_config, + parallel_config=SimpleNamespace(cfg_parallel_size=1, ulysses_degree=1), ) @@ -770,6 +786,7 @@ def test_pipeline_init_uses_flow_unipc_with_cosmos3_defaults(stub_real_pipeline_ assert pipeline._engine_init_flow_shift == 2.5 +@pytest.mark.parametrize("cfg_parallel_size,ulysses_degree", [(1, 1), (1, 2), (2, 1), (2, 2)]) @pytest.mark.parametrize( ("scheduler_class_name", "expected_distilled"), [ @@ -783,6 +800,8 @@ def test_pipeline_resolves_scheduler_class_from_checkpoint_file( monkeypatch: pytest.MonkeyPatch, scheduler_class_name: str, expected_distilled: bool, + cfg_parallel_size: int, + ulysses_degree: int, ) -> None: import json @@ -831,6 +850,20 @@ def from_config(cls, config, **kwargs): od_config = _make_od_config(sound_gen=False) od_config.model = str(tmp_path) + od_config.parallel_config.cfg_parallel_size = cfg_parallel_size + od_config.parallel_config.ulysses_degree = ulysses_degree + if expected_distilled and cfg_parallel_size > 1: + monkeypatch.setattr( + pipeline_cosmos3.AutoTokenizer, + "from_pretrained", + lambda *args, **kwargs: pytest.fail("component loading must not start for distilled CFG parallelism"), + ) + with pytest.raises(ValueError, match="Set --cfg-parallel-size 1 and use --ulysses-degree"): + Cosmos3OmniDiffusersPipeline(od_config=od_config) + assert StubFlowMatchScheduler.from_config_calls == [] + assert StubFlowUniPCScheduler.from_config_calls == [] + return + pipeline = Cosmos3OmniDiffusersPipeline(od_config=od_config) assert pipeline.is_distilled_model is expected_distilled diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 241eae80577..5977e0eeaac 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -948,6 +948,20 @@ def __init__( model_path = od_config.model local_files_only = os.path.exists(model_path) + # Validate guidance parallelism from checkpoint metadata before loading + # components. Distilled checkpoints have only a conditional branch. + scheduler_config = FlowUniPCMultistepScheduler.load_config( + model_path, + subfolder="scheduler", + local_files_only=local_files_only, + ) + self.is_distilled_model = scheduler_config.get("_class_name") == COSMOS3_DISTILLED_CHECKPOINT_SCHEDULER_CLASS + if self.is_distilled_model and od_config.parallel_config.cfg_parallel_size > 1: + raise ValueError( + "Distilled Cosmos3 checkpoints run without classifier-free guidance. " + "Set --cfg-parallel-size 1 and use --ulysses-degree for multi-GPU inference." + ) + # --- Tokenizer --- self.tokenizer = AutoTokenizer.from_pretrained( model_path, @@ -1010,20 +1024,7 @@ def __init__( logger.info("Cosmos3: session state manager enabled (max_sessions=%d)", mm_max_sessions) # --- Scheduler --- - # Distilled model differs from regular one only by scheduler, - # distilled one uses FlowMatchEulerDiscreteScheduler, while - # regular should use FlowUniPCMultistepScheduler - - scheduler_config = FlowUniPCMultistepScheduler.load_config( - model_path, - subfolder="scheduler", - local_files_only=local_files_only, - ) - - scheduler_class_name = scheduler_config.get("_class_name") - - self.is_distilled_model = False - if scheduler_class_name == COSMOS3_DISTILLED_CHECKPOINT_SCHEDULER_CLASS: + if self.is_distilled_model: fixed_step_config = scheduler_config.get("fixed_step_sampler_config") if not isinstance(fixed_step_config, dict) or fixed_step_config.get("sample_type") != "sde": raise ValueError("Cosmos3 distilled scheduler requires fixed_step_sampler_config.sample_type=sde.") @@ -1035,7 +1036,6 @@ def __init__( stochastic_sampling=True, ) self._scheduler_init_t_list = list(t_list) - self.is_distilled_model = True else: # Preserve compatible solver settings from the checkpoint, but keep # the base shift neutral. The concrete request shift is applied when @@ -1811,6 +1811,11 @@ def _resolve_seed(sp: OmniDiffusionSamplingParams, generator: torch.Generator | def _resolve_guidance_scale(self, sp: OmniDiffusionSamplingParams, default: float) -> float: if self.is_distilled_model: + if sp.guidance_scale_provided and float(sp.guidance_scale) != 1.0: + logger.warning_once( + "Distilled Cosmos3 checkpoints run without classifier-free guidance. " + "The requested guidance_scale is overridden to 1.0; negative_prompt does not affect generation." + ) return 1.0 if sp.guidance_scale_provided: return float(sp.guidance_scale) From 220366db005a79e92877c50266358375d828c349 Mon Sep 17 00:00:00 2001 From: DanaerLee Date: Sat, 12 Sep 2026 11:31:05 +0800 Subject: [PATCH 83/93] [Perf][Diffusion] Run MammothModa2 DiT attention through the shared attention layer (#7094) Signed-off-by: MrlixiangWE Co-authored-by: Hongsheng Liu --- .../attention/test_attention_config.py | 170 +++++++++++++++- .../models/ltx2/test_ltx2_transformer.py | 1 + .../models/mammoth_moda2/__init__.py | 0 .../mammoth_moda2/test_dit_attention.py | 189 ++++++++++++++++++ .../mammoth_moda2/test_dit_attention_cuda.py | 151 ++++++++++++++ vllm_omni/diffusion/attention/layer.py | 31 ++- .../mammoth_moda2/mammothmoda2_dit_model.py | 148 +++++--------- 7 files changed, 582 insertions(+), 108 deletions(-) create mode 100644 tests/diffusion/models/mammoth_moda2/__init__.py create mode 100644 tests/diffusion/models/mammoth_moda2/test_dit_attention.py create mode 100644 tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py index 6561615896e..4a053b1bfa8 100644 --- a/tests/diffusion/attention/test_attention_config.py +++ b/tests/diffusion/attention/test_attention_config.py @@ -10,6 +10,7 @@ - AttentionMetadata.extra field """ +from dataclasses import replace from types import SimpleNamespace from unittest.mock import Mock @@ -17,7 +18,12 @@ import torch import vllm_omni.diffusion.attention.layer as layer_mod -from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +import vllm_omni.diffusion.attention.selector as selector_mod +from vllm_omni.diffusion.attention.backends.abstract import ( + AttentionMetadata, + PackedPaddingMetadata, + QueryRange, +) from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionBackend from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.config import ( @@ -477,6 +483,7 @@ class TestAttentionInitUsesCurrentDiffusionConfig: @pytest.mark.parametrize("mindiesd_available", [False, True]) def test_paged_npu_memory_profile_uses_sdpa_only_without_mindiesd(self, monkeypatch, mindiesd_available): attention = Attention.__new__(Attention) + attention.allow_fp32_fallback = False attention._scheduler_paged_kv = True attention.paged_kv_cache_role = "primary" attention.backend_pref = "FLASH_ATTN" @@ -519,6 +526,7 @@ def test_paged_npu_memory_profile_uses_sdpa_only_without_mindiesd(self, monkeypa @pytest.mark.parametrize("fa4_available", [False, True]) def test_paged_cuda_memory_profile_uses_sdpa_without_dense_fa4(self, monkeypatch, fa4_available): attention = Attention.__new__(Attention) + attention.allow_fp32_fallback = False attention._scheduler_paged_kv = True attention.paged_kv_cache_role = "primary" attention.backend_pref = "FLASH_ATTN" @@ -561,6 +569,7 @@ def test_paged_cuda_memory_profile_uses_sdpa_without_dense_fa4(self, monkeypatch def test_dense_flash_import_error_is_not_reclassified(self, monkeypatch): attention = Attention.__new__(Attention) + attention.allow_fp32_fallback = False attention._scheduler_paged_kv = False attention.paged_kv_cache_role = "primary" attention.backend_pref = "FLASH_ATTN" @@ -900,6 +909,7 @@ def _selected_forward(*args): return sentinel fake_attention = SimpleNamespace( + allow_fp32_fallback=False, attention=SimpleNamespace( forward=_selected_forward, ), @@ -938,6 +948,164 @@ def test_ring_attention_rejects_mask_instead_of_ignoring_it(self): Attention._run_ring_attention(fake_attention, query, query, query, metadata) +class TestOptInFloat32Fallback: + @staticmethod + def _make_attention(monkeypatch, *, backend=None, allow=True, kv_cache_dtype=None): + events, state = [], {} + context = object() + + class _SelectedImpl: + def __init__(self, **kwargs): + self.kwargs = kwargs + + @staticmethod + def supports_kv_cache_dtype(*args): + return True + + def forward(self, query, key, value, metadata): + events.append("selected") + state["kernel_call"] = (query, key, value, metadata) + state["kernel_output"] = query + 4 + return state["kernel_output"] + + class _FallbackImpl(_SelectedImpl): + def forward(self, query, key, value, metadata): + output = super().forward(query, key, value, metadata) + events[-1] = "fallback" + return output + + class _FlashBackend(FlashAttentionBackend): + @staticmethod + def get_impl_cls(): + return _SelectedImpl + + class _SDPABackend(layer_mod.SDPABackend): + @staticmethod + def get_impl_cls(): + return _SelectedImpl + + class _Strategy: + def pre_attention(self, query, key, value, metadata): + events.append("pre") + state["prepared"] = ( + query + 1, + key + 2, + value + 3, + replace(metadata) if metadata is not None else None, + ) + return (*state["prepared"], context) + + def post_attention(self, output, ctx): + events.append("post") + assert output is state["kernel_output"] + assert ctx is context + return output + 7 + + # Keep real role/config resolution and construction; replace hardware + # lookup and kernels. These are CPU dispatch tests, not SP collectives. + monkeypatch.setattr( + selector_mod, + "_cached_get_backend_cls", + lambda name, *args, **kwargs: _SDPABackend if name == "TORCH_SDPA" else _FlashBackend, + ) + monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FallbackImpl)) + monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: _Strategy()) + monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False) + config = OmniDiffusionConfig( + diffusion_attention_config={"default": backend} if backend else {}, + diffusion_kv_cache_dtype=kv_cache_dtype, + ) + with set_current_diffusion_config(config): + attention = Attention( + num_heads=4, + num_kv_heads=2, + head_size=8, + causal=False, + softmax_scale=0.25, + **({"allow_fp32_fallback": True} if allow else {}), + ) + return attention, events, state + + @pytest.mark.parametrize( + "backend,allow,dtype,is_cuda,expected", + [ + pytest.param(None, True, torch.float32, True, "fallback", id="auto-fp32"), + pytest.param(None, False, torch.float32, True, "selected", id="default-opt-out"), + pytest.param("FLASH_ATTN", True, torch.float32, True, "selected", id="explicit-flash"), + pytest.param("TORCH_SDPA", True, torch.float32, True, "selected", id="explicit-sdpa"), + pytest.param(None, True, torch.bfloat16, True, "selected", id="bf16"), + pytest.param(None, True, torch.float32, False, "selected", id="cpu-fp32"), + ], + ) + def test_public_forward_preserves_strategy_and_backend_choice( + self, monkeypatch, backend, allow, dtype, is_cuda, expected + ): + attention, events, state = self._make_attention(monkeypatch, backend=backend, allow=allow) + # Only the dispatch predicate is simulated; every tensor stays on CPU. + monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: is_cuda)) + query = torch.ones(1, 2, 4, 8, dtype=dtype) + key = torch.full((1, 2, 2, 8), 2, dtype=dtype) + value = torch.full_like(key, 3) + metadata = AttentionMetadata(attn_mask=torch.tensor([[True, False]])) + + output = attention(query, key, value, metadata) + + assert events == ["pre", expected, "post"] + assert all(got is prepared for got, prepared in zip(state["kernel_call"], state["prepared"])) + assert state["kernel_call"][3] is not metadata + assert state["kernel_call"][3].attn_mask is metadata.attn_mask + torch.testing.assert_close(output, query + 12) + assert attention.backend_explicit is (backend is not None) + impl = attention.sdpa_fallback if expected == "fallback" else attention.attention + assert impl.kwargs["num_kv_heads"] == 2 + assert impl.kwargs["softmax_scale"] == 0.25 + + def test_explicit_backend_failure_is_not_retried_with_sdpa(self, monkeypatch): + attention, events, _ = self._make_attention(monkeypatch, backend="FLASH_ATTN") + monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True)) + selected = Mock(side_effect=RuntimeError("selected backend rejects this dtype")) + monkeypatch.setattr(attention.attention, "forward", selected) + query = torch.ones(1, 2, 4, 8) + + with pytest.raises(RuntimeError, match="selected backend rejects this dtype"): + attention(query, query, query) + + selected.assert_called_once() + assert events == ["pre"] + + @pytest.mark.parametrize("contract", ["packed", "packed-padding", "piecewise", "quantized-kv"]) + def test_backend_metadata_is_not_silently_sent_to_dense_sdpa(self, monkeypatch, contract): + quant_dtype = "fp8" if contract == "quantized-kv" else None + attention, events, state = self._make_attention(monkeypatch, kv_cache_dtype=quant_dtype) + monkeypatch.setattr(torch.Tensor, "is_cuda", property(lambda self: True)) + lengths = torch.tensor([0, 1, 2], dtype=torch.int32) + metadata = AttentionMetadata() + if contract == "packed": + metadata.extra = { + "cu_seqlens_q": lengths, + "cu_seqlens_k": lengths, + "max_seqlen_q": 1, + "max_seqlen_k": 1, + } + elif contract == "packed-padding": + metadata.packed_padding = PackedPaddingMetadata(1, 1, lengths[:2], lengths[:2]) + elif contract == "piecewise": + metadata.full_attn_spans = [[(0, 1)]] + metadata.query_ranges = (QueryRange(0, 2, 0),) + query = torch.ones(1, 2, 4, 8) + + attention(query, query, query, metadata) + + assert events == ["pre", "selected", "post"] + forwarded = state["kernel_call"][3] + if quant_dtype: + # Quantization is injected by the public forward from config. + assert forwarded.extra == {"kv_cache_dtype": "fp8"} + assert metadata.extra == {} + else: + assert forwarded is state["prepared"][3] + + class TestDiffusionKvCacheQuantization: @staticmethod def _install_attention_init_stubs(monkeypatch, *, supports_kv_cache_dtype=True): diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer.py b/tests/diffusion/models/ltx2/test_ltx2_transformer.py index f12fbc466b5..679b7ca37a7 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_transformer.py +++ b/tests/diffusion/models/ltx2/test_ltx2_transformer.py @@ -233,6 +233,7 @@ def forward(query, _key, _value, _metadata): layer = object.__new__(ltx2_transformer._LTX2ParallelAttention) nn.Module.__init__(layer) + layer.allow_fp32_fallback = False layer.attn_backend = Backend layer.attention = SimpleNamespace(forward=run("native")) layer.sdpa_fallback = SimpleNamespace(forward=run("sdpa")) diff --git a/tests/diffusion/models/mammoth_moda2/__init__.py b/tests/diffusion/models/mammoth_moda2/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/diffusion/models/mammoth_moda2/test_dit_attention.py b/tests/diffusion/models/mammoth_moda2/test_dit_attention.py new file mode 100644 index 00000000000..bab2fee7d2f --- /dev/null +++ b/tests/diffusion/models/mammoth_moda2/test_dit_attention.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""MammothModa2 DiT self-attention runs through the shared Omni attention layer. + +``_reference_attention`` is the arithmetic the processor used before it was +wired to the shared layer: SDPA over K/V replicated to the query head count, +padded rows zeroed. Both sides read the same module weights, so any drift in +the projections, the QK RMSNorm, the interleaved real RoPE, the softmax scale +or the padding mask shows up as a mismatch here. +""" + +import pytest +import torch +import torch.nn.functional as F + +from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum +from vllm_omni.diffusion.config import set_current_diffusion_config +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import ( + TransformerBlock, + apply_real_rotary_emb, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# Small stand-in for the real 2520 / 21 / 7 (head_dim 120) layout. +DIM, HEADS = 48, 6 +SEQ, BATCH = 37, 2 + +# The shared layer resolves its backend from the diffusion config at construction; +# without one it takes the platform default, which on a CUDA host is FLASH_ATTN +# and cannot run CPU tensors. Pin SDPA the way a deployment would. +_SDPA_CONFIG = OmniDiffusionConfig(diffusion_attention_config={"default": {"backend": "TORCH_SDPA"}}) + + +def _block(kv_heads: int, modulation: bool = True) -> TransformerBlock: + torch.manual_seed(0) + with set_current_diffusion_config(_SDPA_CONFIG): + block = TransformerBlock( + DIM, + HEADS, + kv_heads, + multiple_of=8, + ffn_dim_multiplier=1.0, + norm_eps=1e-5, + modulation=modulation, + ) + return block.eval() + + +def _inputs(head_dim: int): + torch.manual_seed(1) + hidden = torch.randn(BATCH, SEQ, DIM) + mask = torch.ones(BATCH, SEQ, dtype=torch.bool) + mask[0, SEQ - 9 :] = False + mask[1, SEQ - 2 :] = False + angles = torch.rand(1, SEQ, head_dim) * 6.283 + rotary = (angles.cos(), angles.sin()) + temb = torch.randn(BATCH, min(DIM, 1024)) + return hidden, mask, rotary, temb + + +def _reference_attention(attn, hidden, mask, rotary): + """The pre-change processor, kept verbatim as the oracle.""" + batch, seq, _ = hidden.shape + query, key, value = attn.to_q(hidden), attn.to_k(hidden), attn.to_v(hidden) + head_dim = query.shape[-1] // attn.heads + kv_heads = key.shape[-1] // head_dim + query = attn.norm_q(query.view(batch, seq, attn.heads, head_dim)) + key = attn.norm_k(key.view(batch, seq, kv_heads, head_dim)) + value = value.view(batch, seq, kv_heads, head_dim) + query = apply_real_rotary_emb(query, rotary[0], rotary[1]) + key = apply_real_rotary_emb(key, rotary[0], rotary[1]) + query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2) + if kv_heads < attn.heads: + key = key.repeat_interleave(attn.heads // kv_heads, dim=1) + value = value.repeat_interleave(attn.heads // kv_heads, dim=1) + out = F.scaled_dot_product_attention(query, key, value, attn_mask=mask.view(batch, 1, 1, seq), scale=attn.scale) + out = (out * mask[:, None, :, None]).transpose(1, 2).reshape(batch, seq, attn.heads * head_dim) + return attn.to_out[1](attn.to_out[0](out)) + + +@pytest.mark.parametrize("kv_heads", [2, HEADS], ids=["gqa_3to1", "mha"]) +def test_processor_matches_previous_arithmetic(kv_heads): + block = _block(kv_heads) + hidden, mask, rotary, _ = _inputs(block.head_dim) + with torch.no_grad(): + got = block.attn( + hidden_states=hidden, encoder_hidden_states=hidden, attention_mask=mask, image_rotary_emb=rotary + ) + want = _reference_attention(block.attn, hidden, mask, rotary) + torch.testing.assert_close(got, want, atol=1e-5, rtol=1e-5) + assert torch.count_nonzero(got[~mask]) == 0, "padded rows must stay zero" + + +def test_no_mask_equals_all_valid_mask(): + block = _block(2) + hidden, _, rotary, _ = _inputs(block.head_dim) + full = torch.ones(BATCH, SEQ, dtype=torch.bool) + with torch.no_grad(): + unmasked = block.attn( + hidden_states=hidden, encoder_hidden_states=hidden, attention_mask=None, image_rotary_emb=rotary + ) + masked = block.attn( + hidden_states=hidden, encoder_hidden_states=hidden, attention_mask=full, image_rotary_emb=rotary + ) + torch.testing.assert_close(unmasked, masked, atol=1e-6, rtol=1e-6) + + +@pytest.mark.parametrize("modulation", [True, False], ids=["noise_block", "context_refiner"]) +def test_block_forward_hands_native_kv_heads_to_shared_layer(monkeypatch, modulation): + """Through the block's real forward, on both of its branches (the modulated + noise/refiner blocks and the unmodulated context refiner): the shared layer + is what runs, and it receives K/V with their native head count -- no + replication before the call.""" + block = _block(2, modulation=modulation) + hidden, mask, rotary, temb = _inputs(block.head_dim) + if not modulation: + temb = None + seen = [] + original = block.attn.omni_attn.forward + + def spy(query, key, value, attn_metadata=None): + seen.append((tuple(query.shape), tuple(key.shape), tuple(value.shape), attn_metadata)) + return original(query, key, value, attn_metadata) + + monkeypatch.setattr(block.attn.omni_attn, "forward", spy) + with torch.no_grad(): + out = block(hidden, mask, rotary, temb) + assert out.shape == hidden.shape and torch.isfinite(out).all() + assert len(seen) == 1 + q_shape, k_shape, v_shape, metadata = seen[0] + assert q_shape == (BATCH, SEQ, HEADS, block.head_dim) + assert k_shape == v_shape == (BATCH, SEQ, 2, block.head_dim) + assert metadata is not None and metadata.attn_mask.dtype == torch.bool + assert torch.equal(metadata.attn_mask, mask) + + +def test_backend_selection_reaches_the_dit_layer(): + """The diffusion config selects the backend when the DiT layer is constructed.""" + assert _block(2).attn.omni_attn.attn_backend is DiffusionAttentionBackendEnum.TORCH_SDPA.get_class() + with set_current_diffusion_config( + OmniDiffusionConfig(diffusion_attention_config={"default": {"backend": "FLASH_ATTN"}}) + ): + block = TransformerBlock(DIM, HEADS, 2, multiple_of=8, ffn_dim_multiplier=1.0, norm_eps=1e-5) + assert block.attn.omni_attn.attn_backend is DiffusionAttentionBackendEnum.FLASH_ATTN.get_class() + + +def test_shared_layer_owns_no_parameters_and_keeps_checkpoint_keys(): + block = _block(2) + assert sum(p.numel() for p in block.attn.omni_attn.parameters()) == 0 + assert not list(block.attn.omni_attn.buffers()) + keys = set(block.state_dict()) + assert not any("omni_attn" in key for key in keys) + # The projections and QK norms the checkpoint provides are still where they were. + for expected in ( + "attn.to_q.weight", + "attn.to_k.weight", + "attn.to_v.weight", + "attn.to_out.0.weight", + "attn.norm_q.weight", + "attn.norm_k.weight", + ): + assert expected in keys, expected + + +@pytest.mark.parametrize("with_mask", [True, False], ids=["empty_mask", "no_mask"]) +def test_empty_text_stream_skips_the_kernel(monkeypatch, with_mask): + """CFG's unconditional branch carries zero text tokens (the pipeline's default + negative_prompt_embeds has no rows), so the context refiner attends over an + empty sequence. The flash-attention varlen fallback cannot take that + (arange step 0), so the processor must not reach the shared layer at all.""" + block = _block(kv_heads=2, modulation=False) + attn = block.attn + + def _never(*args, **kwargs): + raise AssertionError("shared layer called with an empty sequence") + + monkeypatch.setattr(attn.omni_attn, "forward", _never) + hidden = torch.randn(BATCH, 0, DIM) + mask = torch.ones(BATCH, 0, dtype=torch.bool) if with_mask else None + angles = torch.rand(1, 0, block.head_dim) + out = attn( + hidden_states=hidden, + encoder_hidden_states=hidden, + attention_mask=mask, + image_rotary_emb=(angles.cos(), angles.sin()), + ) + assert out.shape == (BATCH, 0, DIM) diff --git a/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py b/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py new file mode 100644 index 00000000000..62415953a91 --- /dev/null +++ b/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Real-shape check of the MammothModa2 DiT attention on CUDA: the shared +backend (FLASH_ATTN by default, head_dim 120) against the pre-change arithmetic +in bf16, with the padding mask the joint text+image sequence carries.""" + +import pytest +import torch + +import vllm_omni.diffusion.attention.backends.sdpa as sdpa_backend +from tests.helpers.mark import hardware_marks +from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import TransformerBlock + +from .test_dit_attention import _reference_attention + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.advanced_model, + pytest.mark.cuda, + *hardware_marks(res={"cuda": "L4"}, num_cards=1), +] + +# MammothModa2-Preview gen_dit_config: hidden 2520, 21 heads, 7 KV heads, head_dim 120. +DIM, HEADS, KV_HEADS = 2520, 21, 7 + + +@pytest.mark.parametrize("seq", [77 + 4096, 512], ids=["t2i_1024", "short"]) +def test_real_shape_matches_previous_arithmetic_bf16(seq): + torch.manual_seed(0) + block = ( + TransformerBlock(DIM, HEADS, KV_HEADS, multiple_of=256, ffn_dim_multiplier=1.0, norm_eps=1e-5) + .cuda() + .to(torch.bfloat16) + .eval() + ) + hidden = torch.randn(2, seq, DIM, device="cuda", dtype=torch.bfloat16) + mask = torch.ones(2, seq, dtype=torch.bool, device="cuda") + mask[0, seq - 300 :] = False + mask[1, seq - 17 :] = False + angles = torch.rand(1, seq, block.head_dim, device="cuda") * 6.283 + rotary = (angles.cos().to(torch.bfloat16), angles.sin().to(torch.bfloat16)) + with torch.no_grad(): + got = block.attn( + hidden_states=hidden, encoder_hidden_states=hidden, attention_mask=mask, image_rotary_emb=rotary + ) + want = _reference_attention(block.attn, hidden, mask, rotary) + assert torch.isfinite(got).all() + assert torch.count_nonzero(got[~mask]) == 0 + diff = (got.float() - want.float()).abs()[mask] + # bf16 kernels differ in accumulation order; measured on A800: max 1.0e-3, mean 5e-5. + assert diff.max().item() < 2e-2, diff.max().item() + assert diff.mean().item() < 1e-3, diff.mean().item() + + +def test_empty_text_stream_on_the_default_backend(): + """The recipe's text-to-image request with text_guidance_scale > 1 runs the + context refiner on a zero-token unconditional prompt. Before the guard this + raised ``RuntimeError: step must be nonzero`` from the FA varlen fallback.""" + torch.manual_seed(0) + block = ( + TransformerBlock(DIM, HEADS, KV_HEADS, multiple_of=256, ffn_dim_multiplier=1.0, norm_eps=1e-5, modulation=False) + .cuda() + .to(torch.bfloat16) + .eval() + ) + hidden = torch.randn(1, 0, DIM, device="cuda", dtype=torch.bfloat16) + mask = torch.ones(1, 0, dtype=torch.bool, device="cuda") + angles = torch.rand(1, 0, block.head_dim, device="cuda") + with torch.no_grad(): + out = block.attn( + hidden_states=hidden, + encoder_hidden_states=hidden, + attention_mask=mask, + image_rotary_emb=(angles.cos().to(torch.bfloat16), angles.sin().to(torch.bfloat16)), + ) + assert out.shape == (1, 0, DIM) + + +@pytest.mark.parametrize( + ("seq", "force_gqa_fallback"), + [(512, False), (77 + 4096, False), (32, True)], + ids=["short", "t2i_1024", "unsupported_native_gqa"], +) +def test_fp32_falls_back_to_sdpa_and_matches_reference(monkeypatch, seq, force_gqa_fallback): + """FP32 keeps its dtype and reaches shared SDPA's GQA compatibility check. + + Cover the long masked shape that otherwise falls back to quadratic math + attention, and force the unsupported-GQA branch independently of hardware. + """ + torch.manual_seed(0) + dt = torch.float32 + block = ( + TransformerBlock(DIM, HEADS, KV_HEADS, multiple_of=256, ffn_dim_multiplier=1.0, norm_eps=1e-5) + .cuda() + .to(dt) + .eval() + ) + hidden = torch.randn(1, seq, DIM, device="cuda", dtype=dt) + mask = torch.ones(1, seq, dtype=torch.bool, device="cuda") + mask[:, -13:] = False + angles = torch.rand(1, seq, block.head_dim, device="cuda") + rotary = (angles.cos().to(dt), angles.sin().to(dt)) + capabilities, calls, phases = [], [], [] + can_use_fused_gqa = sdpa_backend.can_sdpa_use_fused_gqa + sdpa = torch.nn.functional.scaled_dot_product_attention + strategy = block.attn.omni_attn._get_active_parallel_strategy() + pre_attention, post_attention = strategy.pre_attention, strategy.post_attention + + def prepare(*args): + phases.append("pre") + return pre_attention(*args) + + def restore(*args): + phases.append("post") + return post_attention(*args) + + def check_gqa(*args): + supported = False if force_gqa_fallback else can_use_fused_gqa(*args) + capabilities.append(supported) + return supported + + def record_sdpa(query, key, value, **kwargs): + phases.append("sdpa") + calls.append((query.shape, key.shape, value.shape, query.dtype, kwargs)) + return sdpa(query, key, value, **kwargs) + + with torch.no_grad(), monkeypatch.context() as patch: + patch.setattr(sdpa_backend, "can_sdpa_use_fused_gqa", check_gqa) + patch.setattr(torch.nn.functional, "scaled_dot_product_attention", record_sdpa) + patch.setattr(strategy, "pre_attention", prepare) + patch.setattr(strategy, "post_attention", restore) + got = block.attn(hidden, hidden, attention_mask=mask, image_rotary_emb=rotary) + + assert phases == ["pre", "sdpa", "post"], "FP32 must preserve the shared parallel dispatch" + assert len(capabilities) == len(calls) == 1, "FP32 must use shared SDPA's runtime GQA check" + q_shape, k_shape, v_shape, dtype, kwargs = calls[0] + expected_kv_heads = KV_HEADS if capabilities[0] else HEADS + assert q_shape == (1, HEADS, seq, block.head_dim) + assert k_shape == v_shape == (1, expected_kv_heads, seq, block.head_dim) + assert dtype == got.dtype == torch.float32 + assert kwargs["enable_gqa"] == capabilities[0] + assert kwargs["scale"] == block.attn.scale + assert kwargs["is_causal"] is False + assert torch.equal(kwargs["attn_mask"], mask[:, None, None, :]) + + with torch.no_grad(): + want = _reference_attention(block.attn, hidden, mask, rotary) + assert torch.isfinite(got).all() + assert torch.count_nonzero(got[~mask]) == 0 + diff = (got.float() - want.float()).abs()[mask] + assert diff.max().item() < 1e-3, diff.max().item() diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 4b8ca8f4980..4c588af0cfb 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -81,6 +81,8 @@ def __init__( # varlen attention with learned sink logits). The shared Attention # layer still owns parallel dispatch and compile boundaries. custom_attention: nn.Module | None = None, + # Preserve dense FP32 inference for models opting into CUDA auto fallback. + allow_fp32_fallback: bool = False, ): super().__init__() @@ -193,8 +195,8 @@ def __init__( role=role, backend_explicit=self.backend_explicit, ) - # Some model-specific paths call this helper directly. The main backend - # dispatch below never switches to it implicitly. + # Compatibility kernels run inside shared dispatch, between the + # parallel strategy's input preparation and output restoration. self.sdpa_fallback = SDPABackend.get_impl_cls()( num_heads=num_heads, head_size=head_size, @@ -221,6 +223,7 @@ def __init__( self.use_sync = use_sync self.causal = causal self.skip_sequence_parallel = skip_sequence_parallel + self.allow_fp32_fallback = allow_fp32_fallback self.use_ring = False self.ring_pg = None @@ -497,13 +500,33 @@ def _run_local_attention(self, query, key, value, attn_metadata): self._assert_metadata_compatible(attn_metadata) + if ( + self.allow_fp32_fallback + and query.is_cuda + and query.dtype == torch.float32 + and query.ndim == 4 + and not self.backend_explicit + and self.attn_backend.get_name() == "FLASH_ATTN" + and ( + attn_metadata is None + or ( + attn_metadata.full_attn_spans is None + and attn_metadata.query_ranges is None + and attn_metadata.video_layout is None + and attn_metadata.packed_padding is None + and not attn_metadata.extra + ) + ) + ): + logger.warning_once("Using SDPA for this layer's FP32 input with automatic CUDA FlashAttention selection.") + return self.sdpa_fallback.forward(query, key, value, attn_metadata) + in_kv_memory_profile = is_forward_context_available() and get_forward_context().in_diffusion_kv_memory_profile # The startup KV-capacity profile needs tensor shapes, not a paged # attention result. If dense FLASH_ATTN deps are absent (NPU MindIE-SD # or CUDA CuTe FA4), SDPA provides that profile forward. Formal paged # requests never use this branch because their Worker adapter is active. - # Do not silently remap float32 (or other dtypes) to SDPA: a - # user-explicit backend must run or raise. + # A user-explicit backend must run or raise. if ( self._scheduler_paged_kv and self.paged_kv_cache_role is not None diff --git a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py b/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py index 000abe05193..e2c086d4a43 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py +++ b/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py @@ -1,4 +1,5 @@ -import math +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import torch import torch.nn.functional as F @@ -6,27 +7,14 @@ from diffusers.models.attention_processor import Attention from diffusers.models.embeddings import TimestepEmbedding, Timesteps from diffusers.models.modeling_utils import ModelMixin -from einops import rearrange, repeat +from einops import rearrange from torch import nn from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm -from .rope_real import RotaryPosEmbedReal +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.layer import Attention as OmniAttention -try: - from transformers.modeling_flash_attention_utils import ( # type: ignore - flash_attn_varlen_func, # pyright: ignore[reportAttributeAccessIssue] - is_flash_attn_available, - ) -except Exception: # pragma: no cover - best-effort compatibility - flash_attn_varlen_func = None # type: ignore[assignment] - - def is_flash_attn_available() -> bool: # type: ignore[override] - return False - - -from .rope_real import apply_real_rotary_emb - -_HAS_FLASH_ATTN_VARLEN = bool(is_flash_attn_available()) and flash_attn_varlen_func is not None +from .rope_real import RotaryPosEmbedReal, apply_real_rotary_emb class LuminaRMSNormZero(nn.Module): @@ -278,9 +266,14 @@ def forward(self, x: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor class AttnProcessor: - def __init__(self) -> None: - if not hasattr(F, "scaled_dot_product_attention"): - raise ImportError("AttnProcessor requires PyTorch 2.0+ (F.scaled_dot_product_attention).") + """Self-attention for the MammothModa2 DiT through the shared Omni attention layer. + + The model-specific parts stay here: the projections, the per-head QK RMSNorm + and the interleaved real RoPE. The kernel, the padding mask and the + grouped-query heads are the shared backend's job, so K/V are handed over + with their native head count instead of being replicated to the query + head count first. + """ def __call__( self, @@ -289,118 +282,55 @@ def __call__( encoder_hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, image_rotary_emb: torch.Tensor | None = None, - base_sequence_length: int | None = None, ) -> torch.Tensor: batch_size, sequence_length, _ = hidden_states.shape - # Get Query-Key-Value Pair query = attn.to_q(hidden_states) key = attn.to_k(encoder_hidden_states) value = attn.to_v(encoder_hidden_states) - query_dim = query.shape[-1] - inner_dim = key.shape[-1] - head_dim = query_dim // attn.heads + head_dim = query.shape[-1] // attn.heads + kv_heads = key.shape[-1] // head_dim dtype = query.dtype - # Get key-value heads - kv_heads = inner_dim // head_dim + if sequence_length == 0 or key.shape[1] == 0: + # The pipeline's default unconditional branch is a text stream with + # zero tokens (negative_prompt_embeds has no rows), so under CFG the + # context refiner attends over an empty sequence. SDPA returns an + # empty tensor for that; the flash-attention varlen fallback builds + # cu_seqlens with arange(step=seq_len) and rejects step 0. Nothing + # to attend to either way, so hand the projection an empty input. + empty = query.new_zeros(batch_size, sequence_length, attn.heads * head_dim) + return attn.to_out[1](attn.to_out[0](empty)) - # Reshape tensors for attention computation query = query.view(batch_size, -1, attn.heads, head_dim) key = key.view(batch_size, -1, kv_heads, head_dim) value = value.view(batch_size, -1, kv_heads, head_dim) - # Apply Query-Key normalization if attn.norm_q is not None: query = attn.norm_q(query) if attn.norm_k is not None: key = attn.norm_k(key) - # Apply Rotary Position Embeddings if image_rotary_emb is not None: query = apply_real_rotary_emb(query, image_rotary_emb[0], image_rotary_emb[1]) key = apply_real_rotary_emb(key, image_rotary_emb[0], image_rotary_emb[1]) query, key = query.to(dtype), key.to(dtype) - # Calculate attention scale - if base_sequence_length is not None: - softmax_scale = math.sqrt(math.log(sequence_length, base_sequence_length)) * attn.scale - else: - softmax_scale = attn.scale - - if _HAS_FLASH_ATTN_VARLEN and attention_mask is not None and hidden_states.is_cuda: - # Flash-Attn varlen expects packed tokens + cu_seqlens. Here we only need - # the self-attention case (q/k/v share the same padding mask). + if attention_mask is not None: attention_mask = attention_mask.to(torch.bool) - seqlens = attention_mask.sum(dim=-1, dtype=torch.int32) - indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() - max_seqlen = int(seqlens.max().item()) - cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0)) - - query_states = query.reshape(batch_size * sequence_length, attn.heads, head_dim)[indices] - key_states = key.reshape(batch_size * sequence_length, kv_heads, head_dim)[indices] - value_states = value.reshape(batch_size * sequence_length, kv_heads, head_dim)[indices] - - if kv_heads < attn.heads: - key_states = repeat(key_states, "l h c -> l (h k) c", k=attn.heads // kv_heads) - value_states = repeat(value_states, "l h c -> l (h k) c", k=attn.heads // kv_heads) - - attn_output_unpad = flash_attn_varlen_func( - query_states, - key_states, - value_states, - cu_seqlens_q=cu_seqlens, - cu_seqlens_k=cu_seqlens, - max_seqlen_q=max_seqlen, - max_seqlen_k=max_seqlen, - dropout_p=0.0, - causal=False, - softmax_scale=softmax_scale, - ) - out = torch.zeros( - (batch_size * sequence_length, attn.heads, head_dim), - device=attn_output_unpad.device, - dtype=attn_output_unpad.dtype, - ) - out[indices] = attn_output_unpad - hidden_states = out.view(batch_size, sequence_length, attn.heads, head_dim).flatten(-2) - hidden_states = hidden_states.type_as(query) - else: - # PyTorch SDPA path. - attn_mask = None - if attention_mask is not None: - attention_mask = attention_mask.to(torch.bool) - attn_mask = attention_mask.view(batch_size, 1, 1, -1) - - query = query.transpose(1, 2) # [B, H, S, D] - key = key.transpose(1, 2) # [B, H_kv, S, D] - value = value.transpose(1, 2) - - if kv_heads < attn.heads: - key = key.repeat_interleave(attn.heads // kv_heads, dim=1) - value = value.repeat_interleave(attn.heads // kv_heads, dim=1) - - hidden_states = F.scaled_dot_product_attention( - query, - key, - value, - attn_mask=attn_mask, - dropout_p=0.0, - is_causal=False, - scale=softmax_scale, - ) + attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None + hidden_states = attn.omni_attn(query, key, value, attn_metadata) - if attention_mask is not None: - # Keep padding tokens consistent with the flash-varlen path (zero output). - hidden_states = hidden_states * attention_mask[:, None, :, None] + if attention_mask is not None: + # Padded rows carry nothing downstream; keep them at zero as the + # previous varlen path did, since neither backend does it for us. + hidden_states = hidden_states * attention_mask[:, :, None, None] - hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) - hidden_states = hidden_states.type_as(query) + hidden_states = hidden_states.reshape(batch_size, sequence_length, attn.heads * head_dim).type_as(query) - # Apply output projection hidden_states = attn.to_out[0](hidden_states) hidden_states = attn.to_out[1](hidden_states) @@ -441,6 +371,18 @@ def __init__( # 显式使用 transformers 的 Qwen2RMSNorm，避免依赖 diffusers 内部创建的 `RMSNorm` 再做递归替换。 self.attn.norm_q = Qwen2RMSNorm(self.head_dim, eps=1e-5) self.attn.norm_k = Qwen2RMSNorm(self.head_dim, eps=1e-5) + # The kernel itself runs through the shared Omni attention layer: backend + # selection, padding-mask handling and native grouped-query heads live + # there. It owns no parameters, so checkpoint keys are unchanged. The + # diffusers Attention above is kept for its projections and QK norms. + self.attn.omni_attn = OmniAttention( + num_heads=num_attention_heads, + head_size=self.head_dim, + causal=False, + softmax_scale=self.attn.scale, + num_kv_heads=num_kv_heads, + allow_fp32_fallback=True, + ) # Initialize feed-forward network self.feed_forward = LuminaFeedForward( From 4a7a8454ba5649af43a7b9d44b92733cc488c604 Mon Sep 17 00:00:00 2001 From: Guangjian Dong <163994576+Hiro208@users.noreply.github.com> Date: Sat, 12 Sep 2026 12:00:15 +0800 Subject: [PATCH 84/93] [Bugfix] Give model CLI flags typed owners in the Omni config (#7390) Signed-off-by: Guangjian Co-authored-by: amy-why-3459 --- tests/config/test_omni_config.py | 50 ++++++++++++++++++++++++++++++++ vllm_omni/config/omni_config.py | 26 +++++++++++++++++ 2 files changed, 76 insertions(+) diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index fdb2cd4b970..a8a2056b688 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -271,6 +271,44 @@ def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy) +_MODEL_CLI_FLAGS = { + "served_model_name": "omni", + "allowed_local_media_path": "/data", + "allowed_media_domains": ["example.com"], + "max_logprobs": 7, + "logprobs_mode": "processed_logprobs", + "mm_processor_kwargs": {"max_slice_nums": 1}, + "mm_processor_cache_type": "shm", + "hf_token": "hf_test", + "hf_config_path": "/models/config", + "generation_config": "vllm", + "override_generation_config": {"temperature": 0.5}, + "enable_prompt_embeds": True, +} + + +@pytest.mark.parametrize("model_type", ["minicpmo_4_5", "qwen3_tts"]) +@pytest.mark.parametrize(("field", "value"), list(_MODEL_CLI_FLAGS.items()), ids=list(_MODEL_CLI_FLAGS)) +def test_from_pipeline_config_owns_model_cli_fields(model_type, field, value): + config = _from_pipeline_key(model_type, cli_overrides={field: value}) + + for stage in config.stage_configs: + assert getattr(stage.model_config, field) == value + + +def test_model_cli_fields_reach_typed_engine_args(monkeypatch): + from vllm_omni.engine import stage_init_utils + from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config + + # Worker discovery requires hardware support; this test covers config transport. + monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda engine_args: None) + config = _from_pipeline_key("minicpmo_4_5", cli_overrides=dict(_MODEL_CLI_FLAGS)) + + engine_args = build_engine_args_dict_from_omni_stage_config(config.stage_by_id(0), model="test-model") + + assert {field: engine_args.get(field) for field in _MODEL_CLI_FLAGS} == _MODEL_CLI_FLAGS + + @pytest.mark.parametrize("stage_id", [0, 2], ids=["ar", "generation"]) def test_llm_additional_config_roundtrip_and_isolation(stage_id, monkeypatch): from vllm_omni.engine import stage_init_utils @@ -723,6 +761,18 @@ def test_sub_config_fields_match_structured_scopes(): "model_subdir", "tokenizer_subdir", "requires_full_payload_input", + "served_model_name", + "allowed_local_media_path", + "allowed_media_domains", + "max_logprobs", + "logprobs_mode", + "mm_processor_kwargs", + "mm_processor_cache_type", + "hf_token", + "hf_config_path", + "generation_config", + "override_generation_config", + "enable_prompt_embeds", } vllm_load_fields = {f.name for f in fields(VllmLoadConfig)} assert issubclass(OmniStageLoadConfig, VllmLoadConfig) diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index f0752077bb4..b588870995f 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -168,6 +168,19 @@ class _ModelEngineOverrides(TypedDict, total=False): enable_broadcast_weight_load: bool num_weight_load_threads: int disable_autocast: bool + # Upstream ModelConfig inputs that users pass as global CLI flags. + served_model_name: str | list[str] + allowed_local_media_path: str + allowed_media_domains: list[str] + max_logprobs: int + logprobs_mode: str + mm_processor_kwargs: dict[str, Any] + mm_processor_cache_type: str + hf_token: bool | str + hf_config_path: str + generation_config: str + override_generation_config: dict[str, Any] + enable_prompt_embeds: bool class _LoadEngineOverrides(TypedDict, total=False): @@ -445,6 +458,19 @@ class OmniStageModelConfig(_TrackExplicitConfigFields): model_subdir: str | None = None tokenizer_subdir: str | None = None requires_full_payload_input: bool = False + # Upstream ModelConfig inputs that users pass as global CLI flags. + served_model_name: str | list[str] | None = None + allowed_local_media_path: str | None = None + allowed_media_domains: list[str] | None = None + max_logprobs: int | None = None + logprobs_mode: str | None = None + mm_processor_kwargs: dict[str, Any] | None = None + mm_processor_cache_type: str | None = None + hf_token: bool | str | None = None + hf_config_path: str | None = None + generation_config: str | None = None + override_generation_config: dict[str, Any] | None = None + enable_prompt_embeds: bool | None = None @_enforce_keyword_only_init From 68294066e3df6cdd0eba65da0ff05c712564d32b Mon Sep 17 00:00:00 2001 From: longguo <107740309+abinggo@users.noreply.github.com> Date: Sat, 12 Sep 2026 12:03:16 +0800 Subject: [PATCH 85/93] [Bugfix] Require a model for `vllm serve --omni` (fixes #4158) (#4167) Signed-off-by: abinggo <107740309+abinggo@users.noreply.github.com> --- tests/entrypoints/test_serve.py | 74 ++++++++++++++++++++++++++++++ vllm_omni/entrypoints/cli/serve.py | 44 ++++++++++++++++++ 2 files changed, 118 insertions(+) diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 083f96d9f4e..1f3b5817099 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -44,6 +44,80 @@ def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None: assert not explicit["async_chunk"] +def _parse_serve_args(argv: list[str]) -> TrackingNamespace: + """Parse a ``serve`` argv through the real Omni parser, returning the + TrackingNamespace (with ``explicit_keys``) that ``validate`` receives.""" + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + OmniServeCommand().subparser_init(subparsers) + return parser.parse_args(argv) + + +def test_omni_serve_requires_model_when_none_provided() -> None: + """Regression for https://github.com/vllm-project/vllm-omni/issues/4158: + ``vllm serve --omni`` with no model must fail fast instead of silently + falling back to vLLM's default model and crashing in the diffusion worker. + """ + args = _parse_serve_args(["serve", "--omni"]) + cmd = OmniServeCommand() + with pytest.raises(ValueError, match="requires an explicit model"): + cmd.validate(args) + + +@pytest.mark.parametrize( + "argv", + [ + # A deploy YAML supplied *without* an explicit model must still trip + # the guard. It carries per-stage engine args only -- the + # checkpoint is always threaded in from ``args.model`` -- so its mere + # presence does not establish that a model was provided. Without an + # explicit model ``args.model`` stays at vLLM's default and the launch + # reproduces the issue-4158 diffusion-registry crash. + ["serve", "--omni", "--deploy-config", "deploy.yaml"], + # An empty/whitespace model must not satisfy the guard either -- this is + # the ``--model "$MODEL"``-with-unset-var footgun. Otherwise the empty + # value flows downstream and crashes with the same confusing error. + ["serve", "", "--omni"], # empty positional + ["serve", "--omni", "--model", ""], # empty --model + ["serve", "--omni", "--model", " "], # whitespace-only --model + ], +) +def test_omni_serve_rejects_missing_or_empty_model(argv: list[str]) -> None: + """Regression for review feedback on PR #4167: a deploy YAML is not a + model source, so ``--deploy-config`` alone must not bypass the + require-a-model guard; neither may an empty/whitespace model value.""" + args = _parse_serve_args(argv) + cmd = OmniServeCommand() + with pytest.raises(ValueError, match="requires an explicit model"): + cmd.validate(args) + + +@pytest.mark.parametrize( + "argv", + [ + ["serve", "fake-model", "--omni"], # positional model + ["serve", "--omni", "--model", "fake-model"], # --model flag + # An explicit model with a deploy YAML is the legitimate multi-stage + # path and must pass the guard. + ["serve", "fake-model", "--omni", "--deploy-config", "deploy.yaml"], + ], +) +def test_omni_serve_accepts_explicit_model(argv: list[str], mocker: MockerFixture) -> None: + """A model supplied positionally or via ``--model`` -- with or without a + deploy YAML -- must NOT trip the require-a-model guard. Downstream + validation is mocked so only the new guard is exercised.""" + mocker.patch( + "vllm_omni.diffusion.utils.hf_utils.is_diffusion_model", + return_value=False, + ) + mocker.patch("vllm_omni.entrypoints.cli.serve.validate_parsed_serve_args") + + args = _parse_serve_args(argv) + cmd = OmniServeCommand() + # Must not raise the "requires a model" error for any of these. + cmd.validate(args) + + def test_serve_parser_accepts_strategy_config() -> None: """``--strategy-config`` must parse onto the ``strategy_config`` dest and be forwarded as an explicit kwarg so the engine can overlay the strategy.""" diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index f9297aa0c20..0166cb4c1a9 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -148,6 +148,50 @@ def validate(self, args: argparse.Namespace) -> None: if args.stage_id is not None and (args.omni_master_address is None or args.omni_master_port is None): raise ValueError("--stage-id requires both --omni-master-address and --omni-master-port to be set") + # Require an explicit model under --omni. ``args.model`` always carries + # vLLM's ModelConfig default (``Qwen/Qwen3-0.6B``), so an omit is silent: + # the default text LLM is routed into the diffusion stage and startup + # crashes deep in the diffusion worker with a confusing registry error. + # Fail fast instead. + # + # The model must come from the CLI -- positionally + # (``vllm serve --omni``) or via ``--model``. A deploy YAML + # (``--deploy-config``) is NOT a model source: it carries per-stage + # engine args only, and the checkpoint is always threaded in from + # ``args.model`` (see + # ``load_and_resolve_stage_configs``, which takes ``model`` as its first + # argument). Treating its mere presence as "model provided" let + # ``vllm serve --omni --deploy-config `` slip through with the + # default model and reproduce the very crash this guard prevents. See + # https://github.com/vllm-project/vllm-omni/issues/4158. + if getattr(args, "omni", False): + explicit_keys = getattr(args, "explicit_keys", None) or frozenset() + # Resolve the model the user actually supplied on the CLI. A + # positional ``model_tag`` takes precedence (``cmd`` later copies it + # onto ``args.model``); otherwise ``--model`` counts only when it was + # explicitly passed. An empty/whitespace value (e.g. ``--model + # "$MODEL"`` with ``MODEL`` unset) is treated as "not provided" so it + # fails here with a clear message instead of the same confusing + # downstream crash. + model_tag = getattr(args, "model_tag", None) + if model_tag is not None: + explicit_model = model_tag + elif "model" in explicit_keys: + explicit_model = getattr(args, "model", None) + else: + explicit_model = None + model_provided = explicit_model is not None and str(explicit_model).strip() != "" + if not model_provided: + raise ValueError( + "`vllm serve --omni` requires an explicit model. Pass it " + "positionally (`vllm serve --omni`) or via `--model`. " + "`--deploy-config` carries per-stage engine args only and " + "does not supply a model; without an " + "explicit model, vLLM's default (Qwen/Qwen3-0.6B) is selected " + "and routed into the diffusion stage, which fails with a " + "confusing diffusion-registry error." + ) + # --omni-replica-address is only consulted in run_headless(); reject it # on the head so a misconfigured launch fails loudly instead of being # silently ignored. From 9ef9fc442a8a85c8e6b42a46412ceffaec7e0999 Mon Sep 17 00:00:00 2001 From: psv666 <150513104+psv666@users.noreply.github.com> Date: Sat, 12 Sep 2026 12:10:37 +0800 Subject: [PATCH 86/93] [Bugfix] Send a downstream terminal chunk when a parked stage ends (#6889) Signed-off-by: psv666 <2693925048@qq.com> --- .../test_chunk_transfer_adapter.py | 350 ++++++++++++++++++ vllm_omni/core/sched/omni_scheduler_mixin.py | 62 ++++ .../chunk_transfer_adapter.py | 120 +++++- 3 files changed, 519 insertions(+), 13 deletions(-) diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 40d29e96919..a80ded5f056 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -2526,6 +2526,356 @@ def _super_finish(_self, request_ids, finished_status): assert order == ["adapter", "super"] +def _parked_sender_scheduler(adapter, session): + """A minimal OmniARScheduler holding one parked streaming session.""" + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler.chunk_transfer_adapter = adapter + scheduler.input_coordinator = None + scheduler.requests = {session.request_id: session} + scheduler.running = [] + scheduler.waiting = DummyWaitingQueue() + scheduler.skipped_waiting = DummyWaitingQueue([session]) + scheduler.num_waiting_for_streaming_input = 1 + scheduler.finished_req_ids_dict = {} + + def _free_request(self, request, delay_free_blocks=False): + del delay_free_blocks + self.requests.pop(request.request_id, None) + return None, None + + scheduler._free_request = MethodType(_free_request, scheduler) + return scheduler + + +@pytest.mark.parametrize("send_first_chunk", [False, True]) +def test_parked_streaming_final_update_emits_downstream_terminal(build_adapter, send_first_chunk): + """#6670: a final update on a parked sender must terminate the receiver. + + Upstream turns "final update while parked in WAITING_FOR_STREAMING_REQ" + into a silent local abort. The downstream stage only learns a stream ended + from a chunk carrying ``finished``, so the abort used to leave it in + WAITING_FOR_CHUNK until the 600s input deadline. + """ + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + session = _req( + "req-parked-final", + RequestStatus.WAITING_FOR_STREAMING_REQ, + external_req_id="ext-parked-final", + ) + session.resumable = True + session._omni_segment_generation = 0 + # A segment stop queues its boundary before the final update arrives. + # The downstream request is already prewarmed even if the sender has not + # dequeued this first task, so it still needs a terminal in both cases. + adapter.save_async(None, session, is_segment_finished=True) + assert session.external_req_id not in adapter.put_req_chunk + if send_first_chunk: + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + connector.put.reset_mock() + + scheduler = _parked_sender_scheduler(adapter, session) + + final_update = _req(session.request_id, RequestStatus.WAITING) + final_update.resumable = False + OmniARScheduler.add_request(scheduler, final_update) + + # Local teardown happened... + assert session.request_id not in scheduler.requests + # ...but only after a terminal chunk was queued for the next stage. + assert session.status == RequestStatus.FINISHED_STOPPED + assert any(task["is_finished"] for task in adapter._pending_save_reqs) + while adapter._pending_save_reqs: + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + + connector.put.assert_called_once() + assert connector.put.call_args.kwargs["put_key"] == f"ext-parked-final_1_{int(send_first_chunk)}" + payload = connector.put.call_args.kwargs["data"] + assert bool(payload.meta.finished.item()) is True + assert bool(payload.meta.is_segment_finished.item()) is False + + +def test_queued_terminal_chunk_survives_sender_cleanup(build_adapter): + """#6670: ``cleanup_sender`` must not discard an already-queued terminal.""" + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + request = _req("req-terminal", RequestStatus.FINISHED_STOPPED, external_req_id="ext-terminal") + request.resumable = False + + adapter.save_async(None, request, is_segment_finished=False) + assert len(adapter._pending_save_reqs) == 1 + + # The scheduler tears the request down before the save_loop drains. + adapter.cleanup_sender("ext-terminal") + + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + connector.put.assert_called_once() + assert bool(connector.put.call_args.kwargs["data"].meta.finished.item()) is True + # The terminal put reclaims the generation it kept alive. + assert "ext-terminal" not in adapter._sender_tokens + + +def test_non_terminal_chunk_is_still_dropped_by_sender_cleanup(build_adapter): + """The #6670 fence is terminal-only: ordinary queued chunks still drop.""" + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + request = _req("req-mid", RequestStatus.RUNNING, external_req_id="ext-mid") + request.resumable = True + + adapter.save_async(None, request, is_segment_finished=False) + adapter.cleanup_sender("ext-mid") + + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + connector.put.assert_not_called() + + +def test_final_update_on_non_sender_stage_keeps_upstream_abort(build_adapter): + """A final stage never put a chunk, so it must keep the upstream path.""" + adapter, _connector = build_adapter(stage_id=2, model_mode="generation") + session = _req( + "req-final-stage", + RequestStatus.WAITING_FOR_STREAMING_REQ, + external_req_id="ext-final-stage", + ) + session.resumable = True + scheduler = _parked_sender_scheduler(adapter, session) + + final_update = _req(session.request_id, RequestStatus.WAITING) + final_update.resumable = False + + with patch.object(VLLMScheduler, "add_request") as base_add: + OmniARScheduler.add_request(scheduler, final_update) + + base_add.assert_called_once() + assert not adapter._pending_save_reqs + + +def _queue_terminal( + adapter, + external_req_id: str, + req_id: str = "req-terminal", + segment_generation: int = 0, +) -> None: + """Queue one request-terminal chunk on the adapter's save queue. + + ``segment_generation`` must match the adapter's dedup watermark, which a + preceding segment stop advances; otherwise ``save_async`` drops the chunk + as a late duplicate. + """ + request = _req(req_id, RequestStatus.FINISHED_STOPPED, external_req_id=external_req_id) + request.resumable = False + request._omni_segment_generation = segment_generation + adapter.save_async(None, request, is_segment_finished=False) + + +def _assert_terminal_put(connector) -> None: + connector.put.assert_called_once() + assert bool(connector.put.call_args.kwargs["data"].meta.finished.item()) is True + + +def test_terminal_chunk_survives_cleanup_after_dequeue(build_adapter): + """#6670: the fence must hold across the connector handoff, not just the queue. + + ``_finish_parked_streaming_session`` queues the terminal and then calls + ``finish_requests``, so ``cleanup_sender`` can land *after* the save_loop + has already dequeued the terminal. Releasing the fence at dequeue left + that interleaving cancelling the chunk before ``put`` -- the same lost + terminal, and the same downstream WAITING_FOR_CHUNK hang. + """ + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + _queue_terminal(adapter, "ext-late-cleanup") + task = adapter._pending_save_reqs.popleft() + + real_send = adapter._send_single_request_for_generation + + def send_with_cleanup_in_flight(inner_task, sender_token=None): + # The scheduler thread reaches cleanup_sender here: the save_loop owns + # the terminal but has not handed it to the connector yet. + adapter.cleanup_sender("ext-late-cleanup") + return real_send(inner_task, sender_token) + + adapter._send_single_request_for_generation = send_with_cleanup_in_flight + adapter._send_single_request(task) + + _assert_terminal_put(connector) + # The generation the fence kept alive is reclaimed, not leaked. + assert "ext-late-cleanup" not in adapter._sender_tokens + assert "ext-late-cleanup" not in adapter.put_req_chunk + + +def test_terminal_chunk_survives_cleanup_racing_an_inflight_sibling(build_adapter): + """#6670: an in-flight ordinary chunk must not reclaim a fenced generation. + + ``_send_single_request``'s ``finally`` also tears a cancelled generation + down. If the stop and the abort both land while an earlier chunk is on the + wire, that teardown would unregister the token and the queued terminal + would then fail the identity check. + """ + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + running = _req("req-sibling", RequestStatus.RUNNING, external_req_id="ext-sibling") + running.resumable = True + adapter.save_async(None, running, is_segment_finished=True) + ordinary_task = adapter._pending_save_reqs.popleft() + + real_send = adapter._send_single_request_for_generation + + def send_then_stop_and_abort(inner_task, sender_token=None): + result = real_send(inner_task, sender_token) + # Terminal stop and abort both land while this chunk is on the wire. + # The segment stop above already armed the next segment's watermark. + _queue_terminal(adapter, "ext-sibling", req_id="req-sibling", segment_generation=1) + adapter.cleanup_sender("ext-sibling") + return result + + adapter._send_single_request_for_generation = send_then_stop_and_abort + adapter._send_single_request(ordinary_task) + adapter._send_single_request_for_generation = real_send + + assert len(adapter._pending_save_reqs) == 1 + connector.put.reset_mock() + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + + _assert_terminal_put(connector) + assert "ext-sibling" not in adapter._sender_tokens + + +@pytest.mark.parametrize("finish_inside_put", [False, True]) +def test_successful_put_keeps_terminal_chunk_sequence(build_adapter, finish_inside_put): + """A cleanup during put must not reuse a key already consumed downstream.""" + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + receiver, receiver_connector = build_adapter(stage_id=2, model_mode="generation") + session = _req("req-sequence", RequestStatus.WAITING_FOR_STREAMING_REQ, external_req_id="ext-sequence") + session.resumable = True + session._omni_segment_generation = 0 + scheduler = _parked_sender_scheduler(adapter, session) + final_update = _req(session.request_id, RequestStatus.WAITING) + final_update.resumable = False + downstream = _req("req-downstream", RequestStatus.WAITING, external_req_id=session.external_req_id) + receiver.request_ids_mapping[downstream.request_id] = session.external_req_id + payloads = {} + + def receive_next(): + entry = _dequeue_load_entry(receiver, downstream) + return receiver._poll_single_request(entry) + + def put_then_finish(**kwargs): + payload = kwargs["data"] + payloads[kwargs["put_key"]] = { + "meta": { + "finished": bool(payload.meta.finished.item()), + "is_segment_finished": bool(payload.meta.is_segment_finished.item()), + } + } + if not payload.meta.finished.item(): + # Consume the ordinary boundary while put is still in flight. + assert receive_next() + if finish_inside_put: + scheduler.add_request(final_update) + return True, 1, {} + + def get_by_key(_from_stage, _to_stage, key): + payload = payloads.pop(key, None) + return (payload, 1) if payload is not None else None + + connector.put.side_effect = put_then_finish + receiver_connector.get.side_effect = get_by_key + adapter.save_async(None, session, is_segment_finished=True) + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + if not finish_inside_put: + scheduler.add_request(final_update) + while adapter._pending_save_reqs: + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + + assert [call.kwargs["put_key"] for call in connector.put.call_args_list] == ["ext-sequence_1_0", "ext-sequence_1_1"] + assert receive_next() + assert downstream.request_id in receiver.upstream_exhausted_requests + assert not payloads + assert session.external_req_id not in adapter._sender_tokens + + +def test_late_successful_put_cannot_advance_replacement_generation(build_adapter): + """The post-put guard must retain identity checks while relaxing cancellation.""" + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + request = _req("req-old", RequestStatus.RUNNING, external_req_id="ext-replaced") + request.resumable = True + adapter.save_async(None, request, is_segment_finished=True) + task = adapter._pending_save_reqs.popleft() + old_token = task["sender_token"] + replacement_token = type(old_token)() + + def replace_during_put(**kwargs): + # Model a completion that no longer owns its external id. The old + # terminal fence must never authorize writes to the successor's state. + with adapter._sender_state_lock: + old_token.cancelled = True + old_token.terminal_pending = True + adapter._sender_tokens[request.external_req_id] = replacement_token + adapter.put_req_chunk[request.external_req_id] = 4 + adapter.ramp_chunk_count[request.external_req_id] = 2 + return True, 1, {} + + connector.put.side_effect = replace_during_put + adapter._send_single_request(task) + assert adapter._sender_tokens[request.external_req_id] is replacement_token + assert adapter.put_req_chunk[request.external_req_id] == 4 + assert adapter.ramp_chunk_count[request.external_req_id] == 2 + + +def test_new_request_cannot_join_a_generation_with_a_queued_terminal(build_adapter): + """#6670: the fence must not readmit a request reusing the external id. + + ``save_async`` rejects chunks whose generation is ``cancelled``; that is + what stops a new request from inheriting a retiring generation's token and + chunk counters. Deferring the reclaim must not also defer that retirement. + """ + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + _queue_terminal(adapter, "ext-shared", req_id="req-old") + adapter.cleanup_sender("ext-shared") + + newcomer = _req("req-new", RequestStatus.RUNNING, external_req_id="ext-shared") + newcomer.resumable = True + adapter.save_async(None, newcomer, is_segment_finished=True) + + # Rejected loudly, so the scheduler fails the newcomer now instead of its + # chunks being silently discarded at send time. + assert len(adapter._pending_save_reqs) == 1 + assert adapter.collect_failed_send_request_ids() == {"req-new": "previous sender generation is still draining"} + + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + _assert_terminal_put(connector) + assert "ext-shared" not in adapter._sender_tokens + + +def test_failed_terminal_put_runs_the_deferred_cleanup(build_adapter): + """#6670: a silent connector drop must not strand the deferred cleanup. + + ``connector.put`` returning False skips the success path that would have + called ``cleanup``. The reclaim ``cleanup_sender`` handed over is owed on + that path too, or the whole per-request sender state leaks for good. + """ + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + _queue_terminal(adapter, "ext-put-failed") + adapter.cleanup_sender("ext-put-failed") + connector.put.return_value = (False, 0, {}) + + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + + assert "ext-put-failed" not in adapter._sender_tokens + assert "ext-put-failed" not in adapter.put_req_chunk + + +def test_raising_terminal_put_runs_the_deferred_cleanup(build_adapter): + """#6670: the same reclaim is owed when the connector raises.""" + adapter, connector = build_adapter(stage_id=1, model_mode="ar") + _queue_terminal(adapter, "ext-put-raised") + adapter.cleanup_sender("ext-put-raised") + connector.put.side_effect = RuntimeError("connector down") + + # ``save_loop`` catches this and records a send failure. + with pytest.raises(RuntimeError): + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + + assert "ext-put-raised" not in adapter._sender_tokens + assert "ext-put-raised" not in adapter.put_req_chunk + + def test_wire_round_trip_struct_to_dict_contract(): """Pin the wire contract: encoding ``OmniPayloadStruct`` and decoding it yields a dict equivalent to ``to_dict(struct)``. diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 9123b794d7c..f33359835ab 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -688,6 +688,68 @@ def _attach_scheduler_stats( engine_core_outputs[0] = output = EngineCoreOutputs() output.scheduler_stats = stats + def add_request(self, request: Request) -> None: + """Route a terminal streaming update that lands on a parked session. + + Upstream ``Scheduler.add_request`` turns "final update (``resumable`` + is False) arrives while the session sits in + ``WAITING_FOR_STREAMING_REQ``" into a local + ``finish_requests(FINISHED_ABORTED)``. That is fine for a stage whose + output the orchestrator forwards, but an async-chunk *sender* also + owes its downstream stage a terminal chunk: the receiver learns that a + stream ended only from a connector payload carrying ``finished``, and + an abort emits none. The downstream stage then parks in + ``WAITING_FOR_CHUNK`` until ``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S`` (600s by + default), long after the client's own deadline -- the realtime + async-chunk hang in vllm-project/vllm-omni#6670. + + Whether this races is timing, not configuration: the stage parks only + when its own stop beats both the upstream terminal chunk and the + orchestrator's terminal update, a window of tens of milliseconds at + the end of the last segment. Emit the terminal chunk here so the + downstream stage terminates on the payload path either way. + """ + existing = self.requests.get(request.request_id) + if existing is not None and existing.status == RequestStatus.WAITING_FOR_STREAMING_REQ: + adapter = None if getattr(request, "resumable", False) else self._adapter_owing_terminal(existing) + if adapter is not None: + self._finish_parked_streaming_session(existing, adapter) + return + super().add_request(request) + + def _adapter_owing_terminal(self, request: Request) -> OmniChunkTransferAdapter | None: + """The adapter that owes *request*'s downstream stage a terminal chunk. + + Sender ownership starts when the first chunk is queued, before the + background thread initializes its chunk counter. A prewarmed receiver + needs a terminal even if that first chunk has not been sent yet. + Final stages never queue sends, so they never match. + """ + adapter = getattr(self, "chunk_transfer_adapter", None) + if adapter is None: + return None + external_req_id = getattr(request, "external_req_id", None) or request.request_id + return adapter if adapter.has_active_sender(external_req_id) else None + + def _finish_parked_streaming_session(self, request: Request, adapter: OmniChunkTransferAdapter) -> None: + """End a parked async-chunk session with a downstream terminal chunk.""" + request.resumable = False + # A segment stop already advanced the adapter's dedup watermark to + # ``generation + 1`` for the segment that never arrived. Claim that + # generation, exactly as ``_update_request_as_session`` would have, + # or ``save_async`` drops this chunk as a late duplicate. + request._omni_segment_generation = int(getattr(request, "_omni_segment_generation", 0) or 0) + 1 + # ``force_request_finished``: the terminal must be enqueued *before* + # ``finish_requests``, which skips requests that already report + # ``is_finished()``. + adapter.save_async( + None, + request, + is_segment_finished=False, + force_request_finished=True, + ) + self.finish_requests(request.request_id, RequestStatus.FINISHED_STOPPED) + def finish_requests( self, request_ids: str | Iterable[str] | None, 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 d71d55bf15e..16336393178 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 @@ -32,6 +32,16 @@ class _SenderGeneration: def __init__(self) -> None: self.in_flight = False self.cancelled = False + # A request-terminal chunk is queued but not yet on the wire. The + # downstream stage has no other way to learn the stream ended -- it + # polls the connector and parks in WAITING_FOR_CHUNK until a chunk + # carrying ``finished`` arrives -- so cleanup must not discard it. + # See vllm-project/vllm-omni#6670. + self.terminal_pending = False + # ``cleanup_sender`` asked to reclaim this generation while the fence + # was up. The terminal task owes that cleanup once it is done with the + # connector -- on every exit path, including a failed or raising put. + self.cleanup_deferred = False class _LoadEntry: @@ -321,6 +331,17 @@ def _accepts_new_token_ids(self, processor: Callable[..., Any]) -> bool: self._processor_accepts_step_tokens[processor] = accepts return accepts + def has_active_sender(self, external_req_id: str) -> bool: + """Whether this request owns a sender, including work still queued. + + ``save_async`` registers the generation before the background thread + can send its first chunk. Chunk counters are initialized later and + cannot determine whether a prewarmed receiver needs a terminal. + """ + with self._sender_state_lock: + sender_token = self._sender_tokens.get(external_req_id) + return sender_token is not None and not sender_token.cancelled + def save_async( self, multimodal_output: dict[str, Any] | None = None, @@ -329,6 +350,7 @@ def save_async( new_token_ids: Iterable[int] | None = None, confirmed_num_computed_tokens: int | None = None, segment_generation: int | None = None, + force_request_finished: bool = False, ): """Build and enqueue one chunk for asynchronous sending. @@ -352,8 +374,14 @@ def save_async( before a streaming transition can mutate ``request`` segment_generation: generation captured before a resumable stop can apply a queued update to the mutable request + force_request_finished: send this chunk as the request terminal + even though ``request`` is not in a finished status yet. The + scheduler needs this when it ends a *parked* streaming session: + the terminal has to be enqueued before ``finish_requests``, + because ``finish_requests`` skips requests that already report + ``is_finished()``. See ``_finish_parked_streaming_session``. """ - is_finished = request.is_finished() and not request.resumable + is_finished = force_request_finished or (request.is_finished() and not request.resumable) if not hasattr(self, "_segment_generation"): self._segment_generation = defaultdict(int) external_req_id = request.external_req_id @@ -414,6 +442,10 @@ def save_async( self._sender_tokens[external_req_id] = sender_token task["sender_token"] = sender_token self._pending_save_reqs.append(task) + if is_finished: + # Fence this generation against cleanup until the terminal + # chunk has been handed to the connector (#6670). + sender_token.terminal_pending = True if is_segment_finished: # The queued FIFO item now owns the old segment. Start the next # segment's deduplication watermark before the worker sends it. @@ -624,32 +656,71 @@ def _send_single_request(self, task: dict): if sender_token is None: self._send_single_request_for_generation(task) return + is_terminal_task = bool(task.get("is_finished")) with self._sender_state_lock: is_current = self._sender_tokens.get(external_req_id) is sender_token - if is_current and not sender_token.cancelled: + # ``cancelled`` does not veto a terminal chunk. It is the + # downstream stage's only end-of-stream signal, so dropping it + # leaves that stage parked in WAITING_FOR_CHUNK until the input + # deadline -- the hang this fence exists to prevent (#6670). + # Losing ownership of the external id still vetoes: the state the + # chunk describes then belongs to a different generation. + if is_current and (is_terminal_task or not sender_token.cancelled): sender_token.in_flight = True else: is_current = False if not is_current: logger.debug("Discarding stale queued chunk for aborted request %s", external_req_id) + if is_terminal_task: + self._release_terminal_fence(external_req_id, sender_token) return try: self._send_single_request_for_generation(task, sender_token) finally: + if is_terminal_task: + # Drop the fence and run the cleanup it deferred, whether the + # put succeeded, reported failure, or raised. Until this runs, + # the block below must not reclaim the generation either. + self._release_terminal_fence(external_req_id, sender_token) with self._sender_state_lock: if self._sender_tokens.get(external_req_id) is sender_token: sender_token.in_flight = False - if sender_token.cancelled: + if sender_token.cancelled and not sender_token.terminal_pending: self._sender_tokens.pop(external_req_id, None) self._clear_sender_state_locked(external_req_id) + def _release_terminal_fence(self, external_req_id: str, sender_token: _SenderGeneration) -> None: + """Drop the terminal fence and run the cleanup it deferred. + + Called on every exit path of a terminal task. ``cleanup_sender`` + cannot reclaim a generation whose terminal chunk is still queued, so + it records the request on the token instead; without this the only + path that ever replayed it was a *successful* put, and a connector + failure leaked the whole per-request sender state (#6670). + """ + with self._sender_state_lock: + sender_token.terminal_pending = False + deferred = sender_token.cleanup_deferred + sender_token.cleanup_deferred = False + # A generation that no longer owns the external id was already + # reclaimed by whoever replaced it; cleaning now would clear the + # successor's state. + still_current = self._sender_tokens.get(external_req_id) is sender_token + if deferred and still_current: + self.cleanup_sender(external_req_id) + def _sender_generation_is_active( self, external_req_id: str, sender_token: _SenderGeneration, + *, + is_terminal: bool = False, ) -> bool: with self._sender_state_lock: - return self._sender_tokens.get(external_req_id) is sender_token and not sender_token.cancelled + if self._sender_tokens.get(external_req_id) is not sender_token: + return False + # Terminal chunks outlive cancellation -- see _send_single_request. + return is_terminal or not sender_token.cancelled def _send_single_request_for_generation( self, @@ -712,7 +783,9 @@ def _send_single_request_for_generation( if payload_data.meta.is_segment_finished is None: payload_data.meta.is_segment_finished = torch.tensor(is_segment_finished, dtype=torch.bool) - if sender_token is not None and not self._sender_generation_is_active(external_req_id, sender_token): + if sender_token is not None and not self._sender_generation_is_active( + external_req_id, sender_token, is_terminal=is_finished + ): logger.debug("Skipping cancelled chunk for request %s before connector put", external_req_id) return @@ -723,16 +796,25 @@ def _send_single_request_for_generation( data=payload_data, ) - if sender_token is not None and not self._sender_generation_is_active(external_req_id, sender_token): - # cleanup_sender() may cancel this generation while put() is - # blocked. Do not let that stale completion recreate state for a - # request whose cleanup is already in progress. - logger.debug("Ignoring completed put for cancelled request %s", external_req_id) - return + with self._sender_state_lock: + if sender_token is not None: + still_current = self._sender_tokens.get(external_req_id) is sender_token + retiring_without_terminal = sender_token.cancelled and not ( + is_finished or sender_token.terminal_pending + ) + if not still_current or retiring_without_terminal: + logger.debug("Ignoring completed put for retired request %s", external_req_id) + return + if success: + # A receiver may already have consumed this key while put was + # in flight. Even if cleanup retired this generation, its + # pending terminal must use the next key. Check ownership and + # advance counters under the same lock so a stale completion + # cannot modify a successor's state. + self.put_req_chunk[external_req_id] += 1 + self.ramp_chunk_count[external_req_id] += 1 if success: - self.put_req_chunk[external_req_id] += 1 - self.ramp_chunk_count[external_req_id] += 1 logger.debug(f"[Stage-{stage_id}] Sent {connector_put_key}") # Sender uses struct attr access here; the receive path in # `_load_one_request` / `_update_request_payload` reads dict keys. @@ -862,6 +944,18 @@ def cleanup_sender(self, external_req_id: str) -> None: if sender_token is None: self._clear_sender_state_locked(external_req_id) return + if sender_token.terminal_pending: + # A request-terminal chunk is still queued. Reclaiming the + # state now would drop it, and the downstream stage would sit + # in WAITING_FOR_CHUNK until VLLM_OMNI_INPUT_WAIT_TIMEOUT_S + # with no way to know the stream ended (#6670). Retire the + # generation anyway -- ``save_async`` rejects new chunks on + # ``cancelled``, which is what keeps a request reusing this + # external id from joining a generation that is going away -- + # and hand the reclaim to the terminal task. + sender_token.cancelled = True + sender_token.cleanup_deferred = True + return sender_token.cancelled = True if sender_token.in_flight: # The sender's finally block is the first point where no From b3ba67d71ea3466f1af76661275fe63a9f990f09 Mon Sep 17 00:00:00 2001 From: Weiming Liao Date: Sat, 12 Sep 2026 12:14:42 +0800 Subject: [PATCH 87/93] [NPU] upgrade to v0.29.0 (#7433) Signed-off-by: Weiming Liao --- docker/Dockerfile.npu | 2 +- docker/Dockerfile.npu.ci | 2 +- .../installation/npu/npu.inc.md | 22 +++++++++---------- .../npu/worker/npu_ar_model_runner.py | 8 ++++++- .../npu/worker/npu_generation_model_runner.py | 21 +++++++++++++----- .../platforms/npu/worker/npu_model_runner.py | 21 +++++++++++++----- 6 files changed, 50 insertions(+), 26 deletions(-) diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index 0d91c0b3273..f83cc512a88 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -1,5 +1,5 @@ ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend -ARG VLLM_ASCEND_VERSION=v0.28.0 +ARG VLLM_ASCEND_VERSION=v0.29.0 ARG VLLM_ASCEND_DEVICE_SUFFIX= FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_VERSION}${VLLM_ASCEND_DEVICE_SUFFIX} diff --git a/docker/Dockerfile.npu.ci b/docker/Dockerfile.npu.ci index 0a150fdd814..b65b803d6bc 100644 --- a/docker/Dockerfile.npu.ci +++ b/docker/Dockerfile.npu.ci @@ -1,5 +1,5 @@ ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend -ARG VLLM_ASCEND_VERSION=v0.28.0 +ARG VLLM_ASCEND_VERSION=v0.29.0 ARG VLLM_ASCEND_DEVICE_SUFFIX= FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_VERSION}${VLLM_ASCEND_DEVICE_SUFFIX} diff --git a/docs/getting_started/installation/npu/npu.inc.md b/docs/getting_started/installation/npu/npu.inc.md index 9ff477522d4..21127caf2ac 100644 --- a/docs/getting_started/installation/npu/npu.inc.md +++ b/docs/getting_started/installation/npu/npu.inc.md @@ -10,25 +10,25 @@ The recommended way to use vLLM-Omni on NPU is through the vLLM-Ascend pre-built === "A2" ```bash - export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.28.0 + export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.29.0 ``` === "A3" ```bash - export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.28.0-a3 + export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.29.0-a3 ``` === "A5" ```bash - export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.28.0-a5 + export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.29.0-a5 ``` === "310P (Experimental)" ```bash - export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.28.0-310p + export IMAGE=quay.io/atlas-ci/vllm-ascend:v0.29.0-310p ``` ```bash @@ -53,7 +53,7 @@ docker run --rm \ # Inside the container, install vLLM-Omni from source cd /vllm-workspace -git clone -b v0.28.0 https://github.com/vllm-project/vllm-omni.git +git clone -b v0.29.0rc1 https://github.com/vllm-project/vllm-omni.git cd vllm-omni pip install -v -e . --no-build-isolation # or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e . @@ -71,10 +71,10 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886 # --8<-- [start:installation-main] -You can also install vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. Use the aligned vLLM-Ascend v0.28.0 image described above, and check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest vLLM-Omni commit on NPU. +You can also install vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. Use the aligned vLLM-Ascend v0.29.0 image described above, and check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest vLLM-Omni commit on NPU. ```bash -# Inside the vLLM-Ascend v0.28.0 container, install vLLM-Omni from main +# Inside the vLLM-Ascend v0.29.0 container, install vLLM-Omni from main cd /vllm-workspace git clone https://github.com/vllm-project/vllm-omni.git cd vllm-omni @@ -103,7 +103,7 @@ Here's an example deployment command that has been verified on 4 x NPUs: === "A2" ```bash - export IMAGE=quay.io/ascend/vllm-omni:v0.28.0 + export IMAGE=quay.io/ascend/vllm-omni:v0.29.0 docker run --rm \ --name vllm-omni-a2 \ --shm-size=64g \ @@ -127,7 +127,7 @@ Here's an example deployment command that has been verified on 4 x NPUs: === "A3" ```bash - export IMAGE=quay.io/ascend/vllm-omni:v0.28.0-a3 + export IMAGE=quay.io/ascend/vllm-omni:v0.29.0-a3 docker run --rm \ --name vllm-omni-a3 \ --shm-size=64g \ @@ -151,7 +151,7 @@ Here's an example deployment command that has been verified on 4 x NPUs: === "A5" ```bash - export IMAGE=quay.io/ascend/vllm-omni:v0.28.0-a5 + export IMAGE=quay.io/ascend/vllm-omni:v0.29.0-a5 docker run --rm \ --name vllm-omni-a5 \ --shm-size=64g \ @@ -175,7 +175,7 @@ Here's an example deployment command that has been verified on 4 x NPUs: === "310P (Experimental)" ```bash - export IMAGE=quay.io/ascend/vllm-omni:v0.28.0-310p + export IMAGE=quay.io/ascend/vllm-omni:v0.29.0-310p docker run --rm \ --name vllm-omni-310p \ --shm-size=64g \ diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index 734178e01ca..7bc6c21c61d 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -334,6 +334,8 @@ def _request_final_stage_id(self, req_id: str) -> int | None: if not isinstance(info, dict): return None val = info.get("omni_final_stage_id") + if val is None: + return None try: return int(val) except (TypeError, ValueError): @@ -512,7 +514,7 @@ def execute_model( # returns True. before returning early here we call # dummy run to ensure coordinate_batch_across_dp # is called into to avoid out of sync issues. - self._dummy_run(1) + self._dummy_run(1, skip_gdn_state_update=True) kv_ids = self.kv_extracted_req_ids self.kv_extracted_req_ids = None @@ -1035,6 +1037,8 @@ def propose_draft_token_ids(sampled_token_ids): propose_draft_token_ids(sampler_output.sampled_token_ids) ( + num_nans_in_logits, + num_nans_device, logprobs_lists, valid_sampled_token_ids, prompt_logprobs_dict, @@ -1278,6 +1282,7 @@ def _unwrap_lists(v): kv_connector_output=kv_connector_output, ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, cudagraph_stats=cudagraph_stats, + num_nans_in_logits=num_nans_in_logits, ) model_runner_output.kv_extracted_req_ids = kv_extracted_req_ids model_runner_output.routed_experts = routed_experts_lists @@ -1318,6 +1323,7 @@ def _unwrap_lists(v): invalid_req_indices=invalid_req_indices, async_output_copy_stream=self.async_output_copy_stream, vocab_size=self.input_batch.vocab_size, + num_nans=num_nans_device, ) self.input_batch.set_async_sampled_token_ids( async_output.sampled_token_ids_cpu, diff --git a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py index fc26882cd5c..59dfc88a3bb 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -166,7 +166,7 @@ def execute_model( # returns True. before returning early here we call # dummy run to ensure coordinate_batch_across_dp # is called into to avoid out of sync issues. - self._dummy_run(1) + self._dummy_run(1, skip_gdn_state_update=True) if not has_kv_transfer_group(): # Return empty ModelRunnerOutput if no work to do. return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT) @@ -636,6 +636,7 @@ def _dummy_run( profile_seq_lens: int | None = None, profile_cpp: bool = False, randomize_inputs: bool = False, + skip_gdn_state_update: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: # only support eager mode and piecewise graph now assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.is_valid_runtime_mode() @@ -768,7 +769,10 @@ def _dummy_run( self.query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens self.query_start_loc.copy_to_gpu() if self._has_gdn: - self.gdn_query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens + if skip_gdn_state_update: + self.gdn_query_start_loc.np.fill(0) + else: + self.gdn_query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens self.gdn_query_start_loc.copy_to_gpu() if not profile_cpp: @@ -786,6 +790,12 @@ def _dummy_run( # rows as well so device-side metadata does not see stale block ids. self.input_batch.block_table.commit_block_table(num_reqs_padded) + # Invalidate slots before backends copy metadata for dummy execution. + if not is_graph_capturing: + for kv_cache_gid in range(len(self.kv_cache_config.kv_cache_groups)): + blk_table = self.input_batch.block_table[kv_cache_gid] + blk_table.slot_mapping.gpu.fill_(-1) + pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL attn_metadata, _ = self._build_attention_metadata( num_tokens=num_tokens_unpadded, @@ -796,11 +806,10 @@ def _dummy_run( ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, for_cudagraph_capture=is_graph_capturing, num_scheduled_tokens_np=num_scheduled_tokens, + cudagraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=batch_desc, + skip_gdn_state_update=skip_gdn_state_update, ) - if not is_graph_capturing: - for kv_cache_gid in range(len(self.kv_cache_config.kv_cache_groups)): - blk_table = self.input_batch.block_table[kv_cache_gid] - blk_table.slot_mapping.gpu.fill_(-1) with self.maybe_dummy_run_with_lora( self.lora_config, diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index 8f03698080a..bd7c7dce40f 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_model_runner.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 functools import partial from typing import Any @@ -126,6 +126,7 @@ def _dummy_run( profile_seq_lens: int | None = None, profile_cpp: bool = False, randomize_inputs: bool = False, + skip_gdn_state_update: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: # only support eager mode and piecewise graph now assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.is_valid_runtime_mode() @@ -257,7 +258,10 @@ def _dummy_run( self.query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens self.query_start_loc.copy_to_gpu() if self._has_gdn: - self.gdn_query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens + if skip_gdn_state_update: + self.gdn_query_start_loc.np.fill(0) + else: + self.gdn_query_start_loc.np[1 : num_reqs_padded + 1] = cum_num_tokens self.gdn_query_start_loc.copy_to_gpu() if not profile_cpp: @@ -275,6 +279,12 @@ def _dummy_run( # rows as well so device-side metadata does not see stale block ids. self.input_batch.block_table.commit_block_table(num_reqs_padded) + # Invalidate slots before backends copy metadata for dummy execution. + if not is_graph_capturing: + for kv_cache_gid in range(len(self.kv_cache_config.kv_cache_groups)): + blk_table = self.input_batch.block_table[kv_cache_gid] + blk_table.slot_mapping.gpu.fill_(-1) + pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL attn_metadata, _ = self._build_attention_metadata( num_tokens=num_tokens_unpadded, @@ -285,11 +295,10 @@ def _dummy_run( ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, for_cudagraph_capture=is_graph_capturing, num_scheduled_tokens_np=num_scheduled_tokens, + cudagraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=batch_desc, + skip_gdn_state_update=skip_gdn_state_update, ) - if not is_graph_capturing: - for kv_cache_gid in range(len(self.kv_cache_config.kv_cache_groups)): - blk_table = self.input_batch.block_table[kv_cache_gid] - blk_table.slot_mapping.gpu.fill_(-1) with self.maybe_dummy_run_with_lora( self.lora_config, From bc0c9f4b45c45c59aa2f92471842e8c18ae403ca Mon Sep 17 00:00:00 2001 From: junpengw67-max Date: Sat, 12 Sep 2026 12:23:24 +0800 Subject: [PATCH 88/93] [Bugfix][Model][Lance] Support decoded video frames in video editing (#5128) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 吴俊鹏 <248679769+junpengw67-max@users.noreply.github.com> Co-authored-by: 吴俊鹏 <248679769+junpengw67-max@users.noreply.github.com> --- .../models/lance/test_pipeline_lance.py | 171 ++++++++++++++++++ .../diffusion/models/lance/pipeline_lance.py | 3 + vllm_omni/model_executor/models/utils.py | 51 +++++- 3 files changed, 224 insertions(+), 1 deletion(-) create mode 100644 tests/diffusion/models/lance/test_pipeline_lance.py diff --git a/tests/diffusion/models/lance/test_pipeline_lance.py b/tests/diffusion/models/lance/test_pipeline_lance.py new file mode 100644 index 00000000000..32ff9bd1062 --- /dev/null +++ b/tests/diffusion/models/lance/test_pipeline_lance.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import numpy as np +import pytest +from PIL import Image + +from vllm_omni.diffusion.models.lance.lance_transformer import LanceBagel +from vllm_omni.diffusion.models.lance.pipeline_lance import LancePipeline +from vllm_omni.model_executor.models.utils import normalize_decoded_video_frames + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class FramesWithMetadata(list[Image.Image]): + def __init__(self, frames: list[Image.Image], **metadata: object) -> None: + super().__init__(frames) + for name, value in metadata.items(): + setattr(self, name, value) + + +class StopAfterVideoPreprocessError(Exception): + pass + + +def test_normalize_decoded_video_frames_preserves_rgb_pixels_and_fps(): + red = Image.new("RGB", (3, 2), color=(255, 0, 0)) + rgba = Image.new("RGBA", (3, 2), color=(0, 255, 0, 64)) + rgba_bytes = rgba.tobytes() + + video, fps = normalize_decoded_video_frames(FramesWithMetadata([red, rgba], fps=8), default_fps=24.0) + + assert video.shape == (2, 2, 3, 3) + assert video.dtype == np.uint8 + assert video.flags.c_contiguous + np.testing.assert_array_equal(video[0, 0, 0], [255, 0, 0]) + np.testing.assert_array_equal(video[1, 0, 0], [0, 255, 0]) + assert fps == 8.0 + assert rgba.mode == "RGBA" + assert rgba.tobytes() == rgba_bytes + + +def test_normalize_decoded_video_frames_uses_frame_rate_when_fps_is_missing(): + frames = FramesWithMetadata([Image.new("RGB", (2, 2))] * 2, frame_rate=12) + + _, fps = normalize_decoded_video_frames(frames, default_fps=24.0) + + assert fps == 12.0 + + +def test_normalize_decoded_video_frames_prefers_valid_fps_over_frame_rate(): + frames = FramesWithMetadata([Image.new("RGB", (2, 2))] * 2, fps=8, frame_rate=12) + + _, fps = normalize_decoded_video_frames(frames, default_fps=24.0) + + assert fps == 8.0 + + +def test_normalize_decoded_video_frames_uses_default_fps_for_plain_list_and_tuple(): + frames = [Image.new("RGB", (2, 2)), Image.new("RGB", (2, 2))] + + assert normalize_decoded_video_frames(frames, default_fps=24.0)[1] == 24.0 + assert normalize_decoded_video_frames(tuple(frames), default_fps=24.0)[1] == 24.0 + + +def test_normalize_decoded_video_frames_accepts_numeric_string_fps(): + frames = FramesWithMetadata([Image.new("RGB", (2, 2))] * 2, fps="12.5") + + _, frame_rate = normalize_decoded_video_frames(frames, default_fps=24.0) + + assert frame_rate == 12.5 + + +@pytest.mark.parametrize( + ("fps", "frame_rate", "expected"), + [ + (float("nan"), 12, 12.0), + (0, 15, 15.0), + (-1, 8, 8.0), + (float("inf"), 10, 10.0), + ("invalid", "6.5", 6.5), + ], +) +def test_normalize_decoded_video_frames_uses_valid_frame_rate_after_invalid_fps(fps, frame_rate, expected): + frames = FramesWithMetadata([Image.new("RGB", (2, 2))] * 2, fps=fps, frame_rate=frame_rate) + + _, frame_rate = normalize_decoded_video_frames(frames, default_fps=24.0) + + assert frame_rate == expected + + +def test_normalize_decoded_video_frames_uses_default_when_both_metadata_values_are_invalid(): + frames = FramesWithMetadata([Image.new("RGB", (2, 2))] * 2, fps="invalid", frame_rate=0) + + _, frame_rate = normalize_decoded_video_frames(frames, default_fps=24.0) + + assert frame_rate == 24.0 + + +@pytest.mark.parametrize( + ("frames", "match"), + [ + ([], "empty decoded video frame sequence"), + ([Image.new("RGB", (2, 2)), object()], "must be a PIL.Image.Image"), + ( + [Image.new("RGB", (2, 2)), Image.new("RGB", (3, 2))], + "must have identical dimensions", + ), + ], +) +def test_normalize_decoded_video_frames_rejects_invalid_sequences(frames, match): + with pytest.raises(ValueError, match=match): + normalize_decoded_video_frames(frames, default_fps=24.0) + + +@pytest.mark.parametrize( + ("frames", "match"), + [ + ([Image.new("RGB", (0, 16))], r"index 0.*0x16"), + ([Image.new("RGB", (16, 0))], r"index 0.*16x0"), + ([Image.new("RGB", (0, 0))], r"index 0.*0x0"), + ( + [Image.new("RGB", (16, 16)), Image.new("RGB", (0, 16))], + r"index 1.*0x16", + ), + ], +) +def test_normalize_decoded_video_frames_rejects_zero_dimensions(frames, match): + with pytest.raises(ValueError, match=match): + normalize_decoded_video_frames(frames, default_fps=24.0) + + +@pytest.mark.parametrize( + ("video_input", "expected_fps"), + [ + (FramesWithMetadata([Image.new("RGB", (3, 2))] * 2, fps=8), 8.0), + (tuple([Image.new("RGBA", (3, 2))] * 2), 12.0), + ], +) +def test_lance_video_edit_normalizes_decoded_frame_sequences(monkeypatch, video_input, expected_fps): + captured = {} + + def capture_preprocessed_input(video, fps): + captured["video"] = video + captured["fps"] = fps + raise StopAfterVideoPreprocessError + + monkeypatch.setattr(LanceBagel, "_lance_video_preprocess", staticmethod(capture_preprocessed_input)) + request = SimpleNamespace( + prompts=[ + { + "modalities": ["video"], + "user_text": "edit the video", + "multi_modal_data": {"video": video_input}, + "extra_args": {"origin_fps": 12.0}, + } + ], + sampling_params=SimpleNamespace(extra_args={}, num_inference_steps=1), + ) + + with pytest.raises(StopAfterVideoPreprocessError): + object.__new__(LancePipeline)._forward_video_edit(request) + + video = captured["video"] + assert isinstance(video, np.ndarray) + assert video.shape == (2, 2, 3, 3) + assert video.dtype == np.uint8 + assert video.flags.c_contiguous + assert captured["fps"] == expected_fps diff --git a/vllm_omni/diffusion/models/lance/pipeline_lance.py b/vllm_omni/diffusion/models/lance/pipeline_lance.py index dc885ac44f9..f33a35b9670 100644 --- a/vllm_omni/diffusion/models/lance/pipeline_lance.py +++ b/vllm_omni/diffusion/models/lance/pipeline_lance.py @@ -62,6 +62,7 @@ from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) +from vllm_omni.model_executor.models.utils import normalize_decoded_video_frames from .lance_transformer import ( LanceBagel, @@ -1454,6 +1455,8 @@ def _forward_video_edit(self, req): frames_bgr.append(_cv2.cvtColor(f, _cv2.COLOR_BGR2RGB)) cap.release() video_raw = _np.stack(frames_bgr, axis=0) + elif isinstance(video_input, (list, tuple)): + video_raw, origin_fps = normalize_decoded_video_frames(video_input, default_fps=origin_fps_default) elif isinstance(video_input, _np.ndarray): video_raw = video_input origin_fps = origin_fps_default diff --git a/vllm_omni/model_executor/models/utils.py b/vllm_omni/model_executor/models/utils.py index fd9ddd41bc9..73e9d5572b9 100644 --- a/vllm_omni/model_executor/models/utils.py +++ b/vllm_omni/model_executor/models/utils.py @@ -1,10 +1,59 @@ -from collections.abc import Callable +import math +from collections.abc import Callable, Sequence from contextlib import contextmanager +import numpy as np import torch +from PIL import Image from vllm.model_executor.models.utils import maybe_prefix +def normalize_decoded_video_frames( + video_input: Sequence[Image.Image], + *, + default_fps: float, +) -> tuple[np.ndarray, float]: + """Convert decoded image frames to a contiguous THWC uint8 video array.""" + if not video_input: + raise ValueError("video_edit received an empty decoded video frame sequence.") + + frame_rate = float(default_fps) + for fps in ( + getattr(video_input, "fps", None), + getattr(video_input, "frame_rate", None), + ): + try: + candidate = float(fps) + except (TypeError, ValueError): + continue + if math.isfinite(candidate) and candidate > 0: + frame_rate = candidate + break + + frames = [] + frame_size = None + for index, frame in enumerate(video_input): + if not isinstance(frame, Image.Image): + raise ValueError( + f"video_edit decoded video frame at index {index} must be a PIL.Image.Image, got {type(frame)}." + ) + width, height = frame.size + if width <= 0 or height <= 0: + raise ValueError( + f"video_edit decoded video frame at index {index} must have positive dimensions, got {width}x{height}." + ) + if frame_size is None: + frame_size = frame.size + elif frame.size != frame_size: + raise ValueError( + "video_edit decoded video frames must have identical dimensions; " + f"frame 0 is {frame_size}, frame {index} is {frame.size}." + ) + frames.append(np.asarray(frame.convert("RGB"), dtype=np.uint8)) + + return np.ascontiguousarray(np.stack(frames)), frame_rate + + @contextmanager def transformers_keys_to_ignore_compat(): """Make ``trust_remote_code`` weight loading robust to the transformers 5.9 From 28291df1795129ae03ce86a7b9b0a230166b5ac5 Mon Sep 17 00:00:00 2001 From: Alicia <115451386+congw729@users.noreply.github.com> Date: Sat, 12 Sep 2026 20:30:03 +0800 Subject: [PATCH 89/93] [Refactor][Diffusion] Remove model-specific names from LoRA and ModelOpt loader defaults (#5907) Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com> Co-authored-by: Hongsheng Liu --- tests/diffusion/lora/test_lora_manager.py | 4 +++- tests/diffusion/models/bagel/test_bagel_lora.py | 9 +++++++-- vllm_omni/diffusion/lora/manager.py | 17 ++++++++++------- .../checkpoint_adapters/modelopt.py | 9 ++++++++- .../diffusion/models/bagel/pipeline_bagel.py | 7 ++++++- .../models/z_image/z_image_transformer.py | 9 +++++++++ 6 files changed, 43 insertions(+), 12 deletions(-) diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py index 10b023bf0ee..b2a0505f302 100644 --- a/tests/diffusion/lora/test_lora_manager.py +++ b/tests/diffusion/lora/test_lora_manager.py @@ -729,8 +729,10 @@ def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs): lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls), ) - # Pipeline with a 'bagel' component (no 'transformer') + # Pipeline with a 'bagel' component (no 'transformer') declared via + # _lora_components, matching how BagelPipeline opts in. pipeline = torch.nn.Module() + pipeline._lora_components = ["bagel"] pipeline.bagel = torch.nn.Module() pipeline.bagel.language_model = torch.nn.Module() pipeline.bagel.language_model.qkv_proj = _FakeLinearBase() diff --git a/tests/diffusion/models/bagel/test_bagel_lora.py b/tests/diffusion/models/bagel/test_bagel_lora.py index c285758fe86..1e73318ce37 100644 --- a/tests/diffusion/models/bagel/test_bagel_lora.py +++ b/tests/diffusion/models/bagel/test_bagel_lora.py @@ -122,8 +122,10 @@ def _fake_from_layer_diffusion(*, layer, **_kwargs): lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls), ) - # Build pipeline with bagel component + # Build pipeline with bagel component declared via _lora_components, + # matching how BagelPipeline opts in. pipeline = torch.nn.Module() + pipeline._lora_components = ["bagel"] pipeline.bagel = torch.nn.Module() lm = torch.nn.Module() lm.stacked_params_mapping = [ @@ -194,8 +196,11 @@ def test_set_active_adapter_loads_and_activates_bagel_lora(self, tmp_path, monke monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", DummyBaseLayerWithLoRA) - # Build pipeline with bagel.language_model.foo (simple non-packed layer) + # Build pipeline with bagel.language_model.foo (simple non-packed layer); + # the bagel component is declared via _lora_components, matching how + # BagelPipeline opts in. pipeline = torch.nn.Module() + pipeline._lora_components = ["bagel"] pipeline.bagel = torch.nn.Module() lm = torch.nn.Module() lm.foo = _FakeLinearBase() diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 11d19353158..3a82305307d 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -394,17 +394,20 @@ def _matches_target(module_name: str) -> bool: fully_sharded_loras=False, ) - # Default denoising components, declared DiT components, and any a - # pipeline opts into via ``_lora_components``. + # Components scanned for LoRA-capable layers: framework defaults, + # declared DiT components, and any a pipeline opts into via + # ``_lora_components``. The defaults only cover the generic diffusers + # naming convention: ``transformer`` (plus ``transformer_2`` for + # dual-DiT pipelines such as Wan2.2) and ``unet`` for SDXL-style + # pipelines. Model-specific attribute names must be declared by the + # pipeline itself via ``_dit_modules`` or ``_lora_components``. # - # NOTE: SDXL-style pipelines expose the denoiser as ``unet``. - # Without scanning this component, adapters can load/activate while - # effectively applying to zero layers, producing base-identical output. + # NOTE: if the denoiser component is not scanned here, adapters can + # load/activate while effectively applying to zero layers, producing + # base-identical output. default_components = ( "transformer", "transformer_2", - "dit", - "bagel", "unet", ) declared_components = tuple(getattr(self.pipeline, "_dit_modules", ()) or ()) diff --git a/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py b/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py index 68ec98d5ae7..f3c703b090e 100644 --- a/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py +++ b/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py @@ -17,10 +17,17 @@ ".weight_scale_2", ".weight_scale_inv", ) +# Fused projections following the generic diffusers Attention naming +# convention. ``to_q``/``to_k``/``to_v`` and the joint-attention +# ``add_q_proj``/``add_k_proj``/``add_v_proj`` shard names are fused into +# ``to_qkv``/``add_kv_proj`` by most vllm-omni DiT implementations, many of +# which do not declare their own ``packed_modules_mapping``; keeping these +# two entries here avoids repeating them in every such model. Model-specific +# fusions (e.g. Z-Image's ``w13``) must be declared on the model class via +# ``packed_modules_mapping``, which is merged on top of these defaults. DEFAULT_PACKED_MODULES_MAPPING = { "to_qkv": ("to_q", "to_k", "to_v"), "add_kv_proj": ("add_q_proj", "add_k_proj", "add_v_proj"), - "w13": ("w1", "w3"), } FP8_DTYPES = tuple( dtype diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index f2760114837..d805fd3b849 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.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 """ BagelPipeline implementation for vLLM-Omni. """ @@ -316,6 +316,11 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf "bagel.connector", "bagel.vit_pos_embed", ] + # LoRA scanning target. ``_dit_modules`` holds the dotted path + # "language_model.model", which DiffusionLoRAManager cannot resolve as a + # pipeline attribute, so the top-level ``bagel`` module (LLM plus ViT, + # connector, and VAE projections) is declared explicitly here. + _lora_components: ClassVar[list[str]] = ["bagel"] supports_step_execution: ClassVar[bool] = True def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): diff --git a/vllm_omni/diffusion/models/z_image/z_image_transformer.py b/vllm_omni/diffusion/models/z_image/z_image_transformer.py index 1399eb0f0b3..f25c01fca87 100644 --- a/vllm_omni/diffusion/models/z_image/z_image_transformer.py +++ b/vllm_omni/diffusion/models/z_image/z_image_transformer.py @@ -612,6 +612,15 @@ class ZImageTransformer2DModel(CachedTransformer): Note: Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in diffusers. """ + # Fused projections and their checkpoint shard names (mirrors + # stacked_params_mapping in load_weights). ``w13`` is the Z-Image FFN + # fusion of ``w1``/``w3``; declared here so quantized-checkpoint adapters + # can resolve shard keys without a framework-level default entry. + packed_modules_mapping = { + "to_qkv": ["to_q", "to_k", "to_v"], + "w13": ["w1", "w3"], + } + _repeated_blocks = ["ZImageTransformerBlock"] _layerwise_offload_blocks_attrs = ["layers"] From 9edc1b4667d4b56d12effe1e5b300315520d1ae3 Mon Sep 17 00:00:00 2001 From: geray <48796550+gerayking@users.noreply.github.com> Date: Sat, 12 Sep 2026 20:31:49 +0800 Subject: [PATCH 90/93] Optimize CosyVoice3 Stage1 flow batching (#4876) Signed-off-by: gerayking <399geray@gmail.com> Co-authored-by: Hongsheng Liu --- .../examples/online_serving/text_to_speech.md | 9 + tests/config/test_environment_variables.py | 6 +- .../cosyvoice3/test_cosyvoice3_components.py | 252 ++++++++++++++++++ .../test_cosyvoice3_model_helpers.py | 64 +++++ .../config/environment_variable_inventory.py | 2 + .../models/cosyvoice3_audio/cosyvoice3_dit.py | 185 +++++++------ .../models/cosyvoice3/code2wav_core/cfm.py | 188 +++++++------ .../models/cosyvoice3/cosyvoice3.py | 97 +++++-- .../models/cosyvoice3/cosyvoice3_code2wav.py | 213 ++++++++++++--- .../models/cosyvoice3/runtime.py | 31 +++ 10 files changed, 837 insertions(+), 210 deletions(-) create mode 100644 vllm_omni/model_executor/models/cosyvoice3/runtime.py diff --git a/docs/user_guide/examples/online_serving/text_to_speech.md b/docs/user_guide/examples/online_serving/text_to_speech.md index 9a33dc8583a..2b4f617f42d 100644 --- a/docs/user_guide/examples/online_serving/text_to_speech.md +++ b/docs/user_guide/examples/online_serving/text_to_speech.md @@ -120,6 +120,15 @@ vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remot Streaming is on by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path. +Cross-request Stage-1 flow batching is opt-in. Enable it in the server environment when concurrent requests should share a flow-estimator call: + +```bash +export COSYVOICE3_BATCH_FLOW=1 +vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remote-code +``` + +Batching preserves output lengths and streaming cache alignment, but the different GEMM shapes can produce small waveform differences compared with processing each request separately. Leave `COSYVOICE3_BATCH_FLOW` unset (or set it to `0`) when request-independent numerical behavior is required. Set `COSYVOICE3_BATCH_FLOW_DEBUG=1` to log the observed group-size distribution and enable detailed Stage-1 profiler scopes; diagnostics are disabled by default to avoid per-step profiling overhead. + ### CLI client The client defaults to the official upstream zero-shot prompt, so it runs without extra flags: diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index 7603945a0bf..cc37cd6b155 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: 24, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, - EnvironmentVariableCategory.MODEL_SPECIFIC: 61, + EnvironmentVariableCategory.MODEL_SPECIFIC: 63, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20, EnvironmentVariableCategory.INTERNAL: 2, } @@ -178,10 +178,10 @@ 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: 37, + ModelEnvironmentVariableDisposition.PROMOTE: 38, ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6, ModelEnvironmentVariableDisposition.EXTERNAL: 0, - ModelEnvironmentVariableDisposition.INTERNALIZE: 13, + ModelEnvironmentVariableDisposition.INTERNALIZE: 14, ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 5, } diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py index 4f01022f8e4..1b1df636728 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for CosyVoice3 components.""" +import types from contextlib import contextmanager from types import SimpleNamespace @@ -9,6 +10,7 @@ import torch import torch.nn as nn +from tests.helpers.mark import hardware_test from vllm_omni.diffusion.config import set_current_diffusion_config from vllm_omni.diffusion.data import AttentionConfig from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( @@ -29,6 +31,23 @@ def _force_torch_sdpa(): yield +def test_batch_flow_runtime_flags_are_opt_in(monkeypatch): + from vllm_omni.model_executor.models.cosyvoice3.runtime import ( + cosyvoice3_batch_flow_debug, + cosyvoice3_batch_flow_enabled, + ) + + monkeypatch.delenv("COSYVOICE3_BATCH_FLOW", raising=False) + monkeypatch.delenv("COSYVOICE3_BATCH_FLOW_DEBUG", raising=False) + assert not cosyvoice3_batch_flow_enabled() + assert not cosyvoice3_batch_flow_debug() + + monkeypatch.setenv("COSYVOICE3_BATCH_FLOW", "1") + monkeypatch.setenv("COSYVOICE3_BATCH_FLOW_DEBUG", "1") + assert cosyvoice3_batch_flow_enabled() + assert cosyvoice3_batch_flow_debug() + + @pytest.fixture def causal_hift(): return CausalHiFTGenerator( @@ -139,6 +158,39 @@ def test_qkv_projections(self, attention): assert attention.to_k.out_features == 512 assert attention.to_v.out_features == 512 + @pytest.mark.core_model + @pytest.mark.cpu + def test_forward_passes_2d_mask_without_casting_qkv(self): + """The float32 fallback receives the per-row validity mask.""" + from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiTAttention + + class CaptureAttention(nn.Module): + def __init__(self): + super().__init__() + self.last_metadata = None + self.last_dtype = None + + def forward(self, q, k, v, attn_metadata=None): + self.last_metadata = attn_metadata + self.last_dtype = q.dtype + return torch.zeros_like(q) + + attention = DiTAttention(dim=16, heads=2, dim_head=8, dropout=0.0) + capture = CaptureAttention() + attention.attn = capture + + x = torch.randn(2, 5, 16) + mask = torch.tensor([[True, True, True, False, False], [True, True, True, True, False]]) + + out = attention(x, mask=mask) + + assert out.shape == x.shape + assert capture.last_metadata is not None + assert torch.equal(capture.last_metadata.attn_mask, mask) + assert capture.last_dtype == x.dtype + assert out.dtype == x.dtype + assert torch.allclose(out[~mask], torch.zeros_like(out[~mask])) + class TestDiTBlock: """Tests for DiTBlock.""" @@ -422,6 +474,52 @@ def reject_stream_context(stream): assert reused_context is context assert reused_stream is stream + @pytest.mark.core_model + @pytest.mark.cpu + def test_causal_conditional_cfm_batches_cfg_estimator(self): + """Batched flow calls should invoke the CFG estimator with 2B rows.""" + from omegaconf import DictConfig + + from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import CausalConditionalCFM + + class RecordingEstimator(nn.Module): + def __init__(self): + super().__init__() + self.shapes: list[tuple[int, ...]] = [] + + def forward(self, x, mask, mu, t, spks=None, cond=None): + self.shapes.append(tuple(x.shape)) + return torch.zeros_like(x) + + estimator = RecordingEstimator() + cfm_params = DictConfig( + { + "sigma_min": 1e-6, + "solver": "euler", + "t_scheduler": "cosine", + "training_cfg_rate": 0.2, + "inference_cfg_rate": 0.7, + } + ) + cfm = CausalConditionalCFM( + in_channels=80, + cfm_params=cfm_params, + n_spks=1, + spk_emb_dim=80, + estimator=estimator, + ) + + batch, mel_dim, seq_len = 3, 80, 16 + mu = torch.randn(batch, mel_dim, seq_len) + mask = torch.ones(batch, 1, seq_len) + spks = torch.randn(batch, 80) + cond = torch.randn(batch, mel_dim, seq_len) + + out, _ = cfm(mu, mask, n_timesteps=2, spks=spks, cond=cond) + + assert out.shape == mu.shape + assert estimator.shapes == [(2 * batch, mel_dim, seq_len)] * 2 + class TestSDPAFallback: """Test SDPA fallback for float32 inputs.""" @@ -484,3 +582,157 @@ def fake_forward_mel(**kwargs): assert out.shape == (1, 1, 8) assert model.hift.finalize_calls == [True] assert forward_mel_calls[0]["token_offset_tokens"] == 0 + + +def test_code2wav_streaming_batch_pads_codec_tokens_and_preserves_lengths(): + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav + + model = object.__new__(CosyVoice3Code2Wav) + model.flow_model = SimpleNamespace(pre_lookahead_len=1, token_mel_ratio=2) + forward_mel_calls = [] + + def fake_forward_mel(self, **kwargs): + forward_mel_calls.append(kwargs) + return torch.arange(2 * 80 * 8, dtype=torch.float32).reshape(2, 80, 8) + + def fake_stream_hift(self, feat, *, cache_state=None, finalize=False): + return feat, None + + model._forward_mel = types.MethodType(fake_forward_mel, model) + model._stream_hift_from_feat = types.MethodType(fake_stream_hift, model) + items = [ + { + "token": torch.ones(1, 3, dtype=torch.int32), + "prompt_token": torch.ones(1, 4, dtype=torch.int32), + "prompt_feat": torch.ones(1, 8, 80), + "embedding": torch.ones(1, 192), + "finalize": False, + }, + { + "token": torch.ones(1, 5, dtype=torch.int32), + "prompt_token": torch.ones(1, 4, dtype=torch.int32), + "prompt_feat": torch.ones(1, 8, 80), + "embedding": torch.ones(1, 192), + "finalize": False, + }, + ] + + results = model.forward_streaming_batch(items) + + assert len(forward_mel_calls) == 1 + call = forward_mel_calls[0] + assert call["token"].shape == (2, 5) + assert torch.equal(call["token_lens"], torch.tensor([3, 5], dtype=torch.int32)) + assert results[0][0].shape[-1] == 4 + assert results[1][0].shape[-1] == 8 + + +@pytest.mark.core_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_code2wav_streaming_batch_matches_ragged_flow_numerics(monkeypatch): + """A padded flow batch must match individual flow calls on valid mels.""" + from omegaconf import DictConfig + + from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiT + from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import ( + CausalConditionalCFM, + CausalMaskedDiffWithDiT, + ) + from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.layers import PreLookaheadLayer + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav + + torch.manual_seed(0) + estimator = DiT( + dim=32, + depth=1, + heads=4, + dim_head=8, + dropout=0.0, + ff_mult=2, + mel_dim=80, + mu_dim=80, + spk_dim=80, + out_channels=80, + ) + decoder = CausalConditionalCFM( + in_channels=80, + cfm_params=DictConfig( + { + "sigma_min": 1e-6, + "solver": "euler", + "t_scheduler": "cosine", + "training_cfg_rate": 0.2, + "inference_cfg_rate": 0.7, + } + ), + n_spks=1, + spk_emb_dim=80, + estimator=estimator, + ) + flow_model = ( + CausalMaskedDiffWithDiT( + input_size=80, + output_size=80, + spk_embed_dim=192, + vocab_size=64, + input_frame_rate=25, + only_mask_loss=True, + token_mel_ratio=2, + pre_lookahead_len=1, + pre_lookahead_layer=PreLookaheadLayer(in_channels=80, channels=80, pre_lookahead_len=1), + decoder=decoder, + ) + .to(device="cuda", dtype=torch.bfloat16) + .eval() + ) + + model = object.__new__(CosyVoice3Code2Wav) + nn.Module.__init__(model) + model.flow_model = flow_model + + def return_mel(self, feat, *, cache_state=None, finalize=False): + return feat, None + + model._stream_hift_from_feat = types.MethodType(return_mel, model) + + original_randn = torch.randn + + def length_consistent_randn(*size, **kwargs): + shape = tuple(size[0]) if len(size) == 1 and isinstance(size[0], (tuple, list)) else tuple(size) + if len(shape) == 3 and shape[1] == 80: + device = kwargs.get("device") + dtype = kwargs.get("dtype", torch.float32) + channels = torch.arange(shape[1], device=device, dtype=torch.float32).view(1, -1, 1) + positions = torch.arange(shape[2], device=device, dtype=torch.float32).view(1, 1, -1) + noise = torch.sin(channels * 0.17 + positions * 0.31) + return noise.expand(shape[0], -1, -1).to(dtype=dtype).clone() + return original_randn(*size, **kwargs) + + monkeypatch.setattr(torch, "randn", length_consistent_randn) + common = { + "prompt_token": torch.tensor([[7, 8]], dtype=torch.int32), + "prompt_feat": torch.linspace(-0.5, 0.5, 4 * 80).reshape(1, 4, 80), + "embedding": torch.linspace(-1.0, 1.0, 192).reshape(1, 192), + "finalize": False, + } + items = [ + {**common, "token": torch.tensor([[1, 2, 3]], dtype=torch.int32)}, + {**common, "token": torch.tensor([[1, 2, 3, 4, 5]], dtype=torch.int32)}, + ] + + batched = model.forward_streaming_batch(items, n_timesteps=2) + individual = [ + model.forward_streaming( + token=item["token"], + prompt_token=item["prompt_token"], + prompt_feat=item["prompt_feat"], + embedding=item["embedding"], + n_timesteps=2, + ) + for item in items + ] + + for (batched_mel, _), (individual_mel, _) in zip(batched, individual): + assert batched_mel.shape == individual_mel.shape + rel_mean = (batched_mel - individual_mel).abs().mean() / individual_mel.abs().mean().clamp_min(1e-6) + assert rel_mean.item() < 0.05 diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py index e1692709629..1d583de6c43 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py @@ -42,6 +42,7 @@ def __init__( self.outputs = list(outputs or []) self.forward_calls: list[dict[str, object]] = [] self.forward_streaming_calls: list[dict[str, object]] = [] + self.forward_streaming_batch_calls: list[list[dict[str, object]]] = [] def forward(self, **kwargs): self.forward_calls.append(kwargs) @@ -65,6 +66,22 @@ def forward_streaming(self, **kwargs): } return audio, new_state + def forward_streaming_batch(self, items, *, n_timesteps: int = 10): + self.forward_streaming_batch_calls.append(items) + return [ + self.forward_streaming( + token=item["token"], + prompt_token=item["prompt_token"], + prompt_feat=item["prompt_feat"], + embedding=item["embedding"], + cache_state=item.get("cache_state"), + n_timesteps=n_timesteps, + token_offset_tokens=int(item.get("token_offset_tokens", 0)), + finalize=bool(item.get("finalize", False)), + ) + for item in items + ] + def _make_code2wav_model( *, @@ -373,6 +390,53 @@ def test_forward_clears_streaming_cache_on_terminal_chunk(): assert "rid-stream" not in model._stream_vocoder_cache_by_req +def test_forward_batches_streaming_flow_items(monkeypatch): + monkeypatch.setenv("COSYVOICE3_BATCH_FLOW", "1") + model = _make_code2wav_model() + runtime_info = [ + { + "embed": { + "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long), + "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32), + "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32), + }, + "meta": { + "req_id": ["rid-a"], + "stream_finished": torch.tensor(False), + "left_context_size": 0, + }, + }, + { + "embed": { + "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long), + "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32), + "embedding": torch.tensor([[0.7, 0.8]], dtype=torch.float32), + }, + "meta": { + "req_id": ["rid-b"], + "stream_finished": torch.tensor(False), + "left_context_size": 1, + }, + }, + ] + + out = model.forward( + input_ids=torch.tensor([0, 1, 2, 0, 1, 2], dtype=torch.long), + positions=torch.arange(6, dtype=torch.long), + model_intermediate_buffer=runtime_info, + seq_token_counts=[3, 3], + ) + + assert len(out.multimodal_outputs["audio"]) == 2 + assert len(model.code2wav.forward_streaming_batch_calls) == 1 + batch_items = model.code2wav.forward_streaming_batch_calls[0] + assert [item["index"] for item in batch_items] == [0, 1] + assert torch.equal(batch_items[0]["token"], torch.tensor([[0, 1, 2]])) + assert batch_items[1]["token_offset_tokens"] == 1 + assert "rid-a" in model._stream_vocoder_cache_by_req + assert "rid-b" in model._stream_vocoder_cache_by_req + + def test_sample_uses_ras_rejection_for_recent_repetition(): model = _make_talker_model() metadata = _make_sampling_metadata(output_token_ids=[[1] * 10]) diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index c1e380a22ae..4a4966d9c2a 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -137,6 +137,7 @@ def is_public_omni(self) -> bool: "AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH", "COSMOS3_SOUND_TOKENIZER_CONFIG_PATH", "COSMOS3_SOUND_TOKENIZER_PATH", + "COSYVOICE3_BATCH_FLOW", "COSYVOICE3_ESTIMATOR_ONNX", "COSYVOICE3_TRT", "COSYVOICE3_TRT_CACHE", @@ -188,6 +189,7 @@ def is_public_omni(self) -> bool: _MODEL_EXTERNAL: tuple[str, ...] = () _MODEL_INTERNALIZE = ( + "COSYVOICE3_BATCH_FLOW_DEBUG", "DOTS_TTS_BETA_TRACE", "DYNIN_S2U_VENDOR_ROOT", "DZ_PHASE_TIMING", diff --git a/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py b/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py index 5ae75b25617..d0538d79f94 100644 --- a/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.py +++ b/vllm_omni/diffusion/models/cosyvoice3_audio/cosyvoice3_dit.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 # Adapted from https://github.com/FunAudioLLM/CosyVoice/tree/main/cosyvoice/flow/DiT # Refactored to use vllm_omni diffusion infrastructure for optimized attention backends @@ -15,8 +15,10 @@ from vllm.logger import init_logger from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention as DiffusionAttention from vllm_omni.model_executor.layers.timestep_embedding import DiTTimestepEmbedding +from vllm_omni.model_executor.models.cosyvoice3.runtime import cosyvoice3_batch_flow_profile logger = init_logger(__name__) @@ -52,6 +54,22 @@ def get_pos_embed_indices(start, length, max_pos, scale=1.0): return pos +def _normalize_attention_mask(mask: torch.Tensor | None, batch_size: int, seq_len: int) -> torch.Tensor | None: + if mask is None: + return None + if mask.dim() == 2: + attn_mask = mask + elif mask.dim() == 3 and mask.shape[1] == 1: + attn_mask = mask[:, 0] + elif mask.dim() == 4: + attn_mask = mask[:, 0, -1] + else: + return None + if attn_mask.shape != (batch_size, seq_len): + return None + return attn_mask.bool() + + class FeedForward(nn.Module): """Feed-forward network with GELU activation.""" @@ -114,42 +132,39 @@ def forward( ) -> torch.Tensor: batch_size, seq_len = x.shape[0], x.shape[1] - # Project to Q, K, V - query = self.to_q(x) - key = self.to_k(x) - value = self.to_v(x) - - # Apply rotary position embedding - if rope is not None: - freqs, xpos_scale = rope - q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) - query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) - key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) - - # Reshape for attention: (batch, seq, heads, head_dim) - query = query.view(batch_size, seq_len, self.heads, self.dim_head) - key = key.view(batch_size, seq_len, self.heads, self.dim_head) - value = value.view(batch_size, seq_len, self.heads, self.dim_head) - - # Use diffusion attention backend - # The diffusion Attention layer expects (batch, seq, heads, head_dim) - out = self.attn(query, key, value, attn_metadata=None) - - # Some attention backends return a non-contiguous layout. - out = out.reshape(batch_size, seq_len, self.inner_dim) - out = out.to(query.dtype) - - # Output projection - out = self.to_out(out) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_attention_qkv_projection"): + query = self.to_q(x) + key = self.to_k(x) + value = self.to_v(x) + + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_attention_rope"): + if rope is not None: + freqs, xpos_scale = rope + q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) + key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_attention_backend"): + # Reshape for attention: (batch, seq, heads, head_dim) + query = query.view(batch_size, seq_len, self.heads, self.dim_head) + key = key.view(batch_size, seq_len, self.heads, self.dim_head) + value = value.view(batch_size, seq_len, self.heads, self.dim_head) + + # The diffusion Attention layer expects (batch, seq, heads, head_dim) + attn_mask = _normalize_attention_mask(mask, batch_size, seq_len) + attn_metadata = AttentionMetadata(attn_mask=attn_mask) if attn_mask is not None else None + out = self.attn(query, key, value, attn_metadata=attn_metadata) + + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_attention_output_projection"): + # Some attention backends return a non-contiguous layout. + out = out.reshape(batch_size, seq_len, self.inner_dim) + out = out.to(x.dtype) + out = self.to_out(out) # Apply mask if provided - if mask is not None: - if mask.dim() == 2: - mask = mask.unsqueeze(-1) - elif mask.dim() == 4: - # (batch, heads, seq, seq) -> use last dim - mask = mask[:, 0, -1].unsqueeze(-1) - out = out.masked_fill(~mask.bool(), 0.0) + attn_mask = _normalize_attention_mask(mask, batch_size, seq_len) + if attn_mask is not None: + out = out.masked_fill(~attn_mask.unsqueeze(-1), 0.0) return out @@ -172,18 +187,19 @@ def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1): self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh") def forward(self, x, t, mask=None, rope=None): - # pre-norm & modulation for attention input - norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_adalayernorm_attention"): + norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t) - # attention - attn_output = self.attn(x=norm, mask=mask, rope=rope) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_attention"): + attn_output = self.attn(x=norm, mask=mask, rope=rope) - # process attention output for input x - x = x + gate_msa.unsqueeze(1) * attn_output + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_attention_residual"): + x = x + gate_msa.unsqueeze(1) * attn_output - ff_norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] - ff_output = self.ff(ff_norm) - x = x + gate_mlp.unsqueeze(1) * ff_output + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_layernorm_ffn"): + ff_norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + ff_output = self.ff(ff_norm) + x = x + gate_mlp.unsqueeze(1) * ff_output return x @@ -205,19 +221,20 @@ def __init__(self, dim, kernel_size=31, groups=16): ) def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None): - if mask is not None: - mask = mask[..., None] - x = x.masked_fill(~mask, 0.0) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_causal_conv_position_embedding"): + if mask is not None: + mask = mask[..., None] + x = x.masked_fill(~mask, 0.0) - x = x.permute(0, 2, 1) - x = F.pad(x, (self.kernel_size - 1, 0, 0, 0)) - x = self.conv1(x) - x = F.pad(x, (self.kernel_size - 1, 0, 0, 0)) - x = self.conv2(x) - out = x.permute(0, 2, 1) + x = x.permute(0, 2, 1) + x = F.pad(x, (self.kernel_size - 1, 0, 0, 0)) + x = self.conv1(x) + x = F.pad(x, (self.kernel_size - 1, 0, 0, 0)) + x = self.conv2(x) + out = x.permute(0, 2, 1) - if mask is not None: - out = out.masked_fill(~mask, 0.0) + if mask is not None: + out = out.masked_fill(~mask, 0.0) return out @@ -327,13 +344,15 @@ def __init__(self, mel_dim, text_dim, out_dim, spk_dim=None): self.conv_pos_embed = CausalConvPositionEmbedding(dim=out_dim) def forward(self, x, cond, text_embed, spks): - to_cat = [x, cond, text_embed] - if self.spk_dim > 0: - spks = repeat(spks, "b c -> b t c", t=x.shape[1]) - to_cat.append(spks) - - x = self.proj(torch.cat(to_cat, dim=-1)) - x = self.conv_pos_embed(x) + x + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_input_projection"): + to_cat = [x, cond, text_embed] + if self.spk_dim > 0: + spks = repeat(spks, "b c -> b t c", t=x.shape[1]) + to_cat.append(spks) + + x = self.proj(torch.cat(to_cat, dim=-1)) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_input_causal_conv"): + x = self.conv_pos_embed(x) + x return x @@ -385,31 +404,37 @@ def __init__( self.num_decoding_left_chunks = num_decoding_left_chunks def forward(self, x, mask, mu, t, spks=None, cond=None): - x = x.transpose(1, 2) - mu = mu.transpose(1, 2) - cond = cond.transpose(1, 2) - spks = spks.unsqueeze(dim=1) - batch, seq_len = x.shape[0], x.shape[1] - if t.ndim == 0: - t = t.repeat(batch) - - # t: conditioning time, c: context (text + masked cond audio), x: noised input audio - t = self.time_embed(t) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_prepare_inputs"): + x = x.transpose(1, 2) + mu = mu.transpose(1, 2) + cond = cond.transpose(1, 2) + spks = spks.unsqueeze(dim=1) + batch, seq_len = x.shape[0], x.shape[1] + if t.ndim == 0: + t = t.repeat(batch) + + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_time_embedding"): + t = self.time_embed(t) x = self.input_embed(x, cond, mu, spks.squeeze(1)) - rope = self.rotary_embed.forward_from_seq_len(seq_len) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_rope_and_mask"): + rope = self.rotary_embed.forward_from_seq_len(seq_len) + attn_mask = mask[:, 0].bool() if mask.dim() == 3 and mask.shape[1] == 1 else mask.bool() if self.long_skip_connection is not None: residual = x - attn_mask = mask.bool().repeat(1, x.size(1), 1).unsqueeze(dim=1) - - for block in self.transformer_blocks: - x = block(x, t, mask=attn_mask.bool(), rope=rope) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_transformer_blocks"): + for block_index, block in enumerate(self.transformer_blocks): + with cosyvoice3_batch_flow_profile(f"cosyvoice3_dit_transformer_block_{block_index}"): + x = block(x, t, mask=attn_mask.bool(), rope=rope) if self.long_skip_connection is not None: - x = self.long_skip_connection(torch.cat((x, residual), dim=-1)) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_long_skip_projection"): + x = self.long_skip_connection(torch.cat((x, residual), dim=-1)) - x = self.norm_out(x, t) - output = self.proj_out(x).transpose(1, 2) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_final_adalayernorm"): + x = self.norm_out(x, t) + with cosyvoice3_batch_flow_profile("cosyvoice3_dit_final_projection"): + output = self.proj_out(x).transpose(1, 2) return output diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py index 2ec12a25124..c86b4a3055b 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.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 # Adopted from https://github.com/FunAudioLLM/CosyVoice/tree/main/cosyvoice/flow """Conditional Flow Matching (CFM) classes for audio generation.""" @@ -11,6 +11,7 @@ from torch.nn import functional as F from vllm.logger import init_logger +from vllm_omni.model_executor.models.cosyvoice3.runtime import cosyvoice3_batch_flow_profile from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask logger = init_logger(__name__) @@ -74,20 +75,23 @@ def forward( shape: (batch_size, n_feats, mel_timesteps) """ - z = torch.randn_like(mu).to(mu.device).to(mu.dtype) * temperature - cache_size = cache.shape[2] - # fix prompt and overlap part mu and z - if cache_size != 0: - z[:, :, :cache_size] = cache[:, :, :, 0] - mu[:, :, :cache_size] = cache[:, :, :, 1] - z_cache = torch.concat([z[:, :, :prompt_len], z[:, :, -34:]], dim=2) - mu_cache = torch.concat([mu[:, :, :prompt_len], mu[:, :, -34:]], dim=2) - cache = torch.stack([z_cache, mu_cache], dim=-1) - - t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) - if self.t_scheduler == "cosine": - t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) - return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), cache + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_noise_cache"): + z = torch.randn_like(mu).to(mu.device).to(mu.dtype) * temperature + cache_size = cache.shape[2] + # fix prompt and overlap part mu and z + if cache_size != 0: + z[:, :, :cache_size] = cache[:, :, :, 0] + mu[:, :, :cache_size] = cache[:, :, :, 1] + z_cache = torch.concat([z[:, :, :prompt_len], z[:, :, -34:]], dim=2) + mu_cache = torch.concat([mu[:, :, :prompt_len], mu[:, :, -34:]], dim=2) + cache = torch.stack([z_cache, mu_cache], dim=-1) + + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_t_span"): + t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) + if self.t_scheduler == "cosine": + t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) + with cosyvoice3_batch_flow_profile(f"cosyvoice3_cfm_euler_{max(1, int(n_timesteps))}_steps"): + return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), cache def solve_euler(self, x, t_span, mu, mask, spks, cond): """ @@ -111,26 +115,39 @@ def solve_euler(self, x, t_span, mu, mask, spks, cond): # Do not use concat, it may cause memory format changed and trt infer with wrong results! # NOTE when flow run in amp mode, x.dtype is float32, which cause nan in trt fp16 - # inference, so set dtype=spks.dtype - x_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=spks.dtype) - mask_in = torch.zeros([2, 1, x.size(2)], device=x.device, dtype=spks.dtype) - mu_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=spks.dtype) - t_in = torch.zeros([2], device=x.device, dtype=spks.dtype) - spks_in = torch.zeros([2, 80], device=x.device, dtype=spks.dtype) - cond_in = torch.zeros([2, 80, x.size(2)], device=x.device, dtype=spks.dtype) + # inference, so set dtype=spks.dtype. The batch is doubled for CFG: + # first B rows are conditioned, second B rows are unconditional. + batch_size = int(x.size(0)) + estimator_batch = 2 * batch_size + estimator_dtype = spks.dtype if spks is not None else x.dtype + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_cfg_allocate_2b"): + x_in = torch.zeros([estimator_batch, 80, x.size(2)], device=x.device, dtype=estimator_dtype) + mask_in = torch.zeros([estimator_batch, 1, x.size(2)], device=x.device, dtype=estimator_dtype) + mu_in = torch.zeros([estimator_batch, 80, x.size(2)], device=x.device, dtype=estimator_dtype) + t_in = torch.zeros([estimator_batch], device=x.device, dtype=estimator_dtype) + spks_in = torch.zeros([estimator_batch, 80], device=x.device, dtype=estimator_dtype) + cond_in = torch.zeros([estimator_batch, 80, x.size(2)], device=x.device, dtype=estimator_dtype) for step in range(1, len(t_span)): # Classifier-Free Guidance inference introduced in VoiceBox - x_in[:] = x - mask_in[:] = mask - mu_in[0] = mu - t_in[:] = t.unsqueeze(0) - spks_in[0] = spks - cond_in[0] = cond - dphi_dt = self.forward_estimator(x_in, mask_in, mu_in, t_in, spks_in, cond_in) - dphi_dt, cfg_dphi_dt = torch.split(dphi_dt, [x.size(0), x.size(0)], dim=0) - dphi_dt = (1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt - x = x + dt * dphi_dt - t = t + dt + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_cfg_prepare_2b"): + x_in[:batch_size] = x + x_in[batch_size:] = x + mask_in[:batch_size] = mask + mask_in[batch_size:] = mask + mu_in[:batch_size] = mu + t_in[:] = t + if spks is not None: + spks_in[:batch_size] = spks + if cond is not None: + cond_in[:batch_size] = cond + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_forward_estimator"): + dphi_dt = self.forward_estimator(x_in, mask_in, mu_in, t_in, spks_in, cond_in) + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_cfg_combine"): + dphi_dt, cfg_dphi_dt = torch.split(dphi_dt, [batch_size, batch_size], dim=0) + dphi_dt = (1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_euler_update"): + x = x + dt * dphi_dt + t = t + dt sol.append(x) if step < len(t_span) - 1: dt = t_span[step + 1] - t @@ -158,12 +175,12 @@ def forward_estimator(self, x, mask, mu, t, spks, cond): spks_e = spks.to(io_dtype).contiguous() cond_e = cond.to(io_dtype).contiguous() out_e = torch.empty_like(x_e) - estimator.set_input_shape("x", (2, 80, x_e.size(2))) - estimator.set_input_shape("mask", (2, 1, x_e.size(2))) - estimator.set_input_shape("mu", (2, 80, x_e.size(2))) - estimator.set_input_shape("t", (2,)) - estimator.set_input_shape("spks", (2, 80)) - estimator.set_input_shape("cond", (2, 80, x_e.size(2))) + estimator.set_input_shape("x", tuple(x_e.shape)) + estimator.set_input_shape("mask", tuple(mask_e.shape)) + estimator.set_input_shape("mu", tuple(mu_e.shape)) + estimator.set_input_shape("t", tuple(t_e.shape)) + estimator.set_input_shape("spks", tuple(spks_e.shape)) + estimator.set_input_shape("cond", tuple(cond_e.shape)) data_ptrs = [ x_e.data_ptr(), mask_e.data_ptr(), @@ -211,22 +228,25 @@ def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None, shape: (batch_size, n_feats, mel_timesteps) """ - z = ( - torch.randn( - (mu.size(0), mu.size(1), mu.size(2)), - device=mu.device, - dtype=mu.dtype, + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_noise_cache"): + z = ( + torch.randn( + (mu.size(0), mu.size(1), mu.size(2)), + device=mu.device, + dtype=mu.dtype, + ) + * temperature ) - * temperature - ) # fix prompt and overlap part mu and z - t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_t_span"): + t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) - if self.t_scheduler == "cosine": - t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) + if self.t_scheduler == "cosine": + t_span = 1 - torch.cos(t_span * 0.5 * torch.pi) - return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), None + with cosyvoice3_batch_flow_profile(f"cosyvoice3_cfm_euler_{max(1, int(n_timesteps))}_steps"): + return self.solve_euler(z, t_span=t_span, mu=mu, mask=mask, spks=spks, cond=cond), None class CausalMaskedDiffWithDiT(torch.nn.Module): @@ -299,36 +319,41 @@ def inference( finalize: bool = False, n_timesteps: int = 10, ): - assert token.shape[0] == 1 - # xvec projection - - embedding = F.normalize(embedding, dim=1) - - embedding = self.spk_embed_affine_layer(embedding) - - # concat text and prompt_text - token, token_len = torch.concat([prompt_token, token], dim=1), prompt_token_len + token_len - mask = (~make_pad_mask(token_len)).unsqueeze(-1).to(embedding) - token = self.input_embedding(torch.clamp(token, min=0)) * mask - # text encode - if finalize is True: - h = self.pre_lookahead_layer(token) - else: - h = self.pre_lookahead_layer( - token[:, : -self.pre_lookahead_len], context=token[:, -self.pre_lookahead_len :] - ) - h = h.repeat_interleave(self.token_mel_ratio, dim=1) - - mel_len1, mel_len2 = prompt_feat.shape[1], h.shape[1] - prompt_feat.shape[1] - - # get conditions - conds = torch.zeros([1, mel_len1 + mel_len2, self.output_size], device=token.device).to(h.dtype) - - conds[:, :mel_len1] = prompt_feat - - conds = conds.transpose(1, 2) + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_speaker_embedding"): + embedding = F.normalize(embedding, dim=1) + embedding = self.spk_embed_affine_layer(embedding) + + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_token_embedding_lookahead"): + # concat text and prompt_text + codec_token_len = token_len + token, total_token_len = torch.concat([prompt_token, token], dim=1), prompt_token_len + codec_token_len + mask = (~make_pad_mask(total_token_len, max_len=token.shape[1])).unsqueeze(-1).to(embedding) + token = self.input_embedding(torch.clamp(token, min=0)) * mask + # text encode + if finalize is True: + h = self.pre_lookahead_layer(token) + else: + h = self.pre_lookahead_layer( + token[:, : -self.pre_lookahead_len], context=token[:, -self.pre_lookahead_len :] + ) + + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_repeat_to_mel_axis"): + h = h.repeat_interleave(self.token_mel_ratio, dim=1) + + batch_size = int(token.shape[0]) + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_cond_prompt_mel"): + mel_len1, mel_len2 = prompt_feat.shape[1], h.shape[1] - prompt_feat.shape[1] + + # get conditions + conds = torch.zeros([batch_size, mel_len1 + mel_len2, self.output_size], device=token.device).to(h.dtype) + conds[:, :mel_len1] = prompt_feat + conds = conds.transpose(1, 2) + + lookahead = 0 if finalize else int(self.pre_lookahead_len) + valid_h_lens = torch.clamp(total_token_len.to(torch.long) - lookahead, min=0) + mel_lens = torch.clamp(valid_h_lens * int(self.token_mel_ratio), max=mel_len1 + mel_len2) + mask = (~make_pad_mask(mel_lens, max_len=mel_len1 + mel_len2)).to(h) - mask = (~make_pad_mask(torch.tensor([mel_len1 + mel_len2]))).to(h) feat, _ = self.decoder( mu=h.transpose(1, 2).contiguous(), mask=mask.unsqueeze(1), @@ -338,6 +363,7 @@ def inference( streaming=streaming, ) - feat = feat[:, :, mel_len1:] - assert feat.shape[2] == mel_len2 + with cosyvoice3_batch_flow_profile("cosyvoice3_cfm_crop_prompt_mel"): + feat = feat[:, :, mel_len1:] + assert feat.shape[2] == mel_len2 return feat.float(), None diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index 558e96b65c2..0ddc6dafa40 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -39,6 +39,10 @@ from vllm_omni.data_entry_keys import EmbeddingsStruct, OmniPayloadStruct, to_dict, to_struct from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor +from vllm_omni.model_executor.models.cosyvoice3.runtime import ( + cosyvoice3_batch_flow_debug, + cosyvoice3_batch_flow_enabled, +) from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer from vllm_omni.model_executor.models.cosyvoice3.utils import ( concat_text_with_prompt_ids, @@ -1070,12 +1074,21 @@ def forward( request_ids_list = self._split_request_ids(flat_ids, seq_token_counts) num_reqs = max(1, len(request_ids_list)) + debug_batch_flow = cosyvoice3_batch_flow_debug() + if debug_batch_flow: + logger.info( + "CosyVoice3 code2wav debug: forward num_reqs=%d seq_token_counts=%s flat_ids=%d", + num_reqs, + seq_token_counts, + int(flat_ids.numel()), + ) sample_rate = torch.tensor(int(self.config.sample_rate), dtype=torch.int32) empty_audio = torch.zeros((0,), dtype=torch.float32, device=input_ids.device) audios: list[torch.Tensor] = [empty_audio] * num_reqs srs: list[torch.Tensor] = [sample_rate] * num_reqs if not isinstance(runtime_info, list): runtime_info = [] + streaming_flow_items: list[dict[str, object]] = [] for idx, req_ids in enumerate(request_ids_list): raw = runtime_info[idx] if idx < len(runtime_info) and isinstance(runtime_info[idx], dict) else {} @@ -1142,23 +1155,21 @@ def forward( with self._stream_audio_cache_lock: cache_state = self._stream_vocoder_cache_by_req.get(req_id) - tts_speech, new_cache_state = self.code2wav.forward_streaming( - token=token.unsqueeze(0), - prompt_token=speech_token[:1], - prompt_feat=speech_feat[:1], - embedding=embedding[:1], - cache_state=cache_state, - n_timesteps=10, - token_offset_tokens=token_offset, - finalize=stream_finished, + streaming_flow_items.append( + { + "index": idx, + "req_id": req_id, + "stream_finished": stream_finished, + "token": token.unsqueeze(0), + "prompt_token": speech_token[:1], + "prompt_feat": speech_feat[:1], + "embedding": embedding[:1], + "cache_state": cache_state, + "token_offset_tokens": token_offset, + "finalize": stream_finished, + } ) - - if req_id is not None and hasattr(self, "_stream_vocoder_cache_by_req"): - with self._stream_audio_cache_lock: - if new_cache_state is None or stream_finished: - self._stream_vocoder_cache_by_req.pop(req_id, None) - else: - self._stream_vocoder_cache_by_req[req_id] = new_cache_state + continue else: token_offset = max(0, meta.talker_prefill_offset or 0) if meta else 0 tts_speech = self.code2wav.forward( @@ -1174,6 +1185,60 @@ def forward( audios[idx] = self._stitch_stream_audio(req_id, audio, stream_finished) + if streaming_flow_items: + if debug_batch_flow: + item_shapes = [ + ( + tuple(item["token"].shape), # type: ignore[union-attr] + tuple(item["prompt_token"].shape), # type: ignore[union-attr] + tuple(item["prompt_feat"].shape), # type: ignore[union-attr] + bool(item.get("finalize", False)), + ) + for item in streaming_flow_items + ] + logger.info( + "CosyVoice3 code2wav debug: streaming_items=%d item_shapes=%s", + len(streaming_flow_items), + item_shapes, + ) + if ( + len(streaming_flow_items) > 1 + and cosyvoice3_batch_flow_enabled() + and hasattr(self.code2wav, "forward_streaming_batch") + ): + streaming_results = self.code2wav.forward_streaming_batch( + streaming_flow_items, + n_timesteps=10, + ) + else: + streaming_results = [ + self.code2wav.forward_streaming( + token=item["token"], # type: ignore[arg-type] + prompt_token=item["prompt_token"], # type: ignore[arg-type] + prompt_feat=item["prompt_feat"], # type: ignore[arg-type] + embedding=item["embedding"], # type: ignore[arg-type] + cache_state=item.get("cache_state"), # type: ignore[arg-type] + n_timesteps=10, + token_offset_tokens=int(item.get("token_offset_tokens", 0)), + finalize=bool(item.get("finalize", False)), + ) + for item in streaming_flow_items + ] + + for item, (tts_speech, new_cache_state) in zip(streaming_flow_items, streaming_results): + idx = int(item["index"]) + req_id = item.get("req_id") + stream_finished = bool(item.get("stream_finished", False)) + if req_id is not None and hasattr(self, "_stream_vocoder_cache_by_req"): + with self._stream_audio_cache_lock: + if new_cache_state is None or stream_finished: + self._stream_vocoder_cache_by_req.pop(req_id, None) + else: + self._stream_vocoder_cache_by_req[req_id] = new_cache_state + + audio = tts_speech.reshape(-1).to(dtype=torch.float32) + audios[idx] = self._stitch_stream_audio(req_id, audio, stream_finished) + return OmniOutput(text_hidden_states=None, multimodal_outputs={"audio": audios, "sr": srs}) else: raise ValueError(f"Unsupported model_stage: {self.model_stage}") diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py index 49e8e515b25..ef706b6c8fc 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.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 """ CosyVoice3 Code2Wav Stage - Converts speech tokens to audio waveforms. @@ -11,6 +11,9 @@ from __future__ import annotations +from collections import Counter +from typing import cast + import numpy as np import torch import torch.nn as nn @@ -27,6 +30,10 @@ CausalHiFTGenerator, ) from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.layers import PreLookaheadLayer +from vllm_omni.model_executor.models.cosyvoice3.runtime import ( + cosyvoice3_batch_flow_debug, + cosyvoice3_batch_flow_profile, +) from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config logger = init_logger(__name__) @@ -159,6 +166,9 @@ def _forward_mel( token_offset_tokens: int = 0, streaming: bool = True, finalize: bool = False, + token_lens: torch.Tensor | None = None, + prompt_token_lens: torch.Tensor | None = None, + prompt_feat_lens: torch.Tensor | None = None, ) -> torch.Tensor: """Generate mel features via the upstream flow-model inference path.""" flow_weight = next(self.flow_model.parameters()) @@ -169,9 +179,22 @@ def _forward_mel( prompt_token = prompt_token.to(device=device, dtype=torch.int32) prompt_feat = prompt_feat.to(device=device, dtype=dtype) embedding = embedding.to(device=device, dtype=dtype) - token_len = torch.tensor([token.shape[1]], device=device, dtype=torch.int32) - prompt_token_len = torch.tensor([prompt_token.shape[1]], device=device, dtype=torch.int32) - prompt_feat_len = torch.tensor([prompt_feat.shape[1]], device=device, dtype=torch.int32) + batch_size = int(token.shape[0]) + token_len = ( + token_lens.to(device=device, dtype=torch.int32) + if token_lens is not None + else torch.full((batch_size,), token.shape[1], device=device, dtype=torch.int32) + ) + prompt_token_len = ( + prompt_token_lens.to(device=device, dtype=torch.int32) + if prompt_token_lens is not None + else torch.full((batch_size,), prompt_token.shape[1], device=device, dtype=torch.int32) + ) + prompt_feat_len = ( + prompt_feat_lens.to(device=device, dtype=torch.int32) + if prompt_feat_lens is not None + else torch.full((batch_size,), prompt_feat.shape[1], device=device, dtype=torch.int32) + ) feat, _ = self.flow_model.inference( token=token, @@ -192,37 +215,13 @@ def _forward_mel( return feat - @torch.inference_mode() - def forward_streaming( + def _stream_hift_from_feat( self, - token: torch.Tensor, - prompt_token: torch.Tensor, - prompt_feat: torch.Tensor, - embedding: torch.Tensor, + feat: torch.Tensor, *, cache_state: dict[str, torch.Tensor] | None = None, - n_timesteps: int = 10, - token_offset_tokens: int = 0, finalize: bool = False, ) -> tuple[torch.Tensor, dict[str, torch.Tensor] | None]: - """Decode streaming audio using cumulative mel + emitted-speech offset. - - This mirrors upstream CosyVoice3 streaming semantics more closely than - waveform-domain overlap-add: keep a cumulative mel history per request, - re-run causal HiFT on the history, and emit only the newly grown speech - suffix. That preserves causal look-right handling without double - trimming or duplicated overlap at chunk boundaries. - """ - feat = self._forward_mel( - token=token, - prompt_token=prompt_token, - prompt_feat=prompt_feat, - embedding=embedding, - n_timesteps=n_timesteps, - token_offset_tokens=token_offset_tokens, - streaming=True, - finalize=finalize, - ) hift_weight = self.hift.m_source.l_linear.weight chunk_mel = feat.to(device=hift_weight.device, dtype=hift_weight.dtype) @@ -257,6 +256,160 @@ def forward_streaming( } return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), new_state + @torch.inference_mode() + def forward_streaming_batch( + self, + items: list[dict[str, object]], + *, + n_timesteps: int = 10, + ) -> list[tuple[torch.Tensor, dict[str, torch.Tensor] | None]]: + """Batch the flow-matching mel path, then run HiFT per request. + + Items are grouped by prompt condition shape and finalization state. + Codec tokens may have different lengths; those are padded within the + group and passed to the flow as per-row token lengths. + """ + results: list[tuple[torch.Tensor, dict[str, torch.Tensor] | None] | None] = [None] * len(items) + groups: dict[tuple[int, int, int, bool], list[tuple[int, dict[str, object]]]] = {} + for index, item in enumerate(items): + token = item["token"] + prompt_token = item["prompt_token"] + prompt_feat = item["prompt_feat"] + embedding = item["embedding"] + assert isinstance(token, torch.Tensor) + assert isinstance(prompt_token, torch.Tensor) + assert isinstance(prompt_feat, torch.Tensor) + assert isinstance(embedding, torch.Tensor) + key = ( + int(prompt_token.shape[1]), + int(prompt_feat.shape[1]), + int(embedding.shape[1]), + bool(item.get("finalize", False)), + ) + groups.setdefault(key, []).append((index, item)) + + if cosyvoice3_batch_flow_debug(): + group_summary = {key: len(group) for key, group in groups.items()} + group_size_distribution = Counter(group_summary.values()) + batchable_items = sum(size for size in group_summary.values() if size > 1) + logger.info( + "CosyVoice3 code2wav debug: forward_streaming_batch items=%d " + "groups=%s group_size_distribution=%s batchable_items=%d", + len(items), + group_summary, + dict(sorted(group_size_distribution.items())), + batchable_items, + ) + + for _key, group in groups.items(): + if len(group) == 1: + index, item = group[0] + result = self.forward_streaming( + token=item["token"], # type: ignore[arg-type] + prompt_token=item["prompt_token"], # type: ignore[arg-type] + prompt_feat=item["prompt_feat"], # type: ignore[arg-type] + embedding=item["embedding"], # type: ignore[arg-type] + cache_state=item.get("cache_state"), # type: ignore[arg-type] + n_timesteps=n_timesteps, + token_offset_tokens=int(item.get("token_offset_tokens", 0)), + finalize=bool(item.get("finalize", False)), + ) + results[index] = result + continue + + token_tensors = [item["token"] for _, item in group] + assert all(isinstance(token, torch.Tensor) for token in token_tensors) + token_lens = torch.tensor( + [int(token.shape[1]) for token in token_tensors], # type: ignore[union-attr] + dtype=torch.int32, + ) + max_token_len = int(token_lens.max().item()) + padded_tokens = [] + for token in token_tensors: + assert isinstance(token, torch.Tensor) + if int(token.shape[1]) == max_token_len: + padded_tokens.append(token) + else: + pad = torch.zeros( + (token.shape[0], max_token_len - int(token.shape[1])), + device=token.device, + dtype=token.dtype, + ) + padded_tokens.append(torch.cat([token, pad], dim=1)) + tokens = torch.cat(padded_tokens, dim=0) + prompt_tokens = torch.cat([item["prompt_token"] for _, item in group], dim=0) # type: ignore[list-item] + prompt_feats = torch.cat([item["prompt_feat"] for _, item in group], dim=0) # type: ignore[list-item] + embeddings = torch.cat([item["embedding"] for _, item in group], dim=0) # type: ignore[list-item] + prompt_token_lens = torch.full((len(group),), prompt_tokens.shape[1], dtype=torch.int32) + prompt_feat_lens = torch.full((len(group),), prompt_feats.shape[1], dtype=torch.int32) + finalize = bool(group[0][1].get("finalize", False)) + + with cosyvoice3_batch_flow_profile(f"cosyvoice3_flow_batch_b{len(group)}_t{tokens.shape[1]}"): + feat = self._forward_mel( + token=tokens, + prompt_token=prompt_tokens, + prompt_feat=prompt_feats, + embedding=embeddings, + n_timesteps=n_timesteps, + token_offset_tokens=0, + streaming=True, + finalize=finalize, + token_lens=token_lens, + prompt_token_lens=prompt_token_lens, + prompt_feat_lens=prompt_feat_lens, + ) + + for row, (index, item) in enumerate(group): + trim_mel = max(0, int(item.get("token_offset_tokens", 0))) * int(self.token_mel_ratio) + valid_tokens = int(token_lens[row].item()) + if not finalize: + valid_tokens = max(0, valid_tokens - int(self.flow_model.pre_lookahead_len)) + valid_mel = valid_tokens * int(self.token_mel_ratio) + row_feat = feat[row : row + 1, :, :valid_mel] + if trim_mel > 0: + row_feat = row_feat[:, :, trim_mel:] + results[index] = self._stream_hift_from_feat( + row_feat, + cache_state=item.get("cache_state"), # type: ignore[arg-type] + finalize=finalize, + ) + + assert all(result is not None for result in results), "every streaming item must produce exactly one result" + return cast(list[tuple[torch.Tensor, dict[str, torch.Tensor] | None]], results) + + @torch.inference_mode() + def forward_streaming( + self, + token: torch.Tensor, + prompt_token: torch.Tensor, + prompt_feat: torch.Tensor, + embedding: torch.Tensor, + *, + cache_state: dict[str, torch.Tensor] | None = None, + n_timesteps: int = 10, + token_offset_tokens: int = 0, + finalize: bool = False, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor] | None]: + """Decode streaming audio using cumulative mel + emitted-speech offset. + + This mirrors upstream CosyVoice3 streaming semantics more closely than + waveform-domain overlap-add: keep a cumulative mel history per request, + re-run causal HiFT on the history, and emit only the newly grown speech + suffix. That preserves causal look-right handling without double + trimming or duplicated overlap at chunk boundaries. + """ + feat = self._forward_mel( + token=token, + prompt_token=prompt_token, + prompt_feat=prompt_feat, + embedding=embedding, + n_timesteps=n_timesteps, + token_offset_tokens=token_offset_tokens, + streaming=True, + finalize=finalize, + ) + return self._stream_hift_from_feat(feat, cache_state=cache_state, finalize=finalize) + @torch.inference_mode() def forward( self, diff --git a/vllm_omni/model_executor/models/cosyvoice3/runtime.py b/vllm_omni/model_executor/models/cosyvoice3/runtime.py new file mode 100644 index 00000000000..d94725148b7 --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/runtime.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Runtime toggles shared by the CosyVoice3 stages.""" + +import os +from contextlib import nullcontext + +import torch + +_FALSE_ENV_VALUES = ("0", "false", "False", "") + + +def _env_flag(name: str, default: str) -> bool: + return os.environ.get(name, default) not in _FALSE_ENV_VALUES + + +def cosyvoice3_batch_flow_enabled() -> bool: + """Return whether cross-request Stage-1 flow batching is enabled.""" + return _env_flag("COSYVOICE3_BATCH_FLOW", "0") + + +def cosyvoice3_batch_flow_debug() -> bool: + """Return whether Stage-1 batching diagnostics are enabled.""" + return _env_flag("COSYVOICE3_BATCH_FLOW_DEBUG", "0") + + +def cosyvoice3_batch_flow_profile(name: str): + """Create a profiler scope only when batching diagnostics are enabled.""" + if cosyvoice3_batch_flow_debug(): + return torch.profiler.record_function(name) + return nullcontext() From 7e520f9cc47e6884520ea8f3f9ab94ab97b34780 Mon Sep 17 00:00:00 2001 From: Anjie Hou <149605198+specture724@users.noreply.github.com> Date: Sat, 12 Sep 2026 20:58:11 +0800 Subject: [PATCH 91/93] [3/N] Encode streamed video on the worker with bounded batching (#7018) Signed-off-by: specture724 Co-authored-by: Claude Sonnet 5 Co-authored-by: Hongsheng Liu --- docs/serving/videos_api.md | 28 +++ .../minimax_h3/test_minimax_h3_contract.py | 163 ++++++++++++++++++ .../test_minimax_h3_step_execution.py | 17 +- .../test_serving_video_output_stream.py | 32 +++- .../openai_api/test_video_server.py | 54 ++++++ .../models/minimax_h3/pipeline_minimax_h3.py | 163 ++++++++++++++++-- vllm_omni/diffusion/utils/media_utils.py | 35 +++- vllm_omni/entrypoints/openai/serving_video.py | 55 ++++-- .../openai/serving_video_output_stream.py | 10 +- 9 files changed, 517 insertions(+), 40 deletions(-) diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md index 09f023d2776..c6e91df60e5 100644 --- a/docs/serving/videos_api.md +++ b/docs/serving/videos_api.md @@ -234,6 +234,34 @@ curl -X POST http://localhost:8091/v1/videos/sync \ -o output.mp4 ``` +## Output Encoding + +These `extra_params` control how the server turns decoded frames into MP4 bytes. + +| Field | Type | Default | Description | +| --- | --- | --- | --- | +| `preencode_mp4` | boolean | false | Encode the MP4 on the worker while the VAE is still decoding, instead of after the full video is materialized | +| `video_codec_options` | object | null | Encoder options passed through to the H.264 encoder, such as `{"preset": "ultrafast", "threads": "0"}` | + +With `preencode_mp4` enabled, each committed VAE chunk leaves the accelerator and +is encoded while later chunks are still decoding, so host transfer and CPU +encoding overlap the remaining decode instead of following it. The response is +unchanged: the same complete MP4, byte-for-byte equivalent frames. + +```bash +curl -X POST http://localhost:8091/v1/videos/sync \ + -F "prompt=A small robot walking through a neon city" \ + -F 'extra_params={"preencode_mp4": true, "video_codec_options": {"preset": "ultrafast"}}' \ + -o output.mp4 +``` + +`preencode_mp4` applies to the complete-MP4 response paths only. The +`/v1/realtime/video` WebSocket endpoint rejects it, because that path already +overlaps encoding through its own incremental fragmented-MP4 encoder. + +Support is per model: MiniMax-H3 implements it, and other models ignore the flag +and take the full-decode path. + ## Storage Set `VLLM_OMNI_SERVER_STORAGE__PATH` to control where asynchronous video outputs are diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 5470cd9c477..7868c255bb0 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -34,6 +34,169 @@ def _turbo_spec(filename: str) -> TurboSpec: return spec +def test_decode_to_mp4_batches_consumer_transfers(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + + class FakeEncoder: + instances = [] + + def __init__(self, **kwargs): + self.pushes = [] + self.kwargs = kwargs + self.__class__.instances.append(self) + + def push(self, frames): + self.pushes.append(np.array(frames, copy=True)) + + def finish(self): + return b"mp4" + + def abort(self): + raise AssertionError("unexpected abort") + + class FakeAudioVAE: + def decode_latent(self, latent): + return torch.zeros(1, 1, 2) + + class FakeVideoVAE: + def decode_with_chunks(self, latent, *, on_chunk): + for value in (0.0, 0.25, 0.5): + on_chunk(torch.full((1, 3, 3, 2, 2), value)) + + pipeline = object.__new__(mod.MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.audio_vae = FakeAudioVAE() + pipeline.video_vae = FakeVideoVAE() + pipeline.device = torch.device("cpu") + monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False) + monkeypatch.setattr( + "vllm_omni.diffusion.utils.media_utils.ChunkedMP4Encoder", + FakeEncoder, + ) + + output = pipeline.decode_to_mp4( + torch.zeros(1), + torch.zeros(1), + height=2, + width=2, + batch_frames=4, + ) + + assert output == b"mp4" + encoder = FakeEncoder.instances[-1] + assert [frames.shape[0] for frames in encoder.pushes] == [6, 3] + assert np.array_equal(encoder.pushes[0][0], np.full((2, 2, 3), 0, dtype=np.uint8)) + assert np.array_equal(encoder.pushes[0][3], np.full((2, 2, 3), 64, dtype=np.uint8)) + assert encoder.pushes[1].shape == (3, 2, 2, 3) + + +def test_request_video_codec_options_reach_the_preencoded_mp4_encoder(monkeypatch): + """A client's encoder options must survive the worker-side pre-encode path.""" + from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod + + class FakeEncoder: + instances = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.__class__.instances.append(self) + + def push(self, frames): + del frames + + def finish(self): + return b"mp4" + + def abort(self): + raise AssertionError("unexpected abort") + + class FakeAudioVAE: + def decode_latent(self, latent): + return torch.zeros(1, 1, 2) + + class FakeVideoVAE: + def decode_with_chunks(self, latent, *, on_chunk): + on_chunk(torch.zeros(1, 3, 1, 2, 2)) + + pipeline = object.__new__(mod.MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.audio_vae = FakeAudioVAE() + pipeline.video_vae = FakeVideoVAE() + pipeline.device = torch.device("cpu") + monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_uses_manual_component_offload", lambda self, component: False) + monkeypatch.setattr("vllm_omni.diffusion.utils.media_utils.ChunkedMP4Encoder", FakeEncoder) + + pipeline.decode_to_mp4( + torch.zeros(1), + torch.zeros(1), + height=2, + width=2, + video_codec_options={"preset": "ultrafast"}, + ) + + assert FakeEncoder.instances[-1].kwargs["video_codec_options"] == {"preset": "ultrafast"} + + +@pytest.mark.parametrize( + ("codec_extra", "expected"), + [ + ({}, {"preset": "ultrafast", "threads": "0"}), + ({"video_codec_options": {"preset": "slow", "threads": 2}}, {"preset": "slow", "threads": "2"}), + ({"video_codec_options": {}}, {}), + ({"video_codec_options": None}, None), + ], +) +@pytest.mark.parametrize("batch_extra, batch_frames", [({}, 17), ({"preencode_batch_frames": 5}, 5)]) +def test_preencode_request_preserves_serving_codec_defaults(codec_extra, expected, batch_extra, batch_frames): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.partition = "fl2va" + pipeline.supported_tasks = frozenset({"t2va"}) + pipeline.default_video_shift = 12.0 + pipeline.default_audio_shift = 3.0 + pipeline.device = torch.device("cpu") + pipeline.od_config = SimpleNamespace() + pipeline.text_encoder = object() + pipeline.encode_prompt = Mock(return_value=(torch.ones(1, 2), torch.ones(1, dtype=torch.long))) + pipeline._quality_policy = Mock() + pipeline._cache_dit_runtime = Mock() + pipeline.diffuse = Mock(return_value=(torch.zeros(1), torch.zeros(1))) + pipeline.decode_to_mp4 = Mock(return_value=b"mp4") + sampling = OmniDiffusionSamplingParams( + width=1344, + height=768, + fps=24, + num_frames=124, + num_inference_steps=50, + extra_args={"task": "t2va", "aspect_ratio": "16:9", "preencode_mp4": True, **codec_extra, **batch_extra}, + ) + batch = DiffusionRequestBatch( + [OmniDiffusionRequest(prompt="test", sampling_params=sampling, request_id="codec-defaults")] + ) + + pipeline.forward(batch) + + assert pipeline.decode_to_mp4.call_args.kwargs["video_codec_options"] == expected + assert pipeline.decode_to_mp4.call_args.kwargs["batch_frames"] == batch_frames + + +def test_video_codec_options_are_normalized_for_the_encoder(): + from vllm_omni.diffusion.utils.media_utils import normalize_video_codec_options + + assert normalize_video_codec_options({"preset": "ultrafast", "threads": 0}) == { + "preset": "ultrafast", + "threads": "0", + } + assert normalize_video_codec_options(None) is None + with pytest.raises(ValueError, match="video_codec_options"): + normalize_video_codec_options("preset=ultrafast") + + def test_h3_prepares_resolved_cache_state_immediately_before_denoise(): from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline from vllm_omni.diffusion.request import OmniDiffusionRequest diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py index 63c3a42f1ff..1b123929c35 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py @@ -284,7 +284,8 @@ def test_mixed_step_batch_leaves_gated_attention_dense(): assert recorder.denoise_timestep is None -def test_prepare_encode_seeds_runner_visible_state(monkeypatch): +@pytest.mark.parametrize("batch_frames", [1, 33]) +def test_prepare_encode_seeds_runner_visible_state(monkeypatch, batch_frames): from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=8) @@ -293,6 +294,8 @@ def test_prepare_encode_seeds_runner_visible_state(monkeypatch): context = { "height": 96, "width": 64, + "preencode_mp4": True, + "preencode_batch_frames": batch_frames, "latent_t": 2, "latent_h": 4, "latent_w": 6, @@ -335,6 +338,18 @@ def test_prepare_encode_seeds_runner_visible_state(monkeypatch): assert state.extra[mod._STEP_BRANCH] is branch assert state.extra[mod._STEP_SHAPE]["height"] == 96 + pipeline.od_config = SimpleNamespace() + monkeypatch.setattr(pipeline, "_unpack_denoised_rows", lambda *args, **kwargs: (torch.zeros(1), torch.zeros(1))) + calls = [] + + def decode_to_mp4(*args, **kwargs): + calls.append(kwargs) + return b"mp4" + + monkeypatch.setattr(pipeline, "decode_to_mp4", decode_to_mp4) + assert pipeline.post_decode(state).output == (b"mp4", None) + assert calls[0]["batch_frames"] == batch_frames + def test_prepare_encode_rejects_request_mode_only_features(): """Multi-output and DLO have no representation in the step contract.""" diff --git a/tests/entrypoints/openai_api/test_serving_video_output_stream.py b/tests/entrypoints/openai_api/test_serving_video_output_stream.py index e57e3e27959..291eada7516 100644 --- a/tests/entrypoints/openai_api/test_serving_video_output_stream.py +++ b/tests/entrypoints/openai_api/test_serving_video_output_stream.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 """Unit tests for `/v1/realtime/video` WebSocket video output streaming.""" from __future__ import annotations @@ -7,12 +7,13 @@ import asyncio import json from collections.abc import AsyncGenerator, Callable +from http import HTTPStatus from types import SimpleNamespace from unittest.mock import MagicMock import numpy as np import pytest -from fastapi import FastAPI, WebSocket +from fastapi import FastAPI, HTTPException, WebSocket from pytest_mock import MockerFixture from starlette.testclient import TestClient @@ -957,3 +958,30 @@ async def _send_interaction_after_delay() -> None: "transition_chunks": 2, }, ) + + +class TestStreamingVideoOutputExtraParams: + """extra_params handling when building streaming generation inputs.""" + + @staticmethod + def _handler() -> OmniStreamingVideoOutputHandler: + handler = object.__new__(OmniStreamingVideoOutputHandler) + handler._engine_client = SimpleNamespace(default_sampling_params_list=[]) + return handler + + async def test_preencode_mp4_is_rejected_for_streaming_sessions(self): + """Worker-side pre-encoding emits one progressive MP4 the fMP4 encoder cannot consume.""" + request = VideoGenerationRequest(prompt="p", extra_params={"preencode_mp4": True}) + + with pytest.raises(HTTPException) as excinfo: + await self._handler()._build_prompt_and_sampling_params(request) + + assert excinfo.value.status_code == HTTPStatus.BAD_REQUEST.value + assert "preencode_mp4" in excinfo.value.detail + + async def test_other_extra_params_still_reach_the_engine(self): + request = VideoGenerationRequest(prompt="p", extra_params={"flow_shift": 3.0}) + + _, gen_params, _ = await self._handler()._build_prompt_and_sampling_params(request) + + assert gen_params.extra_args["flow_shift"] == 3.0 diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index b65b00e59fb..01332beecb7 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -134,6 +134,60 @@ async def _generate_both_response_types(): handler.shutdown() +@pytest.mark.parametrize("batch_frames", [0, -1, True, 1.5, "17", None]) +def test_preencode_rejects_invalid_batch_frames_before_generation(batch_frames): + engine = FakeAsyncOmni() + handler = OmniOpenAIServingVideo.for_diffusion(engine, model_name="test-model") + request = VideoGenerationRequest( + prompt="test", extra_params={"preencode_mp4": True, "preencode_batch_frames": batch_frames} + ) + try: + with pytest.raises(HTTPException, match="preencode_batch_frames") as exc: + asyncio.run(handler.generate_video_bytes(request, "invalid-batch")) + assert exc.value.status_code == 400 + assert engine.captured_prompt is None + finally: + handler.shutdown() + + +def test_preencoded_video_bytes_preserve_metadata(mocker: MockerFixture): + from vllm_omni.entrypoints.openai.serving_video import VideoGenerationArtifacts + + handler = OmniOpenAIServingVideo.for_diffusion(FakeAsyncOmni(), model_name="test-model") + # Resolved frame count differs from anything the request asked for, so the + # metadata has to come from the encoded stream rather than request defaults. + preencoded = _make_test_video_bytes((32, 24), num_frames=7) + artifacts = VideoGenerationArtifacts( + videos=[preencoded], + audios=[None], + actions=[None], + audio_sample_rate=24000, + output_fps=24.0, + stage_durations={"decode": 0.5}, + peak_memory_mb=123.0, + metrics={"generation_time": 1.25}, + ) + mocker.patch.object(handler, "_run_and_extract", return_value=artifacts) + encoder = mocker.patch("vllm_omni.entrypoints.openai.serving_video._encode_video_bytes") + try: + result = asyncio.run(handler.generate_video_bytes(VideoGenerationRequest(prompt="test"), "preencoded")) + assert result == ( + preencoded, + {"decode": 0.5}, + 123.0, + None, + { + "fps": 24.0, + "num_frames": 7, + "duration_s": 7 / 24.0, + "metrics": {"generation_time": 1.25}, + }, + ) + encoder.assert_not_called() + finally: + handler.shutdown() + + def test_resolve_diffusion_od_config_falls_back_to_attribute(): od_config = SimpleNamespace(model_class_name="WanPipeline") handler = OmniOpenAIServingVideo.for_diffusion( diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index fb9f02b28fa..f80513d4eb6 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -64,6 +64,7 @@ DiffusionPipelineProfilerMixin, ) from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule +from vllm_omni.diffusion.utils.media_utils import normalize_preencode_batch_frames, normalize_video_codec_options from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.errors import OmniClientError, client_error_from_metadata from vllm_omni.model_executor.model_loader.weight_utils import ( @@ -367,6 +368,19 @@ def _minimax_h3_post_process(output, output_type: str = "np"): if not isinstance(output, tuple) or len(output) != 2: return output video, audio = output + if isinstance(video, (bytes, bytearray, memoryview)): + video = [video] + if isinstance(video, list) and all(isinstance(item, (bytes, bytearray, memoryview)) for item in video): + encoded_videos = [bytes(item) for item in video] + else: + encoded_videos = None + if encoded_videos is not None: + return { + "video": encoded_videos, + "audio": [None] * len(encoded_videos), + "audio_sample_rate": MINIMAX_H3_AUDIO_SAMPLE_RATE, + "fps": MINIMAX_H3_FPS, + } if video.dtype != torch.uint8 or video.ndim != 5 or video.shape[-1] not in (3, 4): # Float or channel-first frames would reach the muxer as a black or # banded video rather than as an error. @@ -2147,6 +2161,80 @@ def diffuse( audio_t=audio_t, ) + def decode_to_mp4( + self, + video_latent: torch.Tensor, + audio_latent: torch.Tensor, + *, + height: int, + width: int, + max_pending: int = 2, + batch_frames: int = 17, + video_codec_options: dict[str, str] | None = None, + ) -> bytes: + """Decode and encode one output on the worker without full-video materialization. + + Audio is decoded first so the incremental mux session can attach its audio + stream before temporal video chunks arrive. The callback receives committed + float32 ``BCTHW`` frames, performs the requested-size crop and uint8 + conversion in the worker, then applies bounded backpressure to the encoder. + """ + from vllm_omni.diffusion.utils.media_utils import ChunkedMP4Encoder + + if batch_frames <= 0: + raise ValueError("batch_frames must be positive") + + with self._component_on_device(self.audio_vae): + audio = self.audio_vae.decode_latent(audio_latent) + audio_np = audio.detach().float().cpu().numpy() + if audio_np.ndim == 3 and audio_np.shape[0] == 1: + audio_np = audio_np[0] + encoder = ChunkedMP4Encoder( + width=width, + height=height, + fps=MINIMAX_H3_FPS, + audio_waveform=audio_np, + audio_sample_rate=MINIMAX_H3_AUDIO_SAMPLE_RATE, + max_pending=max_pending, + video_codec_options=video_codec_options, + ) + + pending_chunks: list[torch.Tensor] = [] + pending_frames = 0 + + def flush_pending() -> None: + nonlocal pending_frames + if not pending_chunks: + return + batched = torch.cat(pending_chunks, dim=1) + encoder.push(batched[0].cpu().numpy()) + pending_chunks.clear() + pending_frames = 0 + + def on_chunk(frames: torch.Tensor) -> None: + nonlocal pending_frames + prepared = _prepare_minimax_h3_video_output(frames[..., :height, :width]) + if prepared.shape[0] != 1: + raise ValueError("MiniMax H3 chunked MP4 encoding currently expects one output per decoder") + pending_chunks.append(prepared) + pending_frames += int(prepared.shape[1]) + if pending_frames >= batch_frames: + flush_pending() + + try: + with self._component_on_device(self.video_vae): + with current_omni_platform.create_autocast_context( + device_type=self.device.type, + dtype=torch.float16, + enabled=True, + ): + self.video_vae.decode_with_chunks(video_latent, on_chunk=on_chunk) + flush_pending() + return encoder.finish() + except BaseException: + encoder.abort() + raise + def decode( self, video_latent: torch.Tensor, @@ -2230,6 +2318,11 @@ def _prepare_request_inputs( quality = sampling.quality logger.debug("MiniMax H3 request quality=%s", quality) extra = sampling.extra_args or {} + preencode_batch_frames = ( + normalize_preencode_batch_frames(extra.get("preencode_batch_frames", 17)) + if extra.get("preencode_mp4", False) + else 17 + ) turbo_spec = self._active_turbo_spec(sampling) has_native_lora = self._has_active_native_lora(sampling) task = self._resolve_task( @@ -2500,6 +2593,11 @@ def _prepare_request_inputs( "audio_shift": audio_shift, "base_schedule": base_schedule, "num_outputs": num_outputs, + "preencode_mp4": bool(extra.get("preencode_mp4", False)), + "preencode_batch_frames": preencode_batch_frames, + "video_codec_options": normalize_video_codec_options( + extra.get("video_codec_options", {"preset": "ultrafast", "threads": "0"}) + ), } @staticmethod @@ -2526,16 +2624,33 @@ def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput: audios = [] for output_seed in _minimax_h3_output_seeds(context["seed"], num_outputs): video_latent, audio_latent = self.diffuse(**{**denoise_kwargs, "seed": output_seed}) - video, audio = self.decode( - video_latent, - audio_latent, - height=context["height"], - width=context["width"], - ) - videos.append(_prepare_minimax_h3_video_output(video)) - audios.append(audio) - video = videos[0] if len(videos) == 1 else torch.cat(videos, dim=0) - audio = audios[0] if len(audios) == 1 else torch.cat(audios, dim=0) + if context["preencode_mp4"]: + videos.append( + self.decode_to_mp4( + video_latent, + audio_latent, + height=context["height"], + width=context["width"], + video_codec_options=context["video_codec_options"], + batch_frames=context["preencode_batch_frames"], + ) + ) + audios.append(None) + else: + video, audio = self.decode( + video_latent, + audio_latent, + height=context["height"], + width=context["width"], + ) + videos.append(_prepare_minimax_h3_video_output(video)) + audios.append(audio) + if videos and isinstance(videos[0], bytes): + video = videos[0] if len(videos) == 1 else videos + audio = None + else: + video = videos[0] if len(videos) == 1 else torch.cat(videos, dim=0) + audio = audios[0] if len(audios) == 1 else torch.cat(audios, dim=0) return DiffusionOutput( output=(video, audio), post_process_func=get_minimax_h3_post_process_func(self.od_config), @@ -2660,6 +2775,9 @@ def prepare_encode(self, state: StepRequestState, **kwargs: Any) -> StepRequestS "latent_h": context["latent_h"], "latent_w": context["latent_w"], "audio_t": context["audio_t"], + "preencode_mp4": context.get("preencode_mp4", False), + "preencode_batch_frames": context.get("preencode_batch_frames", 17), + "video_codec_options": context.get("video_codec_options"), }, } ) @@ -2831,13 +2949,24 @@ def post_decode(self, state: StepRequestState, **kwargs: Any) -> DiffusionOutput latent_w=shape["latent_w"], audio_t=shape["audio_t"], ) - video, audio = self.decode( - video_latent, - audio_latent, - height=shape["height"], - width=shape["width"], - ) - video = _prepare_minimax_h3_video_output(video) + if shape.get("preencode_mp4", False): + video = self.decode_to_mp4( + video_latent, + audio_latent, + height=shape["height"], + width=shape["width"], + video_codec_options=shape.get("video_codec_options"), + batch_frames=shape.get("preencode_batch_frames", 17), + ) + audio = None + else: + video, audio = self.decode( + video_latent, + audio_latent, + height=shape["height"], + width=shape["width"], + ) + video = _prepare_minimax_h3_video_output(video) return DiffusionOutput( output=(video, audio), post_process_func=get_minimax_h3_post_process_func(self.od_config), diff --git a/vllm_omni/diffusion/utils/media_utils.py b/vllm_omni/diffusion/utils/media_utils.py index 615df090f75..99d20a0af57 100644 --- a/vllm_omni/diffusion/utils/media_utils.py +++ b/vllm_omni/diffusion/utils/media_utils.py @@ -7,7 +7,7 @@ import io import queue import threading -from collections.abc import Iterable +from collections.abc import Iterable, Mapping from fractions import Fraction from typing import Any, cast @@ -17,6 +17,27 @@ _CHUNKED_MP4_DONE = object() +def normalize_preencode_batch_frames(value: Any) -> int: + """Validate the request's MP4 transfer/encoding batch threshold.""" + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError("preencode_batch_frames must be a positive integer") + return value + + +def normalize_video_codec_options(value: Any) -> dict[str, str] | None: + """Coerce a request's ``video_codec_options`` into PyAV's str->str contract. + + The value reaches a worker straight from ``extra_params``, so reject a + non-mapping here rather than letting it fail inside the encoder after the + decode has already run. + """ + if value is None: + return None + if not isinstance(value, Mapping): + raise ValueError("video_codec_options must be a JSON object of encoder option names to values") + return {str(key): str(item) for key, item in value.items()} + + def _validate_video_chunk(chunk: np.ndarray, *, width: int, height: int) -> None: """Validate a ``(T, H, W, 3)`` uint8 RGB chunk against a session's frame size.""" if chunk.ndim != 4 or chunk.shape[-1] != 3: @@ -233,6 +254,18 @@ def _read_new_bytes(self) -> bytes: return chunk +def count_mp4_frames(video_bytes: bytes) -> int | None: + """Read an MP4's frame count from its sample table, without decoding it.""" + if not video_bytes: + return None + try: + with cast(Any, av.open(io.BytesIO(video_bytes), format="mp4")) as container: + frames = int(container.streams.video[0].frames) + except Exception: + return None + return frames or None + + def finalize_streaming_video_bytes( video_bytes: bytes, *, diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index b3afe8b40b7..5ff055fa2b1 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -13,12 +13,14 @@ from types import SimpleNamespace from typing import TYPE_CHECKING, Any, cast +import pybase64 as base64 from fastapi import HTTPException from PIL import Image from vllm.engine.protocol import EngineClient from vllm.logger import init_logger from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata +from vllm_omni.diffusion.utils.media_utils import count_mp4_frames, normalize_preencode_batch_frames from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.openai.protocol.videos import ( VideoAction, @@ -99,7 +101,12 @@ def _video_metadata_from_artifacts(artifacts: VideoGenerationArtifacts) -> dict[ if not artifacts.videos: return metadata - num_frames = count_video_frames(artifacts.videos[0]) + video = artifacts.videos[0] + # Pre-encoded outputs arrive as MP4 bytes, which carry no tensor shape. + if isinstance(video, (bytes, bytearray, memoryview)): + num_frames = count_mp4_frames(bytes(video)) + else: + num_frames = count_video_frames(video) if num_frames is not None and num_frames > 0: metadata["num_frames"] = num_frames if artifacts.output_fps > 0: @@ -402,6 +409,11 @@ async def _run_and_extract( status_code=HTTPStatus.BAD_REQUEST.value, detail="extra_params must be a JSON object/dict.", ) + if request.extra_params.get("preencode_mp4") and "preencode_batch_frames" in request.extra_params: + try: + normalize_preencode_batch_frames(request.extra_params["preencode_batch_frames"]) + except ValueError as exc: + raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(exc)) from exc # Merge extra_params into extra_args gen_params.extra_args.update(request.extra_params) @@ -479,26 +491,22 @@ async def generate_videos( if "video_codec_options" in request.extra_params: video_codec_options = request.extra_params["video_codec_options"] + def encode_video_result(idx: int, video: Any) -> str: + if isinstance(video, bytes): + return base64.b64encode(video).decode("utf-8") + return encode_video_base64( + video, + fps=artifacts.output_fps, + audio=artifacts.audios[idx], + audio_sample_rate=artifacts.audio_sample_rate, + video_codec_options=video_codec_options, + frame_converter=self._video_frame_converter, + ) + _t_encode_start = time.perf_counter() video_data = [ VideoData( - b64_json=( - encode_video_base64( - video, - fps=artifacts.output_fps, - video_codec_options=video_codec_options, - frame_converter=self._video_frame_converter, - ) - if artifacts.audios[idx] is None - else encode_video_base64( - video, - fps=artifacts.output_fps, - audio=artifacts.audios[idx], - audio_sample_rate=artifacts.audio_sample_rate, - video_codec_options=video_codec_options, - frame_converter=self._video_frame_converter, - ) - ), + b64_json=encode_video_result(idx, video), action=artifacts.actions[idx], ) for idx, video in enumerate(artifacts.videos) @@ -549,6 +557,17 @@ async def generate_video_bytes( return b"", artifacts.stage_durations, artifacts.peak_memory_mb, action, video_metadata _t_encode_start = time.perf_counter() + if isinstance(artifacts.videos[0], bytes): + video_bytes = artifacts.videos[0] + _t_encode_ms = (time.perf_counter() - _t_encode_start) * 1000 + logger.info("Video response received pre-encoded MP4 bytes: %.2f ms", _t_encode_ms) + return ( + video_bytes, + artifacts.stage_durations, + artifacts.peak_memory_mb, + artifacts.actions[0], + video_metadata, + ) video_bytes = _encode_video_bytes( artifacts.videos[0], fps=artifacts.output_fps, diff --git a/vllm_omni/entrypoints/openai/serving_video_output_stream.py b/vllm_omni/entrypoints/openai/serving_video_output_stream.py index a9abad146fb..f8f4722ad9b 100644 --- a/vllm_omni/entrypoints/openai/serving_video_output_stream.py +++ b/vllm_omni/entrypoints/openai/serving_video_output_stream.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 """WebSocket handler for streaming generated video chunks. Protocol: @@ -556,6 +556,14 @@ async def _build_prompt_and_sampling_params( status_code=HTTPStatus.BAD_REQUEST.value, detail="extra_params must be a JSON object/dict.", ) + # Worker-side pre-encoding produces one progressive MP4, which this + # path's incremental fMP4 encoder cannot consume; streaming already + # overlaps encoding through that encoder instead. + if request.extra_params.get("preencode_mp4"): + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="preencode_mp4 is not supported for streaming video sessions.", + ) gen_params.extra_args.update(request.extra_params) self._apply_lora(request.lora, gen_params) From f7d9deb45ab56e6a2ccc1690279bd9e6bdefbfe3 Mon Sep 17 00:00:00 2001 From: KANG Qihan <3149604185@qq.com> Date: Sun, 13 Sep 2026 09:39:27 +0800 Subject: [PATCH 92/93] [Kernel][Boogu-Image] Fuse Q/K RMSNorm + interleaved RoPE via fused_qk_norm_rope (#6982) Signed-off-by: Qihan Kang --- docs/configuration/environment_variables.md | 1 + tests/config/test_environment_variables.py | 2 +- .../layers/test_fused_qk_norm_rope.py | 95 ++++- .../test_boogu_image_transformer.py | 153 +++++++- .../config/environment_variable_inventory.py | 1 + vllm_omni/diffusion/envs.py | 6 + .../diffusion/layers/fused_qk_norm_rope.py | 366 ++++++++++++++++-- .../boogu_image/boogu_image_transformer.py | 132 ++++++- 8 files changed, 703 insertions(+), 53 deletions(-) diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md index 90b9ffae092..4ddc2939cf0 100644 --- a/docs/configuration/environment_variables.md +++ b/docs/configuration/environment_variables.md @@ -112,6 +112,7 @@ settings. | Name | Type and default | Applies to and read time | Precedence and invalid values | Lifecycle | | --- | --- | --- | --- | --- | +| `VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS` | Non-negative integer; unset or empty means each consumer's own default (Boogu-Image: `2048`) | Token gate of the shared `fused_qk_norm_rope` op: a rotary table is packed for the fused kernel only when it spans at least this many positions (`B*S`); read at forward time | Environment-only. `0` = always fuse; a very large value disables the fused path. Any other value raises `ValueError` on the first forward. The crossover is host-dependent; re-benchmark before overriding. | Experimental performance control | | `VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; default false | ModelOpt NVFP4 compatibility patch; read when `vllm_omni.patch` imports | Environment-only escape hatch. Any other value means false. Set only to diagnose the upstream NaN-scale issue. | Diagnostic and temporary | | `VLLM_OMNI_USE_QUACK_FP8` | Boolean truthy spellings: `1`, `true`, `yes`, `on`; unset means hardware auto-detection | FP8 scaled matrix multiplication; evaluated when quack capability is selected | A set value overrides auto-detection. Any non-truthy value forces quack off. If quack cannot load, vLLM-Omni warns and falls back to FlashInfer. | Experimental performance control | diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index cc37cd6b155..911cfbc80b0 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -164,7 +164,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): """Make an inventory expansion an explicit review decision.""" category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values()) assert category_counts == { - EnvironmentVariableCategory.PUBLIC_OMNI: 24, + EnvironmentVariableCategory.PUBLIC_OMNI: 25, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27, EnvironmentVariableCategory.MODEL_SPECIFIC: 63, diff --git a/tests/diffusion/layers/test_fused_qk_norm_rope.py b/tests/diffusion/layers/test_fused_qk_norm_rope.py index 40888556a2d..84edce0dcf9 100644 --- a/tests/diffusion/layers/test_fused_qk_norm_rope.py +++ b/tests/diffusion/layers/test_fused_qk_norm_rope.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 pytest import torch @@ -80,3 +80,96 @@ def test_fused_qk_norm_rope_matches_bf16_reference(seq_len): torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02) torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02) + + +# --------------------------------------------------------------------------- +# General geometry (any even head_dim, here Boogu-Image's 120) in both +# pairing modes, against the module's own eager reference at the same +# tolerance as the MiniMax-H3 test above. +# --------------------------------------------------------------------------- + +_BOOGU_HEAD_DIM = 120 + + +def _boogu_inputs(strided: bool): + torch.manual_seed(11) + if strided: + # Slices of a wider head axis: the op must honour q/k strides + # (merged-QKV projections hand the op such views). + q = torch.randn(4139, 35, _BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16)[:, :28] + k = torch.randn(4139, 35, _BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16)[:, :7] + else: + q = torch.randn(4139, 28, _BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16) + k = torch.randn(4139, 7, _BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16) + q_weight = torch.randn(_BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16) + k_weight = torch.randn(_BOOGU_HEAD_DIM, device="cuda", dtype=torch.bfloat16) + freqs = torch.randn(4139, _BOOGU_HEAD_DIM // 2, device="cuda", dtype=torch.float32) + rope_table = torch.cat((torch.cos(freqs), torch.sin(freqs)), dim=-1) + return q, k, q_weight, k_weight, rope_table + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.skipif(not HAS_TRITON, reason="Triton required") +@pytest.mark.parametrize("strided", [False, True]) +def test_fused_qk_norm_rope_interleaved(strided): + from vllm_omni.diffusion.layers.fused_qk_norm_rope import ( + _eager_qk_norm_rope, + _launch_fused_qk_norm_rope, + fused_qk_norm_rope, + ) + + q, k, q_weight, k_weight, rope_table = _boogu_inputs(strided) + expected = _eager_qk_norm_rope(q, k, q_weight, k_weight, rope_table, _EPS, _BOOGU_HEAD_DIM, _BOOGU_HEAD_DIM, True) + # Check the public op AND the launcher directly: the latter cannot fall + # back to eager, so a silent dispatch regression cannot go green here. + for actual in ( + fused_qk_norm_rope(q, k, q_weight, k_weight, rope_table, _EPS, interleaved=True), + _launch_fused_qk_norm_rope(q, k, q_weight, k_weight, rope_table, _EPS, interleaved=True), + ): + torch.testing.assert_close(actual[0], expected[0], atol=0.0625, rtol=0.02) + torch.testing.assert_close(actual[1], expected[1], atol=0.0625, rtol=0.02) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.skipif(not HAS_TRITON, reason="Triton required") +def test_fused_qk_norm_rope_half_split_general_dim(): + """head_dim=120 off the MiniMax-H3 128 pin, half-split pairing. + + Exercises the generalised combined kernel's half-split mode via the + launcher. It is not routed in production (half-split traffic keeps the + untouched pre-existing per-tensor kernel and its 128/96 contract, so the + public op falls back to eager at this geometry), pending the + maintainers' call. + """ + from vllm_omni.diffusion.layers.fused_qk_norm_rope import ( + _eager_qk_norm_rope, + _launch_fused_qk_norm_rope, + ) + + q, k, q_weight, k_weight, rope_table = _boogu_inputs(strided=False) + expected = _eager_qk_norm_rope(q, k, q_weight, k_weight, rope_table, _EPS, _BOOGU_HEAD_DIM, _BOOGU_HEAD_DIM) + actual = _launch_fused_qk_norm_rope(q, k, q_weight, k_weight, rope_table, _EPS, interleaved=False) + torch.testing.assert_close(actual[0], expected[0], atol=0.0625, rtol=0.02) + torch.testing.assert_close(actual[1], expected[1], atol=0.0625, rtol=0.02) + + +def test_fused_qk_norm_rope_min_tokens_resolution(monkeypatch): + """The op-level token-gate override: unset or blank -> caller's default; a + non-negative integer string overrides it (``0`` = always fuse); anything + else is rejected naming the variable.""" + from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope_min_tokens + + env = "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS" + monkeypatch.delenv(env, raising=False) + assert fused_qk_norm_rope_min_tokens(2048) == 2048 + for blank in ("", " "): + monkeypatch.setenv(env, blank) + assert fused_qk_norm_rope_min_tokens(2048) == 2048 + monkeypatch.setenv(env, "0") + assert fused_qk_norm_rope_min_tokens(2048) == 0 + monkeypatch.setenv(env, "4096") + assert fused_qk_norm_rope_min_tokens(2048) == 4096 + for bad in ("-1", "abc"): + monkeypatch.setenv(env, bad) + with pytest.raises(ValueError, match=env): + fused_qk_norm_rope_min_tokens(2048) diff --git a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py index 43e2f12f4cf..79954f1db99 100644 --- a/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py +++ b/tests/diffusion/models/boogu_image/test_boogu_image_transformer.py @@ -26,7 +26,9 @@ def _init_distributed(request): """Initialize the minimal single-rank distributed environment required by the vLLM parallel linear layers (tensor-parallel group must exist).""" - if request.node.name.startswith("test_real_rotary_"): + if request.node.name.startswith( + ("test_real_rotary_", "test_packed_rope_table", "test_fused_qk_norm_rope", "test_fallback_with_no") + ): yield return @@ -600,3 +602,152 @@ def test_transformer_forward_ti2i_shape(): assert out.shape == (batch_size, model.out_channels, latent_h, latent_w) assert torch.isfinite(out).all() + + +# --------------------------------------------------------------------------- +# The three tests below are for the fused qk-norm+RoPE kernel's wiring into +# this transformer: the packed-table layout conversion, the fused path against +# the previous eager chain, and the fallback when no packed table is present. +# They need CUDA (the module-level marks are CPU; these add the cuda mark and +# skip on CPU-only runners) and real Boogu geometry rather than the mock dims +# above. +# --------------------------------------------------------------------------- + +_FUSED_HEAD_DIM = 120 +_FUSED_Q_HEADS = 28 +_FUSED_KV_HEADS = 7 +_FUSED_EPS = 1e-5 + + +def _fused_rotary_pair(tokens: int): + theta = torch.randn(1, tokens, _FUSED_HEAD_DIM // 2, device="cuda", dtype=torch.float32) + return ( + torch.cos(theta).repeat_interleave(2, dim=-1), + torch.sin(theta).repeat_interleave(2, dim=-1), + ) + + +def _fused_norms(): + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.model_executor.layers.layernorm import RMSNorm + + with set_current_vllm_config(VllmConfig()): + norm_q = RMSNorm(_FUSED_HEAD_DIM, eps=_FUSED_EPS).cuda().to(torch.bfloat16) + norm_k = RMSNorm(_FUSED_HEAD_DIM, eps=_FUSED_EPS).cuda().to(torch.bfloat16) + with torch.no_grad(): + norm_q.weight.copy_(torch.randn(_FUSED_HEAD_DIM)) + norm_k.weight.copy_(torch.randn(_FUSED_HEAD_DIM)) + return norm_q, norm_k + + +def _operand_ulp_bound(x, weight, cos, sin): + """One bf16 rounding of a normalized operand, propagated through RoPE.""" + import torch.nn.functional as F + + n = F.rms_norm(x, (_FUSED_HEAD_DIM,), weight, _FUSED_EPS).float().abs() + c = cos[..., ::2].float().abs().unsqueeze(2) + s = sin[..., ::2].float().abs().unsqueeze(2) + even_mag = n[..., ::2] * c + n[..., 1::2] * s + odd_mag = n[..., ::2] * s + n[..., 1::2] * c + return 2.0**-6 * torch.stack((even_mag, odd_mag), dim=-1).flatten(-2) + 1e-6 + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_packed_rope_table_layout(monkeypatch): + from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import ( + _with_packed_rope_table, + ) + + monkeypatch.delenv("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", raising=False) + torch.manual_seed(0) + tokens = 2311 # above _FUSED_MIN_TOKENS + cos, sin = _fused_rotary_pair(tokens) + out = _with_packed_rope_table((cos, sin)) + assert len(out) == 3 + assert out[0] is cos and out[1] is sin + packed = out[2] + assert packed.shape == (tokens, _FUSED_HEAD_DIM) + assert packed.dtype == torch.float32 and packed.is_contiguous() + half = _FUSED_HEAD_DIM // 2 + assert torch.equal(packed[:, :half], cos[0, :, ::2]) + assert torch.equal(packed[:, half:], sin[0, :, ::2]) + # Below _FUSED_MIN_TOKENS the tuple passes through unchanged (host-bound + # regime: the fused path would cost more than it saves). + short = _with_packed_rope_table(_fused_rotary_pair(17)) + assert len(short) == 2 + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_fused_qk_norm_rope_matches_eager_chain(): + """The fused path vs the previous eager chain at real Boogu shapes. + + Every element must sit within one bf16 ulp of a normalized operand + propagated through the rotation (the two paths legitimately differ by + single- vs double-rounding of the norm's weight multiply). + """ + from vllm.triton_utils import HAS_TRITON + + if not HAS_TRITON: + pytest.skip("Triton required") + from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import ( + _qk_norm_rope, + _with_packed_rope_table, + ) + + torch.manual_seed(0) + norm_q, norm_k = _fused_norms() + q = torch.randn(1, 4139, _FUSED_Q_HEADS, _FUSED_HEAD_DIM, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, 4139, _FUSED_KV_HEADS, _FUSED_HEAD_DIM, device="cuda", dtype=torch.bfloat16) + cos, sin = _fused_rotary_pair(4139) + with_table = _with_packed_rope_table((cos, sin)) + + fused = _qk_norm_rope(q, k, norm_q, norm_k, with_table, torch.bfloat16) + eager = _qk_norm_rope(q, k, norm_q, norm_k, (cos, sin), torch.bfloat16) + for fused_t, eager_t, x, w in zip(fused, eager, (q, k), (norm_q.weight, norm_k.weight)): + bound = _operand_ulp_bound(x, w, cos, sin) + diff = (fused_t.float() - eager_t.float()).abs() + assert (diff <= bound).all(), "fused path beyond one operand ulp of the eager chain" + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_fallback_with_no_packed_table(): + """Without a packed table the helper is bit-exact to the old chain.""" + from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import ( + _qk_norm_rope, + apply_rotary_emb, + ) + + torch.manual_seed(0) + norm_q, norm_k = _fused_norms() + q = torch.randn(1, 64, _FUSED_Q_HEADS, _FUSED_HEAD_DIM, device="cuda", dtype=torch.bfloat16) + k = torch.randn(1, 64, _FUSED_KV_HEADS, _FUSED_HEAD_DIM, device="cuda", dtype=torch.bfloat16) + cos, sin = _fused_rotary_pair(64) + + out_q, out_k = _qk_norm_rope(q, k, norm_q, norm_k, (cos, sin), torch.bfloat16) + ref_q = apply_rotary_emb(norm_q(q), (cos, sin)).to(torch.bfloat16) + ref_k = apply_rotary_emb(norm_k(k), (cos, sin)).to(torch.bfloat16) + assert torch.equal(out_q, ref_q) + assert torch.equal(out_k, ref_k) + + +def test_packed_rope_table_env_override(monkeypatch): + """VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS steers the packing gate: 0 packs + even a 43-token table, a huge value leaves a 4139-token table untouched, + unset keeps the 2048 default routing.""" + from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import ( + _with_packed_rope_table, + ) + + env = "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS" + short, long = _identity_rotary_emb(1, 43), _identity_rotary_emb(1, 4139) + + monkeypatch.setenv(env, "0") + assert len(_with_packed_rope_table(short)) == 3 + monkeypatch.setenv(env, str(10**9)) + assert _with_packed_rope_table(long) is long + monkeypatch.delenv(env) + assert len(_with_packed_rope_table(long)) == 3 + assert _with_packed_rope_table(short) is short diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index 4a4966d9c2a..fe439457de1 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -65,6 +65,7 @@ def is_public_omni(self) -> bool: "SPEAKER_SAMPLES_DIR", "VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT", "VLLM_OMNI_EVENT_DRIVEN_ORCH", + "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", "VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "VLLM_OMNI_ORCH_MONITOR_PATH", "VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY", diff --git a/vllm_omni/diffusion/envs.py b/vllm_omni/diffusion/envs.py index 32e006f7ec5..4cbc0907cbd 100644 --- a/vllm_omni/diffusion/envs.py +++ b/vllm_omni/diffusion/envs.py @@ -15,6 +15,7 @@ MASTER_PORT: int | None = None CUDA_HOME: str | None = None LOCAL_RANK: int = 0 + VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS: str | None = None environment_variables: dict[str, Callable[[], Any]] = { # ================== Runtime Env Vars ================== @@ -28,6 +29,11 @@ # local rank of the process in the distributed setting, used to determine # the GPU device id "LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")), + # Minimum rotary-table span (B*S positions) at which consumers of the shared + # fused_qk_norm_rope op take the fused path instead of their eager chain; + # "0" = always fuse; unset = each consumer's own measured default. Raw + # string or None; validated by fused_qk_norm_rope_min_tokens(). + "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS": lambda: os.environ.get("VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS", None), } diff --git a/vllm_omni/diffusion/layers/fused_qk_norm_rope.py b/vllm_omni/diffusion/layers/fused_qk_norm_rope.py index f7e27675cd1..9f66057262e 100644 --- a/vllm_omni/diffusion/layers/fused_qk_norm_rope.py +++ b/vllm_omni/diffusion/layers/fused_qk_norm_rope.py @@ -1,18 +1,26 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Q/K RMSNorm followed by packed non-interleaved RoPE. +"""Q/K RMSNorm followed by packed RoPE (non-interleaved or interleaved). The public contract is shared by diffusion attention implementations: * ``q`` and ``k`` are ``[tokens, heads, head_dim]``; * norm weights are one-dimensional ``[head_dim]`` tensors; * ``rope_table`` is ``[tokens, rotary_dim]`` and stores - ``[cos(theta), sin(theta)]`` with ``theta`` of width ``rotary_dim // 2``. + ``[cos(theta), sin(theta)]`` with ``theta`` of width ``rotary_dim // 2``; + its dtype is either the activation dtype or float32; +* ``interleaved=False`` rotates half-split pairs ``(d, d + rotary_dim/2)`` + (MiniMax-H3); ``interleaved=True`` rotates adjacent pairs ``(2i, 2i + 1)`` + with ``theta_i`` (Boogu-Image, matching its ``apply_rotary_emb``). The CUDA fast path fuses RMSNorm and RoPE without materializing normalized Q/K -or rotary-product intermediates. Ascend composes its RMSNorm and rotary fused -primitives; unsupported inputs use the eager reference. +or rotary-product intermediates. The interleaved mode supports any even +``rotary_dim <= head_dim <= 256`` (``tl.arange`` padding to the next power of +two); the half-split mode keeps the pre-existing kernel and its MiniMax-H3 +geometry contract (``head_dim == 128``, ``rotary_dim == 96``). Ascend +composes its RMSNorm and rotary fused primitives on the MiniMax-H3 geometry; +unsupported inputs use the eager reference. """ from __future__ import annotations @@ -26,31 +34,73 @@ from vllm.triton_utils import HAS_TRITON, tl, triton from vllm.utils.torch_utils import direct_register_custom_op +from vllm_omni.diffusion import envs from vllm_omni.platforms import current_omni_platform _FUSED_HEAD_DIM = 128 _FUSED_ROTARY_DIM = 96 +# Interleaved-mode geometry bounds (any even rotary_dim <= head_dim <= this). +_FUSED_MAX_HEAD_DIM = 256 _HEADS_PER_PROGRAM = 8 +# The combined interleaved kernel: small tiles + one warp measured fastest +# (latency-bound shapes; see the kernel docstring). +_INTERLEAVED_HEADS_PER_PROGRAM = 4 def _apply_rope_table( x: torch.Tensor, rope_table: torch.Tensor, rotary_dim: int, + *, + interleaved: bool = False, + dtype: torch.dtype | None = None, + output_dtype: torch.dtype | None = None, ) -> torch.Tensor: + """Eager RoPE application: the fallback and the unit-test reference. + + Shapes: ``x`` is the normalized ``[tokens, heads, head_dim]`` tensor; + ``rope_table`` is ``[tokens, rotary_dim]`` storing + ``[cos(theta) | sin(theta)]`` with ``theta`` of width ``rotary_dim // 2``. + + ``interleaved`` selects the pairing: half-split ``(d, d + rotary_dim/2)`` + sharing ``theta_d`` (MiniMax-H3) or adjacent pairs ``(2i, 2i + 1)`` + sharing ``theta_i`` (Boogu-Image). ``dtype`` is the precision the + rotation arithmetic runs in and ``output_dtype`` the precision the result + is cast to; both default to ``x``'s dtype (the historical half-split + behaviour). The interleaved caller passes ``dtype=torch.float32`` to + match the Triton kernel's composition — fp32 rotation of the + bf16-rounded normalized value, one final rounding — which is also how + vLLM's own ``fused_qk_norm_rope`` Triton kernel and Boogu's eager + ``apply_rotary_emb`` order their rounding. + """ + compute_dtype = x.dtype if dtype is None else dtype + out_dtype = x.dtype if output_dtype is None else output_dtype half = rotary_dim // 2 - cos = rope_table[..., :half].to(x.dtype).unsqueeze(1) - sin = rope_table[..., half:].to(x.dtype).unsqueeze(1) - first = x[..., :half] - second = x[..., half:rotary_dim] - return torch.cat( - ( - first * cos - second * sin, - second * cos + first * sin, - x[..., rotary_dim:], - ), - dim=-1, - ) + x_c = x.to(compute_dtype) + cos = rope_table[..., :half].to(compute_dtype).unsqueeze(1) + sin = rope_table[..., half:].to(compute_dtype).unsqueeze(1) + first = x_c[..., 0:rotary_dim:2] if interleaved else x_c[..., :half] + second = x_c[..., 1:rotary_dim:2] if interleaved else x_c[..., half:rotary_dim] + rotated_first = first * cos - second * sin + rotated_second = second * cos + first * sin + if interleaved: + # stack(..., dim=-1) pairs (new_even_i, new_odd_i) into + # [..., half, 2]; flatten(-2) lays the pairs back out as + # [e0, o0, e1, o1, ...] — back to their interleaved positions. + rotated = torch.stack( + (rotated_first, rotated_second), + dim=-1, + ).flatten(-2) + return torch.cat((rotated, x_c[..., rotary_dim:]), dim=-1).to(out_dtype) + else: + return torch.cat( + ( + rotated_first, + rotated_second, + x_c[..., rotary_dim:], + ), + dim=-1, + ).to(out_dtype) if HAS_TRITON: @@ -116,6 +166,135 @@ def _rms_norm_rope_kernel( out_offsets = token * out_stride_t + heads[:, None] * out_stride_h + dims[None, :] * out_stride_d tl.store(out_ptr + out_offsets, output, mask=mask) + @triton.jit + def _qk_norm_rope_kernel( + q_ptr, + k_ptr, + q_weight_ptr, + k_weight_ptr, + rope_table_ptr, + out_q_ptr, + out_k_ptr, + q_stride_t, + q_stride_h, + q_stride_d, + k_stride_t, + k_stride_h, + k_stride_d, + rope_stride_t, + out_q_stride_t, + out_q_stride_h, + out_q_stride_d, + out_k_stride_t, + out_k_stride_h, + out_k_stride_d, + num_q_heads: tl.constexpr, + num_kv_heads: tl.constexpr, + head_dim: tl.constexpr, + padded_dim: tl.constexpr, # next power of 2 >= head_dim (tl.arange needs a power of 2) + rotary_half: tl.constexpr, + eps: tl.constexpr, + heads_per_program: tl.constexpr, + q_head_groups: tl.constexpr, + interleaved: tl.constexpr, + ): + """Fused per-head RMSNorm + RoPE for Q and K in one launch. + + Grid axis 1 assigns the first ``q_head_groups`` programs of a token + to Q and the rest to K, so the small K grid hides inside Q's waves + instead of paying its own latency-bound launch. ``interleaved`` + selects the pairing, mirroring ``_apply_rope_table``: + + * ``True`` — adjacent pairs ``(2i, 2i+1)`` sharing ``theta_i``. + * ``False`` — half-split pairs ``(d, d + rotary_half)`` sharing + ``theta_d``. Production traffic still routes half-split inputs + through ``_rms_norm_rope_kernel``; this mode is test-covered and + awaits the maintainers' call before taking over that path. + + The modes differ only in the pair/frequency/selector index + expressions; the arithmetic is shared. Beyond the rotary width + (partial rotary or lane padding) the output is the normalized value + unchanged. + """ + token = tl.program_id(0) + head_group = tl.program_id(1) + dims = tl.arange(0, padded_dim) + dim_mask = dims < head_dim + if head_group < q_head_groups: + heads = head_group * heads_per_program + tl.arange(0, heads_per_program) + # Valid-domain mask: rows beyond the head count (when it is not a + # multiple of heads_per_program) and padded lanes beyond head_dim. + mask = (heads[:, None] < num_q_heads) & dim_mask[None, :] + in_ptr = q_ptr + weight_ptr = q_weight_ptr + out_ptr = out_q_ptr + in_stride_t, in_stride_h, in_stride_d = q_stride_t, q_stride_h, q_stride_d + out_stride_t, out_stride_h, out_stride_d = ( + out_q_stride_t, + out_q_stride_h, + out_q_stride_d, + ) + else: + heads = (head_group - q_head_groups) * heads_per_program + tl.arange(0, heads_per_program) + mask = (heads[:, None] < num_kv_heads) & dim_mask[None, :] + in_ptr = k_ptr + weight_ptr = k_weight_ptr + out_ptr = out_k_ptr + in_stride_t, in_stride_h, in_stride_d = k_stride_t, k_stride_h, k_stride_d + out_stride_t, out_stride_h, out_stride_d = ( + out_k_stride_t, + out_k_stride_h, + out_k_stride_d, + ) + + offsets = token * in_stride_t + heads[:, None] * in_stride_h + dims[None, :] * in_stride_d + x = tl.load(in_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + weight = tl.load(weight_ptr + dims, mask=dim_mask, other=0.0).to(tl.float32) + # Padding lanes load 0 and do not perturb the sum; divide by the true + # head_dim, not the padded width. + inv_rms = tl.rsqrt(tl.sum(x * x, axis=1) / head_dim + eps) + normalized = (x * inv_rms[:, None] * weight[None, :]).to(tl.bfloat16) + + # The two pairings differ only in three index expressions; the pair + # reload, the table gather, the rotation formulas and the write-back + # are shared (the same structure as _apply_rope_table). + rotary_dim = rotary_half * 2 + if interleaved: + # Adjacent pairs (2i, 2i+1) sharing theta_i. + pair_dims = tl.where(dims < rotary_dim, dims ^ 1, dims) + freq_dims = tl.where(dims < rotary_dim, dims // 2, 0) + is_first = (dims % 2) == 0 + else: + # Half-split pairs (d, d + rotary_half) sharing theta_d. + pair_dims = tl.where( + dims < rotary_half, + dims + rotary_half, + tl.where(dims < rotary_dim, dims - rotary_half, dims), + ) + freq_dims = tl.where( + dims < rotary_half, + dims, + tl.where(dims < rotary_dim, dims - rotary_half, 0), + ) + is_first = dims < rotary_half + pair_offsets = token * in_stride_t + heads[:, None] * in_stride_h + pair_dims[None, :] * in_stride_d + pair_x = tl.load(in_ptr + pair_offsets, mask=mask, other=0.0).to(tl.float32) + pair_weight = tl.load(weight_ptr + pair_dims, mask=pair_dims < head_dim, other=0.0).to(tl.float32) + pair_normalized = (pair_x * inv_rms[:, None] * pair_weight[None, :]).to(tl.bfloat16) + table_offsets = token * rope_stride_t + freq_dims + cos = tl.load(rope_table_ptr + table_offsets).to(tl.float32) + sin = tl.load(rope_table_ptr + table_offsets + rotary_half).to(tl.float32) + first = normalized.to(tl.float32) * cos - pair_normalized.to(tl.float32) * sin + second = normalized.to(tl.float32) * cos + pair_normalized.to(tl.float32) * sin + output = tl.where( + dims < rotary_dim, + tl.where(is_first, first, second), + normalized.to(tl.float32), + ) + + out_offsets = token * out_stride_t + heads[:, None] * out_stride_h + dims[None, :] * out_stride_d + tl.store(out_ptr + out_offsets, output, mask=mask) + def _eager_qk_norm_rope( q: torch.Tensor, @@ -126,12 +305,16 @@ def _eager_qk_norm_rope( eps: float, head_dim: int, rotary_dim: int, + interleaved: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: q_norm = F.rms_norm(q, (head_dim,), q_weight, eps) k_norm = F.rms_norm(k, (head_dim,), k_weight, eps) + # fp32 rotation for the interleaved (Boogu) semantics; the half-split + # (H3) reference keeps its historical x-dtype arithmetic via the defaults. + rope_dtype = torch.float32 if interleaved else None return ( - _apply_rope_table(q_norm, rope_table, rotary_dim), - _apply_rope_table(k_norm, rope_table, rotary_dim), + _apply_rope_table(q_norm, rope_table, rotary_dim, interleaved=interleaved, dtype=rope_dtype), + _apply_rope_table(k_norm, rope_table, rotary_dim, interleaved=interleaved, dtype=rope_dtype), ) @@ -206,17 +389,23 @@ def _fused_cuda_supported( k: torch.Tensor, head_dim: int, rotary_dim: int, + interleaved: bool = False, ) -> bool: - return ( + if not ( HAS_TRITON and current_platform.is_cuda() and q.is_cuda and k.is_cuda and q.dtype == torch.bfloat16 and k.dtype == torch.bfloat16 - and head_dim == _FUSED_HEAD_DIM - and rotary_dim == _FUSED_ROTARY_DIM - ) + ): + return False + if interleaved: + return rotary_dim % 2 == 0 and 2 <= rotary_dim <= head_dim <= _FUSED_MAX_HEAD_DIM + # Half-split traffic keeps the pre-existing per-tensor kernel and its + # exact MiniMax-H3 geometry contract (that kernel is untouched by the + # interleaved extension). + return head_dim == _FUSED_HEAD_DIM and rotary_dim == _FUSED_ROTARY_DIM def _fused_npu_supported( @@ -224,10 +413,12 @@ def _fused_npu_supported( k: torch.Tensor, head_dim: int, rotary_dim: int, + interleaved: bool = False, ) -> bool: """Return whether the MiniMax-H3 Ascend fused-op contract is satisfied.""" return ( - current_omni_platform.is_npu() + not interleaved + and current_omni_platform.is_npu() and q.device.type == "npu" and k.device.type == "npu" and q.dtype == torch.bfloat16 @@ -271,6 +462,69 @@ def _launch_fused_rms_norm_rope( return out +def _launch_fused_qk_norm_rope( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rope_table: torch.Tensor, + eps: float, + interleaved: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + tokens, q_heads, head_dim = q.shape + kv_heads = k.shape[1] + rotary_half = rope_table.shape[-1] // 2 + out_q = torch.empty(q.shape, dtype=q.dtype, device=q.device) + out_k = torch.empty(k.shape, dtype=k.dtype, device=k.device) + if tokens == 0: + return out_q, out_k + if interleaved: + heads_per_program = _INTERLEAVED_HEADS_PER_PROGRAM + num_warps, num_stages = 1, 2 + else: + # Match the historical per-tensor launch shape (8 heads per program, + # 8 warps) so the half-split arithmetic keeps its reduction tree. + heads_per_program = _HEADS_PER_PROGRAM + num_warps, num_stages = 8, 3 + q_head_groups = triton.cdiv(q_heads, heads_per_program) + kv_head_groups = triton.cdiv(kv_heads, heads_per_program) + grid = (tokens, q_head_groups + kv_head_groups) + _qk_norm_rope_kernel[grid]( + q, + k, + q_weight, + k_weight, + rope_table, + out_q, + out_k, + q.stride(0), + q.stride(1), + q.stride(2), + k.stride(0), + k.stride(1), + k.stride(2), + rope_table.stride(0), + out_q.stride(0), + out_q.stride(1), + out_q.stride(2), + out_k.stride(0), + out_k.stride(1), + out_k.stride(2), + num_q_heads=q_heads, + num_kv_heads=kv_heads, + head_dim=head_dim, + padded_dim=triton.next_power_of_2(head_dim), + rotary_half=rotary_half, + eps=eps, + heads_per_program=heads_per_program, + q_head_groups=q_head_groups, + interleaved=interleaved, + num_warps=num_warps, + num_stages=num_stages, + ) + return out_q, out_k + + def _fused_qk_norm_rope_impl( q: torch.Tensor, k: torch.Tensor, @@ -280,8 +534,9 @@ def _fused_qk_norm_rope_impl( eps: float, head_dim: int, rotary_dim: int, + interleaved: bool, ) -> tuple[torch.Tensor, torch.Tensor]: - if _fused_npu_supported(q, k, head_dim, rotary_dim): + if _fused_npu_supported(q, k, head_dim, rotary_dim, interleaved): return _npu_qk_norm_rope( q, k, @@ -291,7 +546,7 @@ def _fused_qk_norm_rope_impl( eps, rotary_dim, ) - if not _fused_cuda_supported(q, k, head_dim, rotary_dim): + if not _fused_cuda_supported(q, k, head_dim, rotary_dim, interleaved): return _eager_qk_norm_rope( q, k, @@ -301,7 +556,12 @@ def _fused_qk_norm_rope_impl( eps, head_dim, rotary_dim, + interleaved, ) + if interleaved: + return _launch_fused_qk_norm_rope(q, k, q_weight, k_weight, rope_table, eps, interleaved=True) + # Half-split production traffic keeps the historical per-tensor kernel; + # _qk_norm_rope_kernel's half-split mode awaits the maintainers' call. return ( _launch_fused_rms_norm_rope(q, q_weight, rope_table, eps), _launch_fused_rms_norm_rope(k, k_weight, rope_table, eps), @@ -317,8 +577,9 @@ def _fused_qk_norm_rope_fake( eps: float, head_dim: int, rotary_dim: int, + interleaved: bool, ) -> tuple[torch.Tensor, torch.Tensor]: - del q_weight, k_weight, rope_table, eps, head_dim, rotary_dim + del q_weight, k_weight, rope_table, eps, head_dim, rotary_dim, interleaved return torch.empty_like(q), torch.empty_like(k) @@ -333,6 +594,39 @@ def _fused_qk_norm_rope_fake( ) +_MIN_TOKENS_ENV = "VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS" + + +def fused_qk_norm_rope_min_tokens(default: int) -> int: + """Token span below which a consumer should keep its eager chain. + + Taking the fused path costs a fixed amount of host time per call (the + custom-op dispatch into the Python impl and the Triton launcher), while the + GPU time it saves grows with the number of tokens the call covers. Below a + host- and stack-dependent crossover a call is host-bound and the fused path + can cost more end-to-end than it saves; above it the fused path wins. + Consumers pass the default they measured on their reference hardware + (Boogu-Image: 2048 tokens on one H200) and a deployment on other hardware + overrides it once, for every consumer, through + ``VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS`` (``0`` = always fuse). + + Resolved at call time — consumers call this once per forward, not per + attention site — so tests and operators can flip it without re-importing. + Unset or blank means the caller's default; anything else that is not a + non-negative integer raises ``ValueError``. + """ + raw = envs.VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS + if raw is None or not raw.strip(): + return default + try: + value = int(raw) + except ValueError: + value = -1 + if value < 0: + raise ValueError(f"{_MIN_TOKENS_ENV} must be a non-negative integer, got {raw!r}") + return value + + def fused_qk_norm_rope( q: torch.Tensor, k: torch.Tensor, @@ -343,8 +637,9 @@ def fused_qk_norm_rope( *, head_dim: int | None = None, rotary_dim: int | None = None, + interleaved: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - """Apply Q/K RMSNorm and packed non-interleaved RoPE.""" + """Apply Q/K RMSNorm and packed RoPE (half-split or adjacent pairs).""" if q.ndim != 3 or k.ndim != 3: raise ValueError(f"q and k must be [tokens, heads, head_dim], got {q.shape} and {k.shape}") if q.shape[0] != k.shape[0] or q.shape[2] != k.shape[2]: @@ -364,15 +659,18 @@ def fused_qk_norm_rope( raise ValueError(f"Expected norm weights [{head_dim}], got {tuple(q_weight.shape)} and {tuple(k_weight.shape)}") if q_weight.device != q.device or k_weight.device != q.device: raise ValueError("Q/K norm weights must be on the activation device") - if rope_table.device != q.device or rope_table.dtype != q.dtype: - raise ValueError("rope_table must have the same dtype and device as q/k") + if rope_table.device != q.device or rope_table.dtype not in ( + q.dtype, + torch.float32, + ): + raise ValueError("rope_table must be on q/k's device with their dtype or float32") if rope_table.shape != (q.shape[0], rotary_dim): raise ValueError(f"Expected rope_table [{q.shape[0]}, {rotary_dim}], got {tuple(rope_table.shape)}") q_weight = q_weight.contiguous() k_weight = k_weight.contiguous() rope_table = rope_table.contiguous() - if _fused_npu_supported(q, k, head_dim, rotary_dim): + if _fused_npu_supported(q, k, head_dim, rotary_dim, interleaved): return _npu_qk_norm_rope( q, k, @@ -382,7 +680,7 @@ def fused_qk_norm_rope( eps, rotary_dim, ) - if not _fused_cuda_supported(q, k, head_dim, rotary_dim): + if not _fused_cuda_supported(q, k, head_dim, rotary_dim, interleaved): return _fused_qk_norm_rope_impl( q, k, @@ -392,6 +690,7 @@ def fused_qk_norm_rope( eps, head_dim, rotary_dim, + interleaved, ) return torch.ops.vllm_omni.fused_qk_norm_rope( q, @@ -402,7 +701,8 @@ def fused_qk_norm_rope( eps, head_dim, rotary_dim, + interleaved, ) -__all__ = ["fused_qk_norm_rope"] +__all__ = ["fused_qk_norm_rope", "fused_qk_norm_rope_min_tokens"] diff --git a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py b/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py index 66031ca3a17..93d52b9606e 100644 --- a/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py +++ b/vllm_omni/diffusion/models/boogu_image/boogu_image_transformer.py @@ -31,10 +31,17 @@ RowParallelLinear, ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.platforms import current_platform +from vllm.triton_utils import HAS_TRITON from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.layers.fused_qk_norm_rope import ( + _fused_cuda_supported, + fused_qk_norm_rope, + fused_qk_norm_rope_min_tokens, +) from vllm_omni.diffusion.models.utils import make_attention_mask from vllm_omni.platforms import current_omni_platform @@ -43,7 +50,10 @@ logger = init_logger(__name__) -RotaryEmbedding = tuple[torch.Tensor, torch.Tensor] +# (cos, sin) — optionally (cos, sin, packed_table) where packed_table is the +# fp32 [tokens, head_dim] = [cos(theta) | sin(theta)] layout the fused +# qk-norm+RoPE op consumes; the eager path ignores the third element. +RotaryEmbedding = tuple[torch.Tensor, ...] RotaryFrequencyTables = list[RotaryEmbedding] @@ -56,9 +66,12 @@ def apply_rotary_emb(x: torch.Tensor, rotary_emb: RotaryEmbedding) -> torch.Tens Args: x: Query or key tensor of shape [B, S, H, D]. - rotary_emb: Cosine and sine tensors, each shaped [B, S, D]. + rotary_emb: Cosine and sine tensors, each shaped [B, S, D]. The tuple + may carry a third element (the packed table for the fused + qk-norm+RoPE path); it is ignored here, hence the indexed access + instead of tuple unpacking. """ - freqs_cos, freqs_sin = rotary_emb + freqs_cos, freqs_sin = rotary_emb[0], rotary_emb[1] freqs_cos = freqs_cos.unsqueeze(2) freqs_sin = freqs_sin.unsqueeze(2) @@ -73,6 +86,87 @@ def apply_rotary_emb(x: torch.Tensor, rotary_emb: RotaryEmbedding) -> torch.Tens return x_out.type_as(x) +# The fused qk-norm+RoPE path is always on where the fast path exists; on +# other platforms the eager chain below is the behaviour-identical fallback. +_FUSED_QK_NORM_ROPE = HAS_TRITON and current_platform.is_cuda() + +# Boogu's default for the fused-path token gate: fuse only when a rotary table +# spans at least this many positions (counted as B*S). Below the crossover a +# request is host-bound and the fused path's per-call Python launch (custom-op +# impl + Triton launcher) can cost more end-to-end than the kernel saves. +# Measured on one H200: on the pre-#6871 main, 512^2 (1.1k tokens) paid +# +191 ms/request and the crossover sat at ~2,350 tokens; on the current main +# the 512^2 family measures as noise (+/-25 ms, sign-inconsistent) while +# >=768^2 (2.3k tokens) gains. The crossover is host-dependent — re-benchmark +# on your hardware and override with VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS +# (0 = always fuse); see fused_qk_norm_rope_min_tokens(). +_FUSED_MIN_TOKENS = 2048 + + +def _with_packed_rope_table(rotary_emb: RotaryEmbedding) -> RotaryEmbedding: + """Append the packed fp32 table the fused qk-norm+RoPE op consumes. + + ``freqs_cos``/``freqs_sin`` are theta-repeat-interleaved ``[B, S, D]`` + (``cos[2i] == cos[2i+1] == cos(theta_i)``), so taking every even column + yields the theta-width halves and the packed table is + ``[B*S, D] = [cos(theta) | sin(theta)]``. Tables spanning fewer positions + (counted as ``B*S``) than the token gate — Boogu's ``_FUSED_MIN_TOKENS`` + default, overridable through ``VLLM_OMNI_FUSED_QK_NORM_ROPE_MIN_TOKENS`` + (``0`` = always fuse) — are returned unchanged: without the third element + every consumer stays on the eager chain (this also keeps short tables like + the context refiner's off the fused path). + """ + freqs_cos, freqs_sin = rotary_emb[0], rotary_emb[1] + if freqs_cos.shape[0] * freqs_cos.shape[1] < fused_qk_norm_rope_min_tokens(_FUSED_MIN_TOKENS): + return rotary_emb + packed = torch.cat((freqs_cos[..., ::2], freqs_sin[..., ::2]), dim=-1) + return freqs_cos, freqs_sin, packed.reshape(-1, freqs_cos.shape[-1]).float() + + +def _qk_norm_rope( + query: torch.Tensor, + key: torch.Tensor, + norm_q, + norm_k, + rotary_emb: RotaryEmbedding | None, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + """Per-head Q/K RMSNorm + RoPE + cast, fused into one kernel launch when + a packed rope table is present; otherwise the original eager chain.""" + # Gate on the op's own fast-path predicate: when the fused CUDA kernel + # cannot run (dtype, device, or geometry outside even + # rotary_dim <= head_dim <= 256), the original chain below must be used + # bit-exactly rather than the op's internal eager fallback, whose rounding + # differs from vLLM's RMSNorm by one ulp on a fraction of elements. + if ( + rotary_emb is not None + and len(rotary_emb) > 2 + and rotary_emb[2] is not None + and _fused_cuda_supported(query, key, query.shape[-1], query.shape[-1], interleaved=True) + ): + batch, seq_len, num_heads, head_dim = query.shape + num_kv_heads = key.shape[2] + fused_q, fused_k = fused_qk_norm_rope( + query.view(batch * seq_len, num_heads, head_dim), + key.view(batch * seq_len, num_kv_heads, head_dim), + norm_q.weight, + norm_k.weight, + rotary_emb[2], + norm_q.variance_epsilon, + interleaved=True, + ) + return ( + fused_q.view(batch, seq_len, num_heads, head_dim), + fused_k.view(batch, seq_len, num_kv_heads, head_dim), + ) + query = norm_q(query) + key = norm_k(key) + if rotary_emb is not None: + query = apply_rotary_emb(query, rotary_emb) + key = apply_rotary_emb(key, rotary_emb) + return query.to(dtype), key.to(dtype) + + def swiglu(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: return F.silu(x.float(), inplace=False).to(x.dtype) * y @@ -559,13 +653,7 @@ def forward( key = key.unflatten(-1, (self.num_local_kv_heads, self.head_dim)) value = value.unflatten(-1, (self.num_local_kv_heads, self.head_dim)) - query = self.norm_q(query) - key = self.norm_k(key) - - if rotary_emb is not None: - query = apply_rotary_emb(query, rotary_emb) - key = apply_rotary_emb(key, rotary_emb) - query, key = query.to(dtype), key.to(dtype) + query, key = _qk_norm_rope(query, key, self.norm_q, self.norm_k, rotary_emb, dtype) attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None attn_output = self.attn(query, key, value, attn_metadata) @@ -711,13 +799,7 @@ def forward( key = key.view(batch_size, -1, self.num_local_kv_heads, self.head_dim) value = value.view(batch_size, -1, self.num_local_kv_heads, self.head_dim) - query = self.norm_q(query) - key = self.norm_k(key) - - if rotary_emb is not None: - query = apply_rotary_emb(query, rotary_emb) - key = apply_rotary_emb(key, rotary_emb) - query, key = query.to(dtype), key.to(dtype) + query, key = _qk_norm_rope(query, key, self.norm_q, self.norm_k, rotary_emb, dtype) attn_metadata = AttentionMetadata(attn_mask=joint_attention_mask) if joint_attention_mask is not None else None attn_output = self.attn(query, key, value, attn_metadata) @@ -1449,6 +1531,11 @@ def img_patch_embed_and_refine( shift += ref_img_len idx += 1 + if _FUSED_QK_NORM_ROPE: + # The rebuilt per-reference-image tuple needs its own packed + # table; the forward-level tuples do not flow into this batch. + batch_ref_img_rotary_emb = _with_packed_rope_table(batch_ref_img_rotary_emb) + for layer in self.ref_image_refiner: batch_ref_image_hidden_states = layer( batch_ref_image_hidden_states, batch_ref_img_mask, batch_ref_img_rotary_emb, batch_temb @@ -1534,6 +1621,17 @@ def forward( device, ) + if _FUSED_QK_NORM_ROPE: + # One packed [cos|sin] table per rotary embedding that reaches an + # attention, built once per forward; the tuples grow a third + # element that the fused path consumes and the eager path ignores. + # (ref_img_rotary_emb is packed inside img_patch_embed_and_refine, + # on the rebuilt per-reference-image batch tuple.) + context_rotary_emb = _with_packed_rope_table(context_rotary_emb) + noise_rotary_emb = _with_packed_rope_table(noise_rotary_emb) + rotary_emb = _with_packed_rope_table(rotary_emb) + combined_img_rotary_emb = _with_packed_rope_table(combined_img_rotary_emb) + # Context refinement. for layer in self.context_refiner: instruction_hidden_states = layer(instruction_hidden_states, instruction_attention_mask, context_rotary_emb) From 01a2f93256975c7ff9565c5414b0fe0225bc4765 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Sun, 13 Sep 2026 10:05:56 +0800 Subject: [PATCH 93/93] [Bugfix][Frontend] Honor output_compression on the image generations route (#7447) Signed-off-by: hsliu_ustc Signed-off-by: amy-why-3459 Co-authored-by: amy-why-3459 --- recipes/Boogu/Boogu-Image.md | 33 ++++++++++++ .../openai_api/test_image_api_utils.py | 54 +++++++++++++++++++ vllm_omni/entrypoints/openai/api_server.py | 4 +- .../entrypoints/openai/protocol/images.py | 11 ++++ 4 files changed, 101 insertions(+), 1 deletion(-) create mode 100644 tests/entrypoints/openai_api/test_image_api_utils.py diff --git a/recipes/Boogu/Boogu-Image.md b/recipes/Boogu/Boogu-Image.md index 0353b977802..9a5ce039778 100644 --- a/recipes/Boogu/Boogu-Image.md +++ b/recipes/Boogu/Boogu-Image.md @@ -452,6 +452,39 @@ curl -s http://localhost:8091/v1/chat/completions \ - **Known limitations:** the same single-GPU limitations as Base and Edit apply; CPU offload, Cache-DiT, and multi-GPU parallelism are not yet validated. +## Output format and compression + +Image requests accept `output_format` (`png`, `jpeg`, or `webp`) and +`output_compression` (0-100). The default is `png` with the fastest, +least-compressed PNG encode, which is lossless but produces multi-megabyte +payloads. For latency-sensitive clients, `output_format: "jpeg"` removes +almost the entire response-encoding overhead: on a 1024x1024 Turbo request +the client-observed e2e latency drops from ~900 ms to ~766 ms and the +payload shrinks from ~3.1 MB to ~0.8 MB. Use `output_compression` to trade +encode time against payload size when PNG must stay lossless: lower values +(e.g. 1) compress harder and produce smaller payloads at proportionally +longer encode times. + +## Attention backend + +The diffusion attention backend defaults to `FLASH_ATTN` (FlashAttention-2). +On Hopper-class GPUs (SM90) pass `--diffusion-attention-backend CUDNN_ATTN` +for a measured ~12% engine-time reduction on Boogu-Image workloads: + +```bash +vllm serve Boogu/Boogu-Image-0.1-Turbo \ + --omni \ + --diffusion-attention-backend CUDNN_ATTN \ + --port 8091 +``` + +Measured on a 143 GB SM90 card (1024x1024, 4 steps, 5-run mean): stage time +673 ms with `CUDNN_ATTN` vs 763 ms with the default. Outputs are numerically +different from FLASH_ATTN (different attention kernel) but visually +equivalent, deterministic per seed, and decode identically. On datacenter +Blackwell (SM100 / SM103) use `--diffusion-attention-backend TRTLLM_ATTN` +instead; it is rejected on SM90. + ## Performance validation The checked-in single-device A40 measurements remain the current Boogu diff --git a/tests/entrypoints/openai_api/test_image_api_utils.py b/tests/entrypoints/openai_api/test_image_api_utils.py new file mode 100644 index 00000000000..31bc36c5ed7 --- /dev/null +++ b/tests/entrypoints/openai_api/test_image_api_utils.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""L1 tests for the image response encoder's compression contract.""" + +import base64 +import io + +import pytest +from PIL import Image + +from vllm_omni.entrypoints.openai.image_api_utils import ( + encode_image_base64_with_compression, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture(scope="module") +def image(): + # Photographic-content proxy: noise-plus-gradient content compresses very + # differently across PNG compress levels, which is what the mapping under + # test controls. + import numpy as np + + rng = np.random.default_rng(0) + grad = np.linspace(0, 255, 256, dtype=np.float32) + arr = np.clip(grad[None, :, None] + rng.integers(-24, 25, (256, 256, 3)), 0, 255).astype(np.uint8) + return Image.fromarray(arr, mode="RGB") + + +def test_png_output_compression_maps_to_compress_level(image): + """output_compression 100 -> compress_level 0, 1 -> compress_level 9.""" + + def encode(compression): + buffer = io.BytesIO() + image.save(buffer, format="PNG", compress_level=max(0, min(9, 9 - compression // 11))) + return len(base64.b64encode(buffer.getvalue())) + + fast = encode_image_base64_with_compression(image, format="png", output_compression=100) + small = encode_image_base64_with_compression(image, format="png", output_compression=1) + + # Level-0 PNG barely compresses; the mapped level-9 encode is smaller. + assert len(small) < len(fast) + # The mapping must produce the same sizes as the equivalent compress_level. + assert len(fast) == encode(100) + assert len(small) == encode(1) + + +def test_jpeg_output_compression_maps_to_quality(image): + high = base64.b64decode(encode_image_base64_with_compression(image, format="jpeg", output_compression=100)) + low = base64.b64decode(encode_image_base64_with_compression(image, format="jpeg", output_compression=10)) + + assert len(high) > len(low) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 565362312ba..f4f22f95f1a 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -1605,7 +1605,9 @@ def _build_image_generation_response( output_format = _choose_output_format(request.output_format or "png", None) image_data = [ ImageData( - b64_json=encode_image_base64_with_compression(image, format=output_format), + b64_json=encode_image_base64_with_compression( + image, format=output_format, output_compression=request.output_compression + ), revised_prompt=None, ) for image in images diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py index 0215e1ce0b0..b7800c13e82 100644 --- a/vllm_omni/entrypoints/openai/protocol/images.py +++ b/vllm_omni/entrypoints/openai/protocol/images.py @@ -175,6 +175,17 @@ def validate_use_system_prompt(cls, v): default=None, description="Output image format: 'png', 'jpeg', or 'webp'. Defaults to 'png'.", ) + output_compression: int = Field( + default=100, + ge=0, + le=100, + description=( + "Compression level 0-100. For 'jpeg'/'webp' this is the encoder " + "quality. For 'png' 100 keeps the fastest, least-compressed " + "encode and lower values trade encode time for smaller payloads " + "(100 -> compress_level 0, 1 -> compress_level 9)." + ), + ) return_stage_metrics: bool | None = Field( default=None, description="Return stage metrics for benchmark clients.",