[CI] Add Qwen3-omni offline UT (#216)

Signed-off-by: Rein Yang <ruiruyang2@gmail.com>

.buildkite/pipeline.yml

@@ -50,3 +50,38 @@ steps:
             - "HF_HOME=/fsx/hf_cache"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Omni Model Test with H100"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/multi_stages_h100/
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate

docs/configuration/stage_configs/qwen2_5_omni.yaml

@@ -9,7 +9,7 @@ stage_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true  # Now we only support eager mode
       trust_remote_code: true
@@ -35,7 +35,7 @@ stage_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true
       trust_remote_code: true

examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml

@@ -14,7 +14,7 @@ stage_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
-      scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.9
       enforce_eager: true
       trust_remote_code: true

tests/multi_stages/conftest.py

@@ -101,6 +101,15 @@ class OmniRunner:
                 "generating text and speech."
             )
 
+        video_padding_token = "<|VIDEO|>"
+        image_padding_token = "<|IMAGE|>"
+        audio_padding_token = "<|AUDIO|>"
+
+        if self.model_name == "Qwen/Qwen3-Omni-30B-A3B-Instruct":
+            video_padding_token = "<|video_pad|>"
+            image_padding_token = "<|image_pad|>"
+            audio_padding_token = "<|audio_pad|>"
+
         if isinstance(prompts, str):
             prompts = [prompts]
 
@@ -129,30 +138,30 @@ class OmniRunner:
             if audio is not None:
                 if isinstance(audio, list):
                     for _ in audio:
-                        user_content += "<|audio_bos|><|AUDIO|><|audio_eos|>"
+                        user_content += f"<|audio_bos|>{audio_padding_token}<|audio_eos|>"
                     multi_modal_data["audio"] = audio
                 else:
-                    user_content += "<|audio_bos|><|AUDIO|><|audio_eos|>"
+                    user_content += f"<|audio_bos|>{audio_padding_token}<|audio_eos|>"
                     multi_modal_data["audio"] = audio
 
             image = images_list[i]
             if image is not None:
                 if isinstance(image, list):
                     for _ in image:
-                        user_content += "<|vision_bos|><|IMAGE|><|vision_eos|>"
+                        user_content += f"<|vision_bos|>{image_padding_token}<|vision_eos|>"
                     multi_modal_data["image"] = image
                 else:
-                    user_content += "<|vision_bos|><|IMAGE|><|vision_eos|>"
+                    user_content += f"<|vision_bos|>{image_padding_token}<|vision_eos|>"
                     multi_modal_data["image"] = image
 
             video = videos_list[i]
             if video is not None:
                 if isinstance(video, list):
                     for _ in video:
-                        user_content += "<|vision_bos|><|VIDEO|><|vision_eos|>"
+                        user_content += f"<|vision_bos|>{video_padding_token}<|vision_eos|>"
                     multi_modal_data["video"] = video
                 else:
-                    user_content += "<|vision_bos|><|VIDEO|><|vision_eos|>"
+                    user_content += f"<|vision_bos|>{video_padding_token}<|vision_eos|>"
                     multi_modal_data["video"] = video
 
             user_content += prompt_text
@@ -230,6 +239,87 @@ class OmniRunner:
         )
         return self.generate(omni_inputs, sampling_params_list)
 
+    def generate_audio(
+        self,
+        prompts: list[str] | str,
+        sampling_params_list: list[SamplingParams] | None = None,
+        system_prompt: str | None = None,
+        audios: PromptAudioInput = None,
+        mm_processor_kwargs: dict[str, Any] | None = None,
+    ) -> list[Any]:
+        """
+        Convenience method to generate with multimodal inputs.
+        Args:
+            prompts: Text prompt(s)
+            sampling_params_list: List of sampling parameters for each stage
+            system_prompt: Optional system prompt
+            audios: Audio input(s)
+            mm_processor_kwargs: Optional processor kwargs
+        Returns:
+            List of OmniRequestOutput objects from stages with final_output=True
+        """
+        omni_inputs = self.get_omni_inputs(
+            prompts=prompts,
+            system_prompt=system_prompt,
+            audios=audios,
+            mm_processor_kwargs=mm_processor_kwargs,
+        )
+        return self.generate(omni_inputs, sampling_params_list)
+
+    def generate_video(
+        self,
+        prompts: list[str] | str,
+        sampling_params_list: list[SamplingParams] | None = None,
+        system_prompt: str | None = None,
+        videos: PromptVideoInput = None,
+        mm_processor_kwargs: dict[str, Any] | None = None,
+    ) -> list[Any]:
+        """
+        Convenience method to generate with multimodal inputs.
+        Args:
+            prompts: Text prompt(s)
+            sampling_params_list: List of sampling parameters for each stage
+            system_prompt: Optional system prompt
+            videos: Video input(s)
+            mm_processor_kwargs: Optional processor kwargs
+        Returns:
+            List of OmniRequestOutput objects from stages with final_output=True
+        """
+        omni_inputs = self.get_omni_inputs(
+            prompts=prompts,
+            system_prompt=system_prompt,
+            videos=videos,
+            mm_processor_kwargs=mm_processor_kwargs,
+        )
+        return self.generate(omni_inputs, sampling_params_list)
+
+    def generate_image(
+        self,
+        prompts: list[str] | str,
+        sampling_params_list: list[SamplingParams] | None = None,
+        system_prompt: str | None = None,
+        images: PromptImageInput = None,
+        mm_processor_kwargs: dict[str, Any] | None = None,
+    ) -> list[Any]:
+        """
+        Convenience method to generate with multimodal inputs.
+        Args:
+            prompts: Text prompt(s)
+            sampling_params_list: List of sampling parameters for each stage
+            system_prompt: Optional system prompt
+            images: Image input(s)
+            mm_processor_kwargs: Optional processor kwargs
+        Returns:
+            List of OmniRequestOutput objects from stages with final_output=True
+        """
+        omni_inputs = self.get_omni_inputs(
+            prompts=prompts,
+            system_prompt=system_prompt,
+            images=images,
+            mm_processor_kwargs=mm_processor_kwargs,
+        )
+        return self.generate(omni_inputs, sampling_params_list)
+
     def __enter__(self):
         """Context manager entry."""
         return self

tests/multi_stages_h100/conftest.py

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Pytest configuration for multi_stages_h100 tests.
+
+This conftest imports fixtures from the parent multi_stages conftest.
+"""
+
+# Import the omni_runner fixture from the parent directory's conftest
+# This makes the fixture available to tests in this directory
+from ..multi_stages.conftest import omni_runner  # noqa: F401
+
+# Explicitly expose the fixture for pytest discovery
+__all__ = ["omni_runner"]

tests/multi_stages_h100/stage_configs/qwen3_omni_ci.yaml

@@ -0,0 +1,98 @@
+# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
+# Stage 0: Thinker (multimodal understanding + text generation)
+# Stage 1: Talker (text embeddings → 8-layer RVQ codec codes)
+# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
+
+# The following config has been verified on 2x H100-80G GPUs.
+stage_args:
+  - stage_id: 0
+    runtime:
+      devices: "0,1"
+      max_batch_size: 1
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.6
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent  # Output hidden states for talker
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      hf_config_name: thinker_config
+      tensor_parallel_size: 2
+      load_format: dummy
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 100
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    runtime:
+       devices: "1"
+       max_batch_size: 1
+    engine_args:
+       model_stage: talker
+       model_arch: Qwen3OmniMoeForConditionalGeneration
+       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+       gpu_memory_utilization: 0.5
+       enforce_eager: true
+       trust_remote_code: true
+       engine_output_type: latent  # Output codec codes for code2wav
+      #  tensor_parallel_size: 2
+       enable_prefix_caching: false
+       distributed_executor_backend: "mp"
+       hf_config_name: talker_config
+       load_format: dummy
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
+    # final_output: true
+    # final_output_type: text
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 100
+      seed: 42
+      detokenize: False
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+
+  - stage_id: 2
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: audio  # Final output: audio waveform
+      gpu_memory_utilization: 0.1
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 1000000
+      hf_config_name: thinker_config
+      load_format: dummy
+    engine_input_source: [1]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
+    final_output: true
+    final_output_type: audio
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 200
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1

tests/multi_stages_h100/test_qwen3_omni.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E tests for Qwen3-Omni model with audio/video/image input and audio output.
+"""
+
+import os
+from pathlib import Path
+
+import pytest
+from vllm.assets.video import VideoAsset
+
+from ..multi_stages.conftest import OmniRunner
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
+
+# CI stage config for 2xH100-80G GPUs
+CI_STAGE_CONFIG_PATH = str(Path(__file__).parent / "stage_configs" / "qwen3_omni_ci.yaml")
+
+
+@pytest.mark.parametrize("model", models)
+def test_video_to_audio(omni_runner: type[OmniRunner], model: str) -> None:
+    """Test processing video, generating audio output."""
+    with omni_runner(model, seed=42, stage_configs_path=CI_STAGE_CONFIG_PATH) as runner:
+        # Prepare inputs
+        question = "Describe the video briefly."
+        video = VideoAsset(name="baby_reading", num_frames=4).np_ndarrays
+
+        outputs = runner.generate_multimodal(
+            prompts=question,
+            videos=video,
+        )
+
+        # Verify we got outputs from multiple stages
+        assert len(outputs) > 0
+
+        # Find and verify text output (thinker stage)
+        text_output = None
+        for stage_output in outputs:
+            if stage_output.final_output_type == "text":
+                text_output = stage_output
+                break
+
+        assert text_output is not None
+        assert len(text_output.request_output) > 0
+        text_content = text_output.request_output[0].outputs[0].text
+        assert text_content is not None
+        assert len(text_content.strip()) > 0
+
+        # Find and verify audio output (code2wav stage)
+        audio_output = None
+        for stage_output in outputs:
+            if stage_output.final_output_type == "audio":
+                audio_output = stage_output
+                break
+
+        assert audio_output is not None
+        assert len(audio_output.request_output) > 0
+
+        # Verify audio tensor exists and has content
+        audio_tensor = audio_output.request_output[0].multimodal_output["audio"]
+        assert audio_tensor is not None
+        assert audio_tensor.numel() > 0

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -156,6 +156,9 @@ class Qwen3OmniMoeForConditionalGeneration(
                 ]
             )
 
+            # for CI: Initialize special tokens embeddings early to avoid AttributeError when loading dummy weights
+            self._init_special_tokens_embeddings()
+
         elif self.model_stage == "code2wav":
             self.thinker = None
             self.talker = None