[Test] Add long text output correctness test for Qwen3-Omni (#3539)

Signed-off-by: ZeldaHuang <hzm414167@alibaba-inc.com>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: tzhouam <tzhouam@connect.ust.hk>

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -539,3 +539,24 @@ def test_language_001(omni_server, openai_client) -> None:
     }
 
     openai_client.send_omni_request(request_config)
+
+
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_to_audio_long_output_001(omni_server, openai_client) -> None:
+    """
+    Input Modal: text only (long-form generation prompt).
+    Output Modal: text, audio (default ``modalities``);
+    Input Setting: stream=True
+    Datasets: single request
+    """
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        content_text="Tell a 300-word story.",
+    )
+
+    request_config = {"model": omni_server.model, "messages": messages, "stream": True}
+    responses = openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
+    text = responses[0].text_content if responses else ""
+    word_count = len(text.split())
+    assert word_count >= 200, f"Expected at least 200 words in long output, got {word_count}"

tests/helpers/assertions.py

@@ -477,6 +477,7 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve
         # Verify similarity (Whisper transcript vs streamed/detokenized text)
         if "audio" in modalities:
             audio_ref_text = request_config.get("audio_ref_text")
+            similarity_threshold = request_config.get("similarity_threshold", 0.8)
             if "text" in modalities:
                 transcript = (response.audio_content or "").strip()
                 text_output = (response.text_content or "").strip()
@@ -484,14 +485,14 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve
                     transcript.lower(),
                     text_output.lower(),
                 )
-                assert similarity > 0.9, "The audio content is not same as the text"
                 print(f"similarity is: {similarity}")
+                assert similarity > similarity_threshold, "The audio content is not same as the text"
             if audio_ref_text:
                 audio_similarity = cosine_similarity_text(
                     response.audio_content.lower(),
                     str(audio_ref_text).lower(),
                 )
-                assert audio_similarity > 0.9, (
+                assert audio_similarity > similarity_threshold, (
                     f"The audio content does not match reference text: similarity={audio_similarity:.3f}"
                 )
 

tests/helpers/stage_config.py

@@ -332,15 +332,14 @@ _CI_OVERLAYS: dict[str, dict[str, Any]] = {
             },
             {
                 "stage_id": 1,
-                "gpu_memory_utilization": 0.5,
                 "max_num_seqs": 5,
+                "gpu_memory_utilization": 0.5,
                 "max_model_len": 32768,
                 "default_sampling_params": {"max_tokens": 1000},
             },
             {
                 "stage_id": 2,
                 "max_num_seqs": 5,
-                "max_num_batched_tokens": 100000,
                 "default_sampling_params": {"max_tokens": 2000},
             },
         ],

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -112,6 +112,11 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             num_computed_tokens = request.num_computed_tokens
             required_tokens = len(request.prompt_token_ids) - num_computed_tokens
+            if not self.scheduler_config.enable_chunked_prefill and required_tokens > token_budget:
+                # If chunked_prefill is disabled,
+                # we can stop the scheduling here.
+                break
+            # async_chunk: don't schedule placeholder tokens when no new chunk is available.
             if required_tokens <= 0:
                 if (
                     self.chunk_transfer_adapter is not None

vllm_omni/deploy/qwen3_omni_moe.yaml

@@ -53,12 +53,13 @@ stages:
       repetition_penalty: 1.05
 
   - stage_id: 2
-    max_num_batched_tokens: 51200
+    max_num_batched_tokens: 65536
     max_num_seqs: 64
     gpu_memory_utilization: 0.1
     enforce_eager: false
     trust_remote_code: true
     enable_prefix_caching: false
+    enable_chunked_prefill: false
     async_scheduling: false
     devices: "1"
     input_connectors:

vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml

@@ -55,13 +55,14 @@ stages:
       repetition_penalty: 1.05
 
   - stage_id: 2
-    max_num_batched_tokens: 51200
+    max_num_batched_tokens: 65536
     max_num_seqs: 64
     gpu_memory_utilization: 0.1
     enforce_eager: true
     trust_remote_code: true
     enable_prefix_caching: false
     async_scheduling: false
+    enable_chunked_prefill: false
     devices: "1,2"
     num_replicas: 2
     input_connectors: