[Feature]Support async computation and communication across stages by chunks (#727)

Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>
Signed-off-by: Rein Yang <ruiruyang2@gmail.com>
Co-authored-by: Rein Yang <ruiruyang2@gmail.com>
Co-authored-by: CHEN <116010019@link.cuhk.edu.cn>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

examples/online_serving/qwen3_omni/README.md

@@ -12,6 +12,12 @@ Please refer to [README.md](../../../README.md)
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
+If you want to open async chunking for qwen3-omni, launch the server with command below
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
+```
+
 If you have custom stage configs file, launch the server with command below
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file

examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -1,4 +1,5 @@
 import base64
+import concurrent.futures
 import os
 from typing import NamedTuple
 
@@ -414,51 +415,67 @@ def run_multimodal_generation(args) -> None:
     else:
         output_modalities = None
 
-    chat_completion = client.chat.completions.create(
-        messages=[
-            get_system_prompt(),
-            prompt,
-        ],
-        model=model_name,
-        modalities=output_modalities,
-        extra_body=extra_body,
-        stream=args.stream,
-    )
+    # Test multiple concurrent completions
+    num_concurrent_requests = args.num_concurrent_requests
 
+    with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
+        # Submit multiple completion requests concurrently
+        futures = [
+            executor.submit(
+                client.chat.completions.create,
+                messages=[
+                    get_system_prompt(),
+                    prompt,
+                ],
+                model=model_name,
+                modalities=output_modalities,
+                extra_body=extra_body,
+                stream=args.stream,
+            )
+            for _ in range(num_concurrent_requests)
+        ]
+
+        # Wait for all requests to complete and collect results
+        chat_completions = [future.result() for future in concurrent.futures.as_completed(futures)]
+
+    assert len(chat_completions) == num_concurrent_requests
     count = 0
     if not args.stream:
-        for choice in chat_completion.choices:
-            if choice.message.audio:
-                audio_data = base64.b64decode(choice.message.audio.data)
-                audio_file_path = f"audio_{count}.wav"
-                with open(audio_file_path, "wb") as f:
-                    f.write(audio_data)
-                print(f"Audio saved to {audio_file_path}")
-                count += 1
-            elif choice.message.content:
-                print("Chat completion output from text:", choice.message.content)
-    else:
-        printed_content = False
-        for chunk in chat_completion:
-            for choice in chunk.choices:
-                if hasattr(choice, "delta"):
-                    content = getattr(choice.delta, "content", None)
-                else:
-                    content = None
-
-                if getattr(chunk, "modality", None) == "audio" and content:
-                    audio_data = base64.b64decode(content)
+        # Verify all completions succeeded
+        for chat_completion in chat_completions:
+            for choice in chat_completion.choices:
+                if choice.message.audio:
+                    audio_data = base64.b64decode(choice.message.audio.data)
                     audio_file_path = f"audio_{count}.wav"
                     with open(audio_file_path, "wb") as f:
                         f.write(audio_data)
-                    print(f"\nAudio saved to {audio_file_path}")
+                    print(f"Audio saved to {audio_file_path}")
                     count += 1
+                elif choice.message.content:
+                    print("Chat completion output from text:", choice.message.content)
+    else:
+        printed_content = False
+        for chat_completion in chat_completions:
+            for chunk in chat_completion:
+                for choice in chunk.choices:
+                    if hasattr(choice, "delta"):
+                        content = getattr(choice.delta, "content", None)
+                    else:
+                        content = None
 
-                elif getattr(chunk, "modality", None) == "text":
-                    if not printed_content:
-                        printed_content = True
-                        print("\ncontent:", end="", flush=True)
-                    print(content, end="", flush=True)
+                    if getattr(chunk, "modality", None) == "audio" and content:
+                        audio_data = base64.b64decode(content)
+                        audio_file_path = f"audio_{count}.wav"
+                        with open(audio_file_path, "wb") as f:
+                            f.write(audio_data)
+                        print(f"\nAudio saved to {audio_file_path}")
+                        count += 1
+
+                    elif getattr(chunk, "modality", None) == "text":
+                        if not printed_content:
+                            printed_content = True
+                            print("\ncontent:", end="", flush=True)
+                        print(content, end="", flush=True)
 
 
 def parse_args():
@@ -517,6 +534,12 @@ def parse_args():
         action="store_true",
         help="Stream the response.",
     )
+    parser.add_argument(
+        "--num-concurrent-requests",
+        type=int,
+        default=1,
+        help="Number of concurrent requests to send. Default is 1.",
+    )
 
     return parser.parse_args()