From 35a423c9bd732048c01b9c4634878aabd9dc5caa Mon Sep 17 00:00:00 2001 From: tzhouam Date: Thu, 2 Jul 2026 07:36:14 +0000 Subject: [PATCH] [BugFix][CI] Address review follow-ups from v0.24.0 rebase (#4709) Addresses Gaohan123's post-approval review comments on #4709: - pipeline_dreamzero: restore platform-dispatched synchronize() instead of torch.cuda so DreamZero keeps working on non-CUDA hardware (regression of #4399 introduced during the rebase). - omni_scheduler_mixin: drop the VLLM_OMNI_INPUT_WAIT_TIMEOUT_S env-var override; the full-payload input-wait timeout is now a plain constant. - diffusion_worker: implement _DiffusionVllmModelConfig.is_diffusion by mirroring upstream ModelConfig.is_diffusion (dLLM canvas_length check) and document why omni DiT pipelines must not enable dLLM paths. - qwen3_tts: make _embedding_dtype a class attribute on the talker and prompt-embeds builder so __new__-based test stubs inherit it; remove the (duplicated) manual assignments from the preprocess test. - .buildkite: document why the rebase raised per-step timeouts (steps were killed by the old limits, masking real failures; e.g. CustomVoice passes in ~9-15 min but cold runs were killed at 20m). Co-Authored-By: Claude Fable 5 Signed-off-by: tzhouam --- .buildkite/test-merge.yml | 8 ++++++++ .buildkite/test-nightly.yml | 8 ++++++++ .../qwen3_tts/test_qwen3_tts_talker_preprocess.py | 4 ---- vllm_omni/core/sched/omni_scheduler_mixin.py | 14 ++------------ .../models/dreamzero/pipeline_dreamzero.py | 9 +++++---- vllm_omni/diffusion/worker/diffusion_worker.py | 8 +++++++- .../models/qwen3_tts/prompt_embeds_builder.py | 6 +++++- .../models/qwen3_tts/qwen3_tts_talker.py | 6 +++++- 8 files changed, 40 insertions(+), 23 deletions(-) diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index b073bc542..c2f82519e 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -173,6 +173,8 @@ steps: - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py - vllm_omni/deploy/qwen3_tts.yaml commands: + # 40m (was 20m before the v0.24.0 rebase): cold-cache runs were killed + # at 20m, masking the real failure; healthy runs finish in ~9-15 min. - | timeout 40m bash -c " export VLLM_LOGGING_LEVEL=DEBUG @@ -200,6 +202,8 @@ steps: - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py - vllm_omni/deploy/qwen3_tts.yaml commands: + # 40m (was 20m before the v0.24.0 rebase): cold-cache runs were killed + # at 20m, masking the real failure; healthy runs finish in ~9-15 min. - | timeout 40m bash -c " export VLLM_LOGGING_LEVEL=DEBUG @@ -262,6 +266,8 @@ steps: - vllm_omni/entrypoints/openai/tts_adapters/indextts2.py - vllm_omni/config/pipeline_registry.py - vllm_omni/model_executor/models/registry.py + # 60m/45m (was 20m/20m before the v0.24.0 rebase): the 20m kill fired + # before pytest could report the underlying failure on this 4-GPU step. timeout_in_minutes: 60 commands: - | @@ -290,6 +296,8 @@ steps: - vllm_omni/model_executor/models/qwen3_omni/ - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py - vllm_omni/deploy/qwen3_omni_moe.yaml + # 50m (was 30m before the v0.24.0 rebase): multi-stage Qwen3-Omni runs + # exceeded 30m and were killed, masking the real failure. timeout_in_minutes: 50 commands: - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model' diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index 5eede7ae0..c0b455b77 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -12,6 +12,8 @@ steps: if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "omni-test" steps: - label: ":full_moon: Omni · Function Test with H100" + # 120m (was 90m before the v0.24.0 rebase): the full_model sweep grew + # past 90m and was killed, masking the real failure. timeout_in_minutes: 120 commands: - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py @@ -153,6 +155,8 @@ steps: - label: ":full_moon: Omni · Perf Test · No Async Chunk" key: nightly-omni-performance-no-async-chunk + # 300m (was 180m before the v0.24.0 rebase): the no-async-chunk + # benchmark sweep exceeded 180m and was killed mid-run. timeout_in_minutes: 300 commands: - export BENCHMARK_DIR=tests/dfx/perf/results @@ -376,6 +380,8 @@ steps: if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test" steps: - label: ":full_moon: TTS · Function Test with L4" + # 120m (was 90m before the v0.24.0 rebase): the full_model TTS sweep + # grew past 90m and was killed, masking the real failure. timeout_in_minutes: 120 commands: - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" @@ -826,6 +832,8 @@ steps: - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4" + # 120m (was 60m before the v0.24.0 rebase): the X2I function sweep + # exceeded 60m and was killed mid-run. timeout_in_minutes: 120 commands: - pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py index 7358c7893..99b409758 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py @@ -34,7 +34,6 @@ def _make_minimal_talker( """ model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration) model.talker_config = SimpleNamespace(codec_pad_id=7, num_code_groups=16) - model._embedding_dtype = torch.bfloat16 if tts_pad_embed is None: tts_pad_embed = torch.zeros((1, 4), dtype=torch.bfloat16) model._tts_pad_embed = tts_pad_embed @@ -42,7 +41,6 @@ def _make_minimal_talker( def _default_raise(**_kwargs): raise AssertionError("build_prompt_embeds was not stubbed in this test") - model._embedding_dtype = torch.bfloat16 model._prompt_builder = SimpleNamespace( build_prompt_embeds=build_prompt_embeds if build_prompt_embeds is not None else _default_raise, ) @@ -83,7 +81,6 @@ def _make_minimal_builder( builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device) builder._residual_code_embeddings = lambda: [] builder._speaker_encoder = None - builder._embedding_dtype = torch.bfloat16 builder._tts_pad_embed_buffer = ( tts_pad_embed if tts_pad_embed is not None else torch.zeros((1, 4), dtype=torch.bfloat16) ) @@ -92,7 +89,6 @@ def _make_minimal_builder( ) builder._speaker_cache = None builder._text_tokenizer = None - builder._embedding_dtype = torch.bfloat16 builder._ref_audio_artifact_cache_max_entries = 256 builder._ref_audio_artifact_cache = OrderedDict() builder._resampler_cache = OrderedDict() diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index be73db9c3..e4e50fee8 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -1,6 +1,5 @@ from __future__ import annotations -import os import time from collections.abc import Iterable from typing import Any @@ -20,21 +19,12 @@ _STATS_INTERVAL_S = 1.0 # (the state ``OmniSchedulingCoordinator`` records via ``_waiting_since``) # before the scheduler force-fails it. Defends against stuck consumer-side # requests when the producer drops a full-payload, send fails, or recv -# never arrives. Override per-deployment via -# VLLM_OMNI_INPUT_WAIT_TIMEOUT_S; set <=0 to disable the safety net. +# never arrives. A value <= 0 disables the safety net. # # Scope: this constant only covers the full-payload coordinator path # (``input_coordinator``). The async-chunk path uses # ``chunk_transfer_adapter`` and is not affected by this constant. -_INPUT_WAIT_TIMEOUT_RAW = os.environ.get("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "600") -try: - DEFAULT_INPUT_WAIT_TIMEOUT_S: float = float(_INPUT_WAIT_TIMEOUT_RAW) -except ValueError: - logger.warning( - "Invalid VLLM_OMNI_INPUT_WAIT_TIMEOUT_S=%r; falling back to 600 seconds.", - _INPUT_WAIT_TIMEOUT_RAW, - ) - DEFAULT_INPUT_WAIT_TIMEOUT_S = 600.0 +DEFAULT_INPUT_WAIT_TIMEOUT_S: float = 600.0 class OmniSchedulerMixin: diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py index 1c1902a71..65e1aac70 100644 --- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py +++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py @@ -53,6 +53,7 @@ from vllm_omni.diffusion.models.dreamzero.utils import ( ) from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import FlowUniPCMultistepScheduler from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) MAX_DREAMZERO_SESSIONS = 64 @@ -486,12 +487,12 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): latents: tuple[torch.Tensor, torch.Tensor], do_true_cfg: bool, ) -> tuple[torch.Tensor, torch.Tensor]: - """Post-step sync: .contiguous() + cuda.synchronize()""" + """Post-step sync: .contiguous() + platform synchronize.""" latents = tuple(t.contiguous() for t in latents) if do_true_cfg and get_classifier_free_guidance_world_size() > 1: - device = next((t.device for t in latents if t.is_cuda), None) - if device is not None: - torch.cuda.current_stream(device).synchronize() + # Platform-dispatched sync (not torch.cuda directly) so DreamZero + # keeps working on non-CUDA hardware (XPU/NPU, see PR #4399). + current_omni_platform.synchronize() return latents # ----------------------------------------------------------------------- diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 172f48cf1..efb280a65 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -83,7 +83,13 @@ class _DiffusionVllmModelConfig: @property def is_diffusion(self) -> bool: - return False + # Duck-types ``vllm.config.ModelConfig.is_diffusion``, which detects + # *discrete diffusion (dLLM) text models* via ``hf_config.canvas_length`` + # and gates dLLM-only scheduler/spec-decode paths in vLLM. Omni's DiT + # pipelines run in their own engine and must never enable those paths, + # so mirror the upstream check (False for every DiT model) instead of + # leaving the attribute missing. + return getattr(self.hf_config, "canvas_length", None) is not None def _make_diffusion_vllm_model_config(od_config: OmniDiffusionConfig) -> _DiffusionVllmModelConfig: 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 0d8f03294..a59a0c866 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 @@ -302,6 +302,11 @@ class Qwen3TTSPromptEmbedsBuilder: artifacts. ``0`` disables the cache. """ + # The embedding path always runs in bf16 regardless of the configured + # model dtype (mirrors ``Qwen3TTSTalkerForConditionalGeneration``). + # Class attribute so test stubs built via ``__new__`` inherit it. + _embedding_dtype = torch.bfloat16 + def __init__( self, *, @@ -329,7 +334,6 @@ class Qwen3TTSPromptEmbedsBuilder: self._tts_pad_embed_buffer = tts_pad_embed self._encode_ref_audio_batch_fn = encode_ref_audio_batch self._speaker_cache = speaker_cache - self._embedding_dtype = torch.bfloat16 self._text_tokenizer: Any | None = None 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 c1c7f4264..94c8b2b7b 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 @@ -266,6 +266,11 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): """vLLM-AR talker: step-wise layer-0 codec decoding. Predicts residual codebooks (1..Q-1) into `audio_codes` and streams text via `tailing_text_hidden`.""" + # The embedding path always runs in bf16 regardless of the configured + # model dtype (encoder weights are cast to bf16 in ``__init__``). Class + # attribute so test stubs built via ``__new__`` inherit it. + _embedding_dtype = torch.bfloat16 + hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ # Talker backbone (Qwen3 decoder-only). @@ -411,7 +416,6 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module): torch.zeros(1, int(self.talker_config.hidden_size), dtype=model_dtype), persistent=False, ) - self._embedding_dtype = torch.bfloat16 tokenizer_config = Qwen3TTSTokenizerV2Config.from_pretrained( self.model_path,