[CI] Add online e2e test for MIMO-Audio (#2129)

Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com>

.buildkite/test-merge.yml

@@ -216,6 +216,7 @@ steps:
       - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
       - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
       - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
+      - pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
     agents:
       queue: "mithril-h100-pool"
     plugins:

.buildkite/test-ready.yml

@@ -255,6 +255,49 @@ steps:
   #                 path: /mnt/hf-cache
   #                 type: DirectoryOrCreate
 
+  - label: "MiMo-Audio E2E Test with H100"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 30m bash -c '
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model"
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            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
+
   - label: "Qwen3-TTS E2E Test"
     depends_on: upload-ready-pipeline
     commands:

docs/contributing/ci/CI_5levels.md

@@ -229,6 +229,7 @@ vllm_omni/                                    tests/
                                                │   ├── test_async_omni.py
                                                │   ├── test_qwen3_omni.py
                                                │   ├── test_qwen3_omni_expansion.py
+                                               │   ├── test_mimo_audio.py
                                                │   ├── test_image_gen_edit.py
                                                │   ├── test_images_generations_lora.py
                                                │   └── stage_configs/

docs/contributing/ci/tests_style.md

@@ -134,6 +134,7 @@ vllm_omni/                                    tests/
                                                │   ├── test_async_omni.py
                                                │   ├── test_qwen3_omni.py
                                                │   ├── test_qwen3_omni_expansion.py
+                                               │   ├── test_mimo_audio.py
                                                │   ├── test_image_gen_edit.py
                                                │   ├── test_images_generations_lora.py
                                                │   └── stage_configs/

tests/conftest.py

@@ -1495,15 +1495,15 @@ def assert_omni_response(response: OmniResponse, request_config: dict[str, Any],
 
     modalities = request_config.get("modalities", ["text", "audio"])
 
-    if "audio" in modalities:
-        assert response.audio_content is not None, "No audio output is generated"
-        print(f"audio content is: {response.audio_content}")
-
-    if "text" in modalities:
-        assert response.text_content is not None, "No text output is generated"
-        print(f"text content is: {response.text_content}")
-
     if run_level == "advanced_model":
+        if "audio" in modalities:
+            assert response.audio_content is not None, "No audio output is generated"
+            print(f"audio content is: {response.audio_content}")
+
+        if "text" in modalities:
+            assert response.text_content is not None, "No text output is generated"
+            print(f"text content is: {response.text_content}")
+
         # Verify image description
         word_types = ["text", "image", "audio", "video"]
         keywords_dict = request_config.get("key_words", {})

tests/e2e/online_serving/test_mimo_audio.py

@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E Online tests for MiMo-Audio model with audio/text input and audio output.
+"""
+
+import os
+from pathlib import Path
+
+import pytest
+
+from tests.conftest import (
+    OmniServerParams,
+    dummy_messages_from_mix_data,
+    generate_synthetic_audio,
+    modify_stage_config,
+)
+from tests.utils import hardware_test
+from vllm_omni.model_executor.model_loader.weight_utils import (
+    download_weights_from_hf_specific,
+)
+
+MIMO_AUDIO_TOKENIZER_REPO = "XiaomiMiMo/MiMo-Audio-Tokenizer"
+CHAT_TEMPLATE_PATH = str(
+    Path(__file__).parent.parent.parent.parent / "examples" / "online_serving" / "mimo_audio" / "chat_template.jinja"
+)
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+models = ["XiaomiMiMo/MiMo-Audio-7B-Instruct"]
+
+
+def get_chunk_config():
+    path = modify_stage_config(
+        str(Path(__file__).parent.parent / "stage_configs" / "mimo_audio_ci.yaml"),
+        updates={
+            "async_chunk": True,
+            "stage_args": {
+                0: {
+                    "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav_async_chunk"
+                },
+                1: {
+                    "engine_args.max_model_len": 8192,
+                    "engine_args.max_num_batched_tokens": 8192,
+                },
+            },
+        },
+        deletes={"stage_args": {1: ["custom_process_input_func"]}},
+    )
+    return path
+
+
+def download_tokenizer():
+    tokenizer_path = os.environ.get("MIMO_AUDIO_TOKENIZER_PATH", MIMO_AUDIO_TOKENIZER_REPO)
+    if os.path.exists(tokenizer_path):
+        return tokenizer_path
+    local_path = download_weights_from_hf_specific(
+        model_name_or_path=MIMO_AUDIO_TOKENIZER_REPO,
+        cache_dir=None,
+        allow_patterns=["*"],
+        require_all=True,
+    )
+    return local_path
+
+
+# CI stage config for H100 / MI325
+stage_configs = [get_chunk_config()]
+tokenizer_path = download_tokenizer()
+os.environ["MIMO_AUDIO_TOKENIZER_PATH"] = tokenizer_path
+
+# Create parameter combinations for model and stage config
+test_params = [
+    OmniServerParams(
+        model=model,
+        stage_config_path=stage_config,
+        server_args=["--chat-template", CHAT_TEMPLATE_PATH],
+    )
+    for model in models
+    for stage_config in stage_configs
+]
+
+
+def get_prompt(prompt_type="text_only"):
+    prompts = {
+        "text_only": "What is the capital of China? Answer in 20 words.",
+        "audio": "What is recited in the audio?",
+    }
+    return prompts.get(prompt_type, prompts["text_only"])
+
+
+def get_max_batch_size(size_type="few"):
+    batch_sizes = {"few": 5, "medium": 100, "large": 256}
+    return batch_sizes.get(size_type, 5)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_audio_to_text_audio_001(omni_server, openai_client) -> None:
+    """
+    Test audio and text input processing and text/audio output generation via OpenAI API.
+    Deploy Setting: default yaml
+    Input Modal: text + audio
+    Output Modal: text + audio
+    Input Setting: stream=True
+    Datasets: single request
+    """
+
+    audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1, sample_rate=24000)['base64']}"
+    messages = dummy_messages_from_mix_data(
+        audio_data_url=audio_data_url,
+        content_text=get_prompt("audio"),
+    )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": True,
+        "key_words": {
+            "audio": ["test"],
+        },
+    }
+
+    # Test single completion
+    openai_client.send_omni_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_to_text_001(omni_server, openai_client) -> None:
+    """
+    Test text input processing and text-only output generation via OpenAI API.
+    Deploy Setting: default yaml
+    Input Modal: text
+    Output Modal: text
+    Datasets: few requests
+    """
+    messages = dummy_messages_from_mix_data(content_text=get_prompt())
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": False,
+        "modalities": ["text"],
+        "key_words": {"text": ["beijing"]},
+    }
+
+    openai_client.send_omni_request(request_config, request_num=get_max_batch_size())

tests/e2e/stage_configs/mimo_audio_ci.yaml

@@ -0,0 +1,71 @@
+# CI stage config for running MiMo-Audio in multi-stage omni runtime.
+# Based on mimo_audio.yaml with load_format: dummy for CI testing.
+
+async_chunk: false
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      process: true
+      devices: "0"
+    engine_args:
+      dtype: bfloat16
+      max_num_seqs: 1
+      model_stage: fused_thinker_talker
+      model_arch: MiMoAudioForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      tensor_parallel_size: 1
+      gpu_memory_utilization: 0.5
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: latent
+      max_model_len: 8192
+      max_num_batched_tokens: 8192
+      load_format: dummy
+    is_comprehension: true
+    final_output: true
+    final_output_type: text
+    default_sampling_params:
+      temperature: 0.6
+      top_p: 0.95
+      top_k: 50
+      max_tokens: 18192
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      process: true
+      devices: "0"
+    engine_args:
+      model_stage: code2wav
+      max_num_seqs: 1
+      model_arch: MiMoAudioForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      tensor_parallel_size: 1
+      gpu_memory_utilization: 0.2
+      enforce_eager: true
+      trust_remote_code: true
+      enable_prefix_caching: false
+      engine_output_type: audio
+      max_model_len: 18192
+      max_num_batched_tokens: 18192
+      async_scheduling: false
+      load_format: dummy
+    engine_input_source: [ 0 ]
+    is_comprehension: false
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav
+    final_output: true
+    final_output_type: audio
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 18192
+      seed: 42
+      detokenize: false

vllm_omni/model_executor/stage_configs/mimo_audio.yaml

@@ -39,7 +39,7 @@ stage_args:
     stage_type: llm
     runtime:
       process: true            # Run this stage in a separate process
-      devices: "0"            # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
+      devices: "1"            # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
     engine_args:
       model_stage: code2wav
       max_num_seqs: 1