Merge branch 'main' into fix/pr4709-review-followup

Only post-#4709 nightlies (#2474/#2486, both on rebased code) are a valid
baseline; the rebase changed step runtimes substantially. Revert the raises
that post-merge data no longer supports and keep the ones it confirms:

- Omni Function H100: keep 120m (post-merge 59-71 min vs old 90m ceiling)
- Omni Perf No Async Chunk: 300m -> 180m (post-merge 41-44 min)
- TTS Function L4: 120m -> 90m (post-merge ~27 min)
- Diffusion X2I L4: 120m -> 60m (post-merge ~18 min)
- IndexTTS2: 60m/45m -> 40m/35m (steady ~12.5 min warm, cold kill at 20.7m)
- CustomVoice/Base 40m and Qwen3-Omni 50m: keep (9-21 min warm + cold-cache
  headroom the old limits lacked)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

.buildkite/test-merge.yml

@@ -173,6 +173,9 @@ steps:
           - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
           - vllm_omni/deploy/qwen3_tts.yaml
         commands:
+          # 40m (was 20m): post-rebase runs take 9-12 min warm, but cold-cache
+          # runs exceed 20m and were killed without failure output (a cold run
+          # died at 25.2m wall under the old limit).
           - |
             timeout 40m bash -c "
               export VLLM_LOGGING_LEVEL=DEBUG
@@ -200,6 +203,9 @@ steps:
           - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
           - vllm_omni/deploy/qwen3_tts.yaml
         commands:
+          # 40m (was 20m): post-rebase runs take 9-13 min warm, but cold-cache
+          # runs exceed 20m and were killed without failure output (a cold run
+          # died at 25.2m wall under the old limit).
           - |
             timeout 40m bash -c "
               export VLLM_LOGGING_LEVEL=DEBUG
@@ -262,10 +268,13 @@ steps:
           - vllm_omni/entrypoints/openai/tts_adapters/indextts2.py
           - vllm_omni/config/pipeline_registry.py
           - vllm_omni/model_executor/models/registry.py
-        timeout_in_minutes: 60
+        # 40m/35m (was 20m/20m): post-rebase runs take a steady ~12.5 min warm,
+        # but a cold run was killed at 20.7m under the old limit with no
+        # failure output; the rebase-time raise to 60m/45m was oversized.
+        timeout_in_minutes: 40
         commands:
           - |
-            timeout 45m bash -c "
+            timeout 35m bash -c "
               export VLLM_LOGGING_LEVEL=DEBUG
               pytest -s -v tests/e2e/online_serving/test_indextts2.py tests/e2e/offline_inference/test_indextts2.py -m 'advanced_model and cuda' --run-level 'advanced_model'
             "
@@ -290,6 +299,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): post-rebase runs take 15-21 min warm; 30m left no
+        # headroom for cold-cache runs, which were killed without output.
         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,9 @@ 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): post-rebase nightlies (#2474/#2486) run 59-71 min —
+        # too close to the old 90m ceiling, which #2444 already hit (killed at
+        # 90.8m with no failure output).
         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,7 +156,10 @@ steps:
 
       - label: ":full_moon: Omni · Perf Test · No Async Chunk"
         key: nightly-omni-performance-no-async-chunk
-        timeout_in_minutes: 300
+        # 180m: post-rebase nightlies (#2474/#2486) run 41-44 min; the
+        # rebase-time raise to 300m targeted a pre-rebase 179.4m outlier and
+        # is no longer needed.
+        timeout_in_minutes: 180
         commands:
           - export BENCHMARK_DIR=tests/dfx/perf/results
           - |
@@ -376,7 +382,10 @@ 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"
-        timeout_in_minutes: 120
+        # 90m: post-rebase nightlies (#2474/#2486) run ~27 min; the rebase-time
+        # raise to 120m targeted pre-rebase 72-87 min runs and is no longer
+        # needed.
+        timeout_in_minutes: 90
         commands:
           - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
           - export VLLM_USE_DEEP_GEMM="0"
@@ -861,7 +870,10 @@ steps:
 
 
       - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
-        timeout_in_minutes: 120
+        # 60m: post-rebase nightlies (#2474/#2486) run ~18 min; the rebase-time
+        # raise to 120m targeted pre-rebase 48-58 min runs and is no longer
+        # needed.
+        timeout_in_minutes: 60
         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
         agents:

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
@@ -610,12 +611,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).
@@ -443,7 +448,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,