[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 <noreply@anthropic.com>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

.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'

.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

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()

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:

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
 
     # -----------------------------------------------------------------------

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:

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
 

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,