[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>

docs/api/README.md

@@ -53,8 +53,10 @@ Engine classes for offline and online inference.
 
 Core scheduling and caching components.
 
+- [vllm_omni.core.sched.omni_ar_scheduler.KVCacheTransferData][]
 - [vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler][]
 - [vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler][]
+- [vllm_omni.core.sched.output.OmniCachedRequestData][]
 - [vllm_omni.core.sched.output.OmniNewRequestData][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedGroupResidualVectorQuantization][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedResidualVectorQuantization][]

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -15,6 +15,12 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
 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/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()
 

vllm_omni/config/model.py

@@ -1,4 +1,5 @@
 import warnings
+from dataclasses import field
 from typing import Any
 
 import torch
@@ -42,6 +43,7 @@ class OmniModelConfig(ModelConfig):
 
     Attributes:
         stage_id: Identifier for the stage in a multi-stage pipeline (default: 0)
+        async_chunk: If set to True, perform async chunk
         model_stage: Stage type identifier, e.g., "thinker" or "talker"
             (default: "thinker")
         model_arch: Model architecture name
@@ -49,6 +51,8 @@ class OmniModelConfig(ModelConfig):
         engine_output_type: Optional output type specification for the engine.
             Used to route outputs to appropriate processors (e.g., "image",
             "audio", "latents"). If None, output type is inferred.
+        stage_connector_config: Stage connector configuration dictionary.
+            Contains "name" (connector name), "extra" (extra connector config).
 
     Example:
         >>> config = OmniModelConfig(
@@ -59,10 +63,18 @@ class OmniModelConfig(ModelConfig):
     """
 
     stage_id: int = 0
+    async_chunk: bool = False
     model_stage: str = "thinker"
     model_arch: str = "Qwen2_5OmniForConditionalGeneration"
     engine_output_type: str | None = None
     hf_config_name: str | None = None
+    custom_process_next_stage_input_func: str | None = None
+    stage_connector_config: dict[str, Any] = field(
+        default_factory=lambda: {
+            "name": "SharedMemoryConnector",
+            "extra": {},
+        }
+    )
     omni_kv_config: dict | None = None
 
     @property

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import importlib
 from collections import defaultdict
 from dataclasses import asdict, dataclass
 from time import time
@@ -19,6 +20,10 @@ from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus
 from vllm.v1.spec_decode.metrics import SpecDecodingStats
 
+from vllm_omni.distributed.omni_connectors.adapter import get_chunk, put_chunk
+from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
+from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
+
 logger = init_logger(__name__)
 
 
@@ -59,6 +64,25 @@ class OmniARScheduler(VLLMScheduler):
 
         # Track requests that have already triggered prefill transfer to avoid duplicates
         self.transfer_triggered_requests: set[str] = set()
+        model_config = self.vllm_config.model_config
+        self.omni_connector = None
+        if model_config.async_chunk:
+            connector_config = model_config.stage_connector_config
+            connector_specs = ConnectorSpec(
+                name=connector_config.get("name", "SharedMemoryConnector"),
+                extra=connector_config.get("extra", {}),
+            )
+            self.omni_connector = OmniConnectorFactory.create_connector(connector_specs)
+
+            custom_process_next_stage_input_func = getattr(
+                self.vllm_config.model_config, "custom_process_next_stage_input_func", None
+            )
+            if custom_process_next_stage_input_func:
+                module_path, func_name = custom_process_next_stage_input_func.rsplit(".", 1)
+                module = importlib.import_module(module_path)
+                self.custom_process_next_stage_input_func = getattr(module, func_name)
+
+        self.stage_id = getattr(self.vllm_config.model_config, "stage_id", None)
 
     def _get_kv_transfer_criteria(self) -> dict | None:
         # Note: vllm_config is available in Scheduler after super().__init__
@@ -156,6 +180,8 @@ class OmniARScheduler(VLLMScheduler):
                 new_list.append(omni_nr)
 
             scheduler_output.scheduled_new_reqs = new_list  # type: ignore[assignment]
+            if self.omni_connector is not None:
+                get_chunk(self.omni_connector, scheduler_output)
 
             # Add information about requests needing KV cache transfer
             scheduler_output.finished_requests_needing_kv_transfer = self.get_finished_requests_needing_kv_transfer()
@@ -337,6 +363,9 @@ class OmniARScheduler(VLLMScheduler):
                         num_nans_in_logits=request.num_nans_in_logits,
                     )
                 )
+                if self.omni_connector is not None:
+                    custom_process_next_stage_input_func = self.custom_process_next_stage_input_func
+                    put_chunk(self.omni_connector, pooler_output, request, custom_process_next_stage_input_func)
             else:
                 # Invariant: EngineCore returns no partial prefill outputs.
                 assert not prompt_logprobs_tensors

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -12,11 +12,27 @@ from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutp
 from vllm.v1.request import Request, RequestStatus
 from vllm.v1.spec_decode.metrics import SpecDecodingStats
 
-from vllm_omni.core.sched.output import OmniNewRequestData
+from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData
+from vllm_omni.distributed.omni_connectors.adapter import get_chunk_for_generation
+from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
+from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
 from vllm_omni.outputs import OmniModelRunnerOutput
 
 
 class OmniGenerationScheduler(VLLMScheduler):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        model_config = self.vllm_config.model_config
+        self.omni_connector = None
+        if model_config.async_chunk:
+            connector_config = model_config.stage_connector_config
+            connector_specs = ConnectorSpec(
+                name=connector_config.get("name", "SharedMemoryConnector"),
+                extra=connector_config.get("extra", {}),
+            )
+            self.omni_connector = OmniConnectorFactory.create_connector(connector_specs)
+        self.stage_id = getattr(self.vllm_config.model_config, "stage_id", None)
+
     def schedule(self) -> SchedulerOutput:
         """Diffusion fast path:
         - Feed all input tokens of the request at once
@@ -32,27 +48,54 @@ class OmniGenerationScheduler(VLLMScheduler):
 
         req_to_new_blocks: dict[str, KVCacheBlocks] = {}
         num_scheduled_tokens: dict[str, int] = {}
+        scheduled_running_reqs: list[Request] = []
         scheduled_spec_decode_tokens: dict[str, list[int]] = {}
         scheduled_encoder_inputs: dict[str, list[int]] = {}
+        cached_prompt_token_ids: dict[str, list[int]] = {}
 
         # Temporary queue: preserve waiting order, do not disturb non-diffusion requests
         skipped_waiting_requests = create_request_queue(self.policy)
+        req_index = 0
+        while req_index < len(self.running) and token_budget > 0:
+            request = self.running[req_index]
+            if self.omni_connector is not None:
+                get_chunk_for_generation(self.omni_connector, request)
+            num_computed_tokens = request.num_computed_tokens
+            required_tokens = max(len(request.prompt_token_ids) - num_computed_tokens, 1)
+            num_new_tokens = min(required_tokens, token_budget)
+            new_blocks = self.kv_cache_manager.allocate_slots(
+                request,
+                num_new_tokens,
+                num_lookahead_tokens=self.num_lookahead_tokens,
+            )
+            if new_blocks is None:
+                # Allocation failed (e.g., VRAM pressure); stop fast path and
+                # fall back to default scheduling
+                # Put the current request back to the head of the waiting queue
+                # Note: the original queue order is preserved
+                break
+            if self.log_stats:
+                request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp)
+            req_to_new_blocks[request.request_id] = new_blocks
+            num_scheduled_tokens[request.request_id] = num_new_tokens
+            cached_prompt_token_ids[request.request_id] = request.prompt_token_ids
+            token_budget -= num_new_tokens
+            scheduled_running_reqs.append(request)
+            req_index += 1
 
         # Fast path selection and scheduling (treat all as diffusion requests,
         # independent of pooling_params)
         while self.waiting and token_budget > 0 and len(self.running) < self.max_num_running_reqs:
             request = self.waiting.peek_request()
+            if self.omni_connector is not None:
+                get_chunk_for_generation(self.omni_connector, request)
             # Uniformly treat as diffusion. A feature flag can be added later
             # via config or request tag.
 
             # Allocate all input tokens for the request in one shot
             # (allocate 1 placeholder if zero)
-            required_tokens = max(getattr(request, "num_prompt_tokens", 0), 1)
-            if required_tokens > token_budget:
-                # Insufficient budget to process all inputs at once;
-                # stop fast path attempt
-                break
-            num_new_tokens = required_tokens
+            required_tokens = max(len(request.prompt_token_ids), 1)
+            num_new_tokens = min(required_tokens, token_budget)
             new_blocks = self.kv_cache_manager.allocate_slots(
                 request,
                 num_new_tokens,
@@ -68,7 +111,6 @@ class OmniGenerationScheduler(VLLMScheduler):
             # Officially schedule this request
             request = self.waiting.pop_request()
             self.running.append(request)
-            request.status = RequestStatus.RUNNING
             if self.log_stats:
                 request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp)
 
@@ -109,13 +151,24 @@ class OmniGenerationScheduler(VLLMScheduler):
             ]
         # No running/resumed reqs scheduled in our fast path
         cached_reqs_data = self._make_cached_request_data(
-            running_reqs=[],
+            running_reqs=scheduled_running_reqs,
             resumed_reqs=[],
             num_scheduled_tokens=num_scheduled_tokens,
             spec_decode_tokens=scheduled_spec_decode_tokens,
             req_to_new_blocks=req_to_new_blocks,
         )
 
+        cached_reqs_data = OmniCachedRequestData(
+            req_ids=cached_reqs_data.req_ids,
+            resumed_req_ids=cached_reqs_data.resumed_req_ids,
+            new_token_ids=cached_reqs_data.new_token_ids,
+            all_token_ids=cached_reqs_data.all_token_ids,
+            new_block_ids=cached_reqs_data.new_block_ids,
+            num_computed_tokens=cached_reqs_data.num_computed_tokens,
+            num_output_tokens=cached_reqs_data.num_output_tokens,
+            prompt_token_ids=cached_prompt_token_ids,
+        )
+
         total_num_scheduled_tokens = sum(num_scheduled_tokens.values())
         scheduler_output = SchedulerOutput(
             scheduled_new_reqs=new_reqs_data,
@@ -257,21 +310,20 @@ class OmniGenerationScheduler(VLLMScheduler):
             new_logprobs = None
             new_token_ids = generated_token_ids
             kv_transfer_params = None
-            status_before_stop = request.status
             pooler_output = None
             if pooler_outputs:
                 pooler_output = pooler_outputs[req_index]
 
             # Diffusion request: completes in one step; mark finished and free resources
-            request.status = RequestStatus.FINISHED_STOPPED
-            # Optional: set a stop_reason for front-end clarity
-            # (does not affect protocol)
-            request.stop_reason = request.stop_reason  # or "generation_done"
-            kv_transfer_params = self._free_request(request)
-            if status_before_stop == RequestStatus.RUNNING:
+            if request.status == RequestStatus.FINISHED_STOPPED or (
+                self.omni_connector is None and request.num_computed_tokens >= request.num_prompt_tokens
+            ):
+                request.status = RequestStatus.FINISHED_STOPPED
+                # Optional: set a stop_reason for front-end clarity
+                # (does not affect protocol)
+                request.stop_reason = request.stop_reason  # or "generation_done"
+                kv_transfer_params = self._free_request(request)
                 stopped_running_reqs.add(request)
-            else:
-                stopped_preempted_reqs.add(request)
 
             # Extract sample logprobs if needed.
             if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs:

vllm_omni/core/sched/output.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 
-from vllm.v1.core.sched.output import NewRequestData
+from vllm.v1.core.sched.output import CachedRequestData, NewRequestData
 from vllm.v1.request import Request
 
 from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
@@ -53,3 +53,14 @@ class OmniNewRequestData(NewRequestData):
             prefill_token_ids=prefill_token_ids,
             additional_information=request.additional_information,
         )
+
+
+@dataclass
+class OmniCachedRequestData(CachedRequestData):
+    """Cached request data for omni models with embeddings support.
+
+    Args:
+        prompt_token_ids: Mapping from request ID to list of prompt token IDs
+    """
+
+    prompt_token_ids: dict[int, list[int]]

vllm_omni/diffusion/worker/gpu_diffusion_model_runner.py

@@ -207,10 +207,11 @@ class GPUDiffusionModelRunner:
                 start_time = time.time()
 
                 while True:
+                    get_key = f"omni_{from_stage}_to_{to_stage}_kv_cache_{req.request_id}"
                     result = self.connector.get(
                         from_stage=from_stage,
                         to_stage=to_stage,
-                        request_id=f"kv_cache_{req.request_id}",
+                        get_key=get_key,
                     )
                     if result:
                         break

vllm_omni/distributed/omni_connectors/adapter.py

@@ -5,7 +5,14 @@
 
 import time
 from collections.abc import Callable
-from typing import Any
+from typing import TYPE_CHECKING, Any
+
+import torch
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.request import Request, RequestStatus
+
+if TYPE_CHECKING:
+    from .connectors.base import OmniConnectorBase
 
 from vllm_omni.entrypoints.stage_utils import OmniStageTaskType
 
@@ -170,3 +177,225 @@ def try_recv_via_connector(
             # but for Stage-0 seed it should be there.
             # We'll return None to let caller handle error if strictly required.
             return None, None
+
+
+def get_chunk(
+    connector: "OmniConnectorBase",
+    scheduler_output: SchedulerOutput,
+) -> None:
+    """Retrieve a chunk of pooling output.
+
+    Args:
+        connector: OmniConnectorBase instance
+        scheduler_output: Partial scheduler output dictionary
+
+    Returns:
+        None: This function modifies scheduler_output in place
+    """
+    stage_id = connector.stage_id
+    if stage_id == 0:
+        return
+
+    target_stage_id = stage_id - 1
+    # Handle new requests
+    for new_req_data in scheduler_output.scheduled_new_reqs:
+        req_id = new_req_data.req_id[0:25]
+        chunk_id = connector.get_requests[req_id]
+        connector_get_key = f"{req_id}_{target_stage_id}_{chunk_id}"
+        payload_data = get_through_connector(connector, target_stage_id, stage_id, req_id, connector_get_key)
+        if payload_data:
+            new_req_data.additional_information = payload_data
+            if payload_data.get("finished"):
+                connector.finished_requests.add(req_id)
+
+    # Handle cached/running requests
+    cached_reqs = scheduler_output.scheduled_cached_reqs
+    if not hasattr(cached_reqs, "additional_information"):
+        cached_reqs.additional_information = {}
+
+    for i, cached_req_id in enumerate(cached_reqs.req_ids):
+        req_id = cached_req_id[0:25]
+        if req_id in connector.finished_requests:
+            continue
+        chunk_id = connector.get_requests[req_id]
+        connector_get_key = f"{req_id}_{target_stage_id}_{chunk_id}"
+        payload_data = get_through_connector(connector, target_stage_id, stage_id, req_id, connector_get_key)
+        if payload_data:
+            cached_reqs.additional_information[cached_req_id] = payload_data
+            if payload_data.get("finished"):
+                connector.finished_requests.add(req_id)
+
+
+def get_through_connector(connector, target_stage_id, stage_id, req_id, connector_get_key):
+    # Wait for data from previous stage
+    import time
+
+    # TODO: add correct check mechanism for the payload_data
+    max_wait = 300
+    for _ in range(max_wait):
+        result = connector.get(
+            from_stage=str(target_stage_id),
+            to_stage=str(stage_id),
+            get_key=connector_get_key,
+        )
+        payload_data = None
+        if result:
+            payload_data, size = result
+            if payload_data:
+                connector.request_prompt_token_ids[req_id] = payload_data.get("thinker_input_ids", [])
+                connector.get_requests[req_id] += 1
+                logger.debug(f"[Stage-{stage_id}] Received one chunk for request {connector_get_key}")
+                break
+        time.sleep(0.1)
+    return payload_data
+
+
+def get_chunk_for_generation(
+    connector: "OmniConnectorBase",
+    request: Request,
+) -> None:
+    """Retrieve a chunk of pooling output.
+
+    Args:
+        connector: OmniConnectorBase instance
+        request: Request object
+
+    Returns:
+        None: This function modifies request in place
+    """
+    stage_id = connector.stage_id
+    target_stage_id = stage_id - 1
+    request_id = request.request_id[0:25]
+
+    if request_id in connector.finished_requests:
+        return
+
+    chunk_id = connector.get_requests[request_id]
+    connector_get_key = f"{request_id}_{target_stage_id}_{chunk_id}"
+    payload_data = get_through_connector(connector, target_stage_id, stage_id, request_id, connector_get_key)
+    if not payload_data:
+        return
+
+    if payload_data.get("finished"):
+        connector.finished_requests.add(request_id)
+        request.status = RequestStatus.FINISHED_STOPPED
+
+    # TODO: remove special handling for prompt token ids ?
+    if chunk_id == 0:
+        request.prompt_token_ids = payload_data.get("code_predictor_codes", [])
+    else:
+        request.prompt_token_ids += payload_data.get("code_predictor_codes", [])
+
+
+def put_chunk(
+    connector: "OmniConnectorBase",
+    pooling_output: dict[str, Any],
+    request: Request,
+    custom_process_input_func: Callable[[dict[str, Any], Request], dict[str, Any] | None] | None = None,
+) -> None:
+    """Store a chunk of pooling output.
+
+    Args:
+        connector: OmniConnectorBase instance
+        pooling_output: Partial pooling output dictionary
+        request: Request object
+        custom_process_input_func: Optional custom function to process input
+
+    Returns:
+        None: This function sends data via connector
+    """
+    stage_id = connector.stage_id
+    next_stage_id = stage_id + 1
+    request_id = request.request_id[0:25]
+    prompt_token_ids = request.prompt_token_ids
+    connector.request_prompt_token_ids[request_id] = prompt_token_ids
+    chunk_id = connector.put_requests[request_id]
+    connector_put_key = f"{request_id}_{stage_id}_{chunk_id}"
+    payload_data = None
+
+    # TODO: add default process_input_func to handle the payload_data ?
+    if custom_process_input_func:
+        try:
+            payload_data = custom_process_input_func(
+                pooling_output=pooling_output,
+                request=request,
+            )
+        except Exception as e:
+            logger.error(f"Failed to use custom_process_input_func for payload extraction: {e}")
+
+        if not payload_data:
+            logger.warning(f"[Stage-{stage_id}] No payload data to send for request {request_id}")
+            return
+
+        if stage_id == 0 and chunk_id == 0:
+            if connector.request_payload.get(request_id) is None:
+                connector.request_payload[request_id] = payload_data
+                return
+            else:
+                save_payload = connector.request_payload.get(request_id)
+                payload_data["thinker_embeddings"] = torch.cat(
+                    (save_payload.get("thinker_embeddings"), payload_data.get("thinker_embeddings")), dim=0
+                )
+                payload_data["thinker_hidden_states"] = torch.cat(
+                    (save_payload.get("thinker_hidden_states"), payload_data.get("thinker_hidden_states")), dim=0
+                )
+                logger.info(f"[Stage-{stage_id}] Merged embeddings and hidden states for request {request_id}")
+
+        if stage_id == 1:
+            # TODO: Make parameters configurable and optimize algorithms
+            chunk_size = left_context_size = 25
+            connector.code_prompt_token_ids[request_id].append(payload_data.get("code_predictor_codes", []))
+            length = len(connector.code_prompt_token_ids[request_id])
+            chunk_length = length % chunk_size
+            if chunk_length != 0 and not payload_data.get("finished"):
+                return
+
+            context_length = chunk_length if chunk_length != 0 else chunk_size
+            end_index = min(length, left_context_size + context_length)
+            payload_data["code_predictor_codes"] = (
+                torch.tensor(connector.code_prompt_token_ids[request_id][-end_index:])
+                .transpose(0, 1)
+                .reshape(-1)
+                .tolist()
+            )
+
+        success, size, metadata = connector.put(
+            from_stage=str(stage_id), to_stage=str(next_stage_id), put_key=connector_put_key, data=payload_data
+        )
+
+        if success:
+            connector.put_requests[request_id] += 1
+            logger.info(f"[Stage-{stage_id}] Sent {connector_put_key}")
+
+
+def compute_talker_prompt_ids_length(prompt_ids: list[int]) -> int:
+    """Compute the length of the talker prompt ids.
+
+    Args:
+        prompt_ids: The prompt ids tensor.
+
+    Returns:
+        The length of the talker prompt ids.
+    """
+    im_start_token_id = 151644
+    system_token_id = 8948
+    user_token_id = 872
+    assistant_token_id = 77091
+    im_start_indexes = [i for i in range(len(prompt_ids)) if prompt_ids[i] == im_start_token_id]
+    im_start_indexes.append(len(prompt_ids))
+    sum_user_len = 0
+    assistant_len = 0
+    for i in range(len(im_start_indexes) - 1):
+        s = im_start_indexes[i]
+        e = im_start_indexes[i + 1]
+        role = prompt_ids[s + 1]
+        if role == system_token_id:
+            continue
+        elif role == user_token_id:
+            sum_user_len += e - s
+        elif role == assistant_token_id and i == len(im_start_indexes) - 2:
+            assistant_len += 9  # 3 + 4 + 1 + 1
+        else:
+            pass
+
+    return sum_user_len + assistant_len

vllm_omni/distributed/omni_connectors/connectors/base.py

@@ -13,15 +13,13 @@ class OmniConnectorBase(ABC):
     """Base class for all OmniConnectors."""
 
     @abstractmethod
-    def put(
-        self, from_stage: str, to_stage: str, request_id: str, data: Any
-    ) -> tuple[bool, int, dict[str, Any] | None]:
+    def put(self, from_stage: str, to_stage: str, put_key: str, data: Any) -> tuple[bool, int, dict[str, Any] | None]:
         """Store Python object, internal serialization handled by connector.
 
         Args:
             from_stage: Source stage identifier
             to_stage: Destination stage identifier
-            request_id: Unique request identifier
+            put_key: Unique request identifier
             data: Python object to store
 
         Returns:
@@ -31,16 +29,13 @@ class OmniConnectorBase(ABC):
         pass
 
     @abstractmethod
-    def get(
-        self, from_stage: str, to_stage: str, request_id: str, metadata: dict[str, Any] | None = None
-    ) -> tuple[Any, int] | None:
+    def get(self, from_stage: str, to_stage: str, get_key: str, metadata=None) -> tuple[Any, int] | None:
         """Retrieve Python object and payload size (bytes).
 
         Args:
             from_stage: Source stage identifier
             to_stage: Destination stage identifier
-            request_id: Unique request identifier
-            metadata: Optional transport-specific metadata from the put operation
+            get_key: Unique request identifier
 
         Returns:
             Tuple of (Python object, serialized byte size) if found, None otherwise

vllm_omni/distributed/omni_connectors/connectors/shm_connector.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import time
+from collections import defaultdict
 from typing import Any
 
 from vllm_omni.entrypoints.stage_utils import shm_read_bytes, shm_write_bytes
@@ -19,6 +21,14 @@ class SharedMemoryConnector(OmniConnectorBase):
 
     def __init__(self, config: dict[str, Any]):
         self.config = config
+        self.stage_id = config.get("stage_id", -1)
+        self.device = config.get("device", "cuda:0")
+        self.put_requests: dict[str, int] = defaultdict(int)
+        self.get_requests: dict[str, int] = defaultdict(int)
+        self.finished_requests: set[str] = set()
+        self.request_payload = {}
+        self.request_prompt_token_ids: dict[str, list[int]] = defaultdict(list)
+        self.code_prompt_token_ids: dict[str, list[list[int]]] = defaultdict(list)
         # Default threshold matches legacy behavior (64KB)
         self.threshold = int(config.get("shm_threshold_bytes", 65536))
         self._metrics = {
@@ -29,9 +39,7 @@ class SharedMemoryConnector(OmniConnectorBase):
             "inline_writes": 0,
         }
 
-    def put(
-        self, from_stage: str, to_stage: str, request_id: str, data: Any
-    ) -> tuple[bool, int, dict[str, Any] | None]:
+    def put(self, from_stage: str, to_stage: str, put_key: str, data: Any) -> tuple[bool, int, dict[str, Any] | None]:
         try:
             # Always serialize first to check size (and for SHM writing)
             # Note: For extremely large objects in "inline" mode (e.g. Ray),
@@ -40,18 +48,19 @@ class SharedMemoryConnector(OmniConnectorBase):
             payload = self.serialize_obj(data)
             size = len(payload)
 
-            if size > self.threshold:
+            metadata = {}
+            # if size > self.threshold:
+            if True:  # TODO: correct put & get logic
                 # Use Shared Memory
-                shm_name = f"omni_{from_stage}_to_{to_stage}_{request_id}"
-                meta = shm_write_bytes(payload, name=shm_name)
+                meta = shm_write_bytes(payload, name=put_key)
                 # meta contains {'name': ..., 'size': ...}
-                metadata = {"shm": meta, "size": size}
+                metadata[put_key] = {"shm": meta, "size": size}
                 self._metrics["shm_writes"] += 1
             else:
                 # Inline - pass bytes directly to avoid double serialization of the object
                 # We already serialized it to check size, so we pass the bytes.
                 # The Queue will pickle these bytes (fast), avoiding re-serializing the complex object.
-                metadata = {"inline_bytes": payload, "size": size}
+                metadata[put_key] = {"inline_bytes": payload, "size": size}
                 self._metrics["inline_writes"] += 1
 
             self._metrics["puts"] += 1
@@ -60,83 +69,40 @@ class SharedMemoryConnector(OmniConnectorBase):
             return True, size, metadata
 
         except Exception as e:
-            logger.error(f"SharedMemoryConnector put failed for req {request_id}: {e}")
+            logger.error(f"SharedMemoryConnector put failed for req {put_key}: {e}")
             return False, 0, None
 
-    def get(
-        self, from_stage: str, to_stage: str, request_id: str, metadata: dict[str, Any] | None = None
-    ) -> tuple[Any, int] | None:
-        # Helper to try reading SHM with a specific name
-        def try_read_shm(name: str):
-            # _shm is imported at module level as: import multiprocessing.shared_memory as _shm
-            # but to be safe inside this closure if not picked up:
-            from multiprocessing import shared_memory as shm_pkg
+    def get(self, from_stage: str, to_stage: str, get_key: str, metadata=None) -> tuple[Any, int] | None:
+        from multiprocessing import shared_memory as shm_pkg
 
-            shm = shm_pkg.SharedMemory(name=name)
+        # Wait for shared memory to be available (with retry logic)
+        max_retries = 30
+        retry_delay = 0.1  # 100ms between retries
+        shm = None
+
+        for attempt in range(max_retries):
             try:
-                # We assume the buffer size matches the data size
-                return shm_read_bytes({"name": name, "size": shm.size})
-            finally:
-                shm.close()
+                shm = shm_pkg.SharedMemory(name=get_key)
+                break  # Successfully opened, exit retry loop
+            except FileNotFoundError:
+                if attempt < max_retries - 1:
+                    time.sleep(retry_delay)
+                else:
+                    # Max retries reached, return None
+                    logger.warning(f"Shared memory '{get_key}' not found after {max_retries} retries")
+                    return None, 0
 
-        if not metadata:
-            # Try to infer metadata from request_id for KV cache transfer scenarios
-            # where metadata is not passed out-of-band.
-            shm_name = f"omni_{from_stage}_to_{to_stage}_{request_id}"
-
-            # We need the size to read from SHM using stage_utils.shm_read_bytes
-            try:
-                # Use local import to avoid potential scope issues if _shm isn't available
-                from multiprocessing import shared_memory as shm_pkg
-
-                temp_shm = shm_pkg.SharedMemory(name=shm_name)
-                size = temp_shm.size
-                temp_shm.close()
-
-                metadata = {"shm": {"name": shm_name, "size": size}, "is_fallback": True}
-                logger.debug(f"SharedMemoryConnector inferred metadata for {request_id}: {metadata}")
-            except Exception as e:
-                # If file not found or other error, standard path will likely fail too, but we let it try or return None
-                logger.debug(f"SharedMemoryConnector failed to infer metadata for {request_id} (fallback): {e}")
-                # We can't proceed without size if we rely on shm_read_bytes
-                return None
+        if shm is None:
+            return None, 0
 
         try:
-            obj = None
-            size = 0
+            data_bytes = shm_read_bytes({"name": get_key, "size": shm.size})
+            obj = self.deserialize_obj(data_bytes)
+            return obj, shm.size
+        finally:
+            shm.close()
 
-            if "shm" in metadata:
-                meta = metadata["shm"]
-                # shm_read_bytes handles reading and unlinking
-                data_bytes = shm_read_bytes(meta)
-                obj = self.deserialize_obj(data_bytes)
-                size = metadata.get("size", len(data_bytes))
-            elif "inline_bytes" in metadata:
-                # Deserialize bytes back to object
-                payload = metadata["inline_bytes"]
-                obj = self.deserialize_obj(payload)
-                size = metadata.get("size", len(payload))
-            elif "inline" in metadata:
-                obj = metadata["inline"]
-                size = metadata.get("size", 0)
-                if size == 0:
-                    # Fallback if size wasn't recorded
-                    try:
-                        size = len(self.serialize_obj(obj))
-                    except Exception:
-                        pass
-            else:
-                logger.error(
-                    f"Unknown metadata format in SharedMemoryConnector for req {request_id}: {list(metadata.keys())}"
-                )
-                return None
-
-            self._metrics["gets"] += 1
-            return obj, size
-
-        except Exception as e:
-            logger.error(f"SharedMemoryConnector get failed for req {request_id}: {e}")
-            return None
+        # TODO: update another read method
 
     def cleanup(self, request_id: str) -> None:
         # SHM segments are automatically unlinked during 'get' (shm_read_bytes).

vllm_omni/distributed/omni_connectors/utils/initialization.py

@@ -91,6 +91,8 @@ def get_connectors_config_for_stage(transfer_config: OmniTransferConfig | None,
         # (Worker needs to create connectors to receive data)
         if to_stage == target_stage:
             stage_connectors_config[f"from_stage_{from_stage}"] = {"spec": {"name": spec.name, "extra": spec.extra}}
+        elif from_stage == target_stage and target_stage == "0":
+            stage_connectors_config[f"to_stage_{to_stage}"] = {"spec": {"name": spec.name, "extra": spec.extra}}
 
     return stage_connectors_config
 

vllm_omni/engine/arg_utils.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeTextConfig
 from vllm.engine.arg_utils import EngineArgs
@@ -56,6 +56,11 @@ class OmniEngineArgs(EngineArgs):
         engine_output_type: Optional output type specification for the engine.
             Used to route outputs to appropriate processors (e.g., "image",
             "audio", "latents"). If None, output type is inferred.
+        custom_process_next_stage_input_func: Optional path to a custom function for processing
+            inputs from previous stages
+            If None, default processing is used.
+        stage_connector_spec: Extra configuration for stage connector
+        async_chunk: If set to True, perform async chunk
     """
 
     stage_id: int = 0
@@ -63,6 +68,9 @@ class OmniEngineArgs(EngineArgs):
     model_arch: str = "Qwen2_5OmniForConditionalGeneration"
     engine_output_type: str | None = None
     hf_config_name: str | None = None
+    custom_process_next_stage_input_func: str | None = None
+    stage_connector_spec: dict[str, any] = field(default_factory=dict)
+    async_chunk: bool = False
     omni_kv_config: dict | None = None
 
     def draw_hf_text_config(self, config_dict: dict) -> Qwen3OmniMoeTextConfig:
@@ -112,10 +120,20 @@ class OmniEngineArgs(EngineArgs):
 
         # Add the new omni-specific fields
         config_dict["stage_id"] = self.stage_id
+        config_dict["async_chunk"] = self.async_chunk
         config_dict["model_stage"] = self.model_stage
         config_dict["model_arch"] = self.model_arch
         config_dict["engine_output_type"] = self.engine_output_type
+        # Build stage_connector_config from stage_connector_spec
+        stage_connector_config = {
+            "name": self.stage_connector_spec.get("name", "SharedMemoryConnector"),
+            "extra": self.stage_connector_spec.get("extra", {}).copy(),
+        }
+        stage_connector_config["extra"]["stage_id"] = self.stage_id
+        config_dict["stage_connector_config"] = stage_connector_config
+
         config_dict["hf_config_name"] = self.hf_config_name
+        config_dict["custom_process_next_stage_input_func"] = self.custom_process_next_stage_input_func
         config_dict["omni_kv_config"] = self.omni_kv_config
         if self.hf_config_name is not None:
             config_dict["hf_text_config"] = self.draw_hf_text_config(config_dict)
@@ -140,6 +158,7 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
         engine_output_type: Optional output type specification for the engine.
             Used to route outputs to appropriate processors (e.g., "image",
             "audio", "latents"). If None, output type is inferred.
+        stage_connector_spec: Extra configuration for stage connector
     """
 
     stage_id: int = 0
@@ -147,6 +166,9 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
     model_arch: str = "Qwen2_5OmniForConditionalGeneration"
     engine_output_type: str | None = None
     hf_config_name: str | None = None
+    custom_process_next_stage_input_func: str | None = None
+    stage_connector_spec: dict[str, any] = field(default_factory=dict)
+    async_chunk: bool = False
     omni_kv_config: dict | None = None
 
     def draw_hf_text_config(self, config_dict: dict) -> Qwen3OmniMoeTextConfig:
@@ -186,11 +208,19 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
 
         # Add the new omni-specific fields
         config_dict["stage_id"] = self.stage_id
+        config_dict["async_chunk"] = self.async_chunk
         config_dict["model_stage"] = self.model_stage
         config_dict["model_arch"] = self.model_arch
         config_dict["engine_output_type"] = self.engine_output_type
+        stage_connector_config = {
+            "name": self.stage_connector_spec.get("name", "SharedMemoryConnector"),
+            "extra": self.stage_connector_spec.get("extra", {}).copy(),
+        }
+        stage_connector_config["extra"]["stage_id"] = self.stage_id
+        config_dict["stage_connector_config"] = stage_connector_config
 
         config_dict["hf_config_name"] = self.hf_config_name
+        config_dict["custom_process_next_stage_input_func"] = self.custom_process_next_stage_input_func
         config_dict["omni_kv_config"] = self.omni_kv_config
         if self.hf_config_name is not None:
             config_dict["hf_text_config"] = self.draw_hf_text_config(config_dict)

vllm_omni/engine/output_processor.py

@@ -111,9 +111,15 @@ class OmniRequestState(RequestState):
             for k, v in self.mm_accumulated.items():
                 if isinstance(v, list) and v and isinstance(v[0], torch.Tensor):
                     try:
-                        self.mm_accumulated[k] = torch.cat(v, dim=0)
+                        if k == "audio":
+                            # When the audio tensor shape is inconsistent, torch.cat will fail.
+                            # We need to use torch.cat in -1 dimension.
+                            self.mm_accumulated[k] = torch.cat(v, dim=-1)
+                        else:
+                            self.mm_accumulated[k] = torch.cat(v, dim=0)
                     except Exception:
                         # Keep last tensor on failure
+                        logger.warning(f"Error concatenating tensor for key {k}; keeping last tensor")
                         self.mm_accumulated[k] = v[-1]
                 elif isinstance(v, dict):
                     for sk, sv in v.items():

vllm_omni/entrypoints/async_omni.py

@@ -20,7 +20,7 @@ from vllm.v1.engine.exceptions import EngineDeadError
 # Internal imports (our code)
 from vllm_omni.config import OmniModelConfig
 from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.distributed.omni_connectors.adapter import try_send_via_connector
+from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ids_length, try_send_via_connector
 from vllm_omni.distributed.ray_utils.utils import try_close_ray
 from vllm_omni.engine.input_processor import OmniInputProcessor
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
@@ -325,139 +325,76 @@ class AsyncOmni(OmniBase):
                 self._enable_stats,
                 _wall_start_ts,
             )
-            # Seed stage-0 queue with all requests
-            logger.debug(f"[{self._name}] Seeding request into stage-0")
             req_state = ClientRequestState(request_id)
             req_state.metrics = metrics
             self.request_states[request_id] = req_state
-            # Mark first input time for stage-0
-            metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time()
 
-            sp0: SamplingParams = sampling_params_list[0]  # type: ignore[index]
-            task = {
-                "request_id": request_id,
-                "engine_inputs": prompt,
-                "sampling_params": sp0,
-            }
-            self.stage_list[0].submit(task)
-            _req_start_ts[request_id] = time.time()
-            logger.debug(f"[{self._name}] Enqueued request {request_id} to stage-0")
+            if self.async_chunk:
+                stage_queues = {stage_id: asyncio.Queue() for stage_id in range(num_stages)}
+                req_state.stage_queues = stage_queues
+                for i in range(num_stages):
+                    sp: SamplingParams = sampling_params_list[i]
+                    engine_inputs = prompt
+                    if i != 0:
+                        prompt_token_ids = prompt["prompt_token_ids"]
+                        prompt_1 = prompt.copy()
+                        prompt_1["prompt_token_ids"] = [0] * compute_talker_prompt_ids_length(prompt_token_ids)
+                        prompt_1["multi_modal_data"] = prompt_1["mm_processor_kwargs"] = None
+                        engine_inputs = prompt_1
 
-            logger.debug(f"[{self._name}] Entering scheduling loop: stages={num_stages}")
-            for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]):
-                finished = False
-                while not finished:
-                    result = await req_state.queue.get()
-                    assert stage_id == req_state.stage_id
+                    task = {
+                        "request_id": request_id,
+                        "engine_inputs": engine_inputs,
+                        "sampling_params": sp,
+                    }
+                    self.stage_list[i].submit(task)
+                    metrics.stage_first_ts[i] = metrics.stage_first_ts[0] or time.time()
 
-                    req_id = result.get("request_id")
-                    if "error" in result:
-                        logger.error(
-                            f"[{self._name}] Stage {stage_id} error on request {req_id}: {result['error']}",
-                        )
-                        raise RuntimeError(result)  # Request Finished due to error
+                    logger.info(f"[{self._name}] Enqueued request {request_id} to stage-{str(i)}")
 
-                    engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
-                    if isinstance(engine_outputs, list):
-                        engine_outputs = engine_outputs[0]
-                    finished = engine_outputs.finished
+                _req_start_ts[request_id] = time.time()
 
-                    # Mark last output time for this stage whenever we receive outputs
-                    metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
-                    try:
-                        _m = asdict(result.get("metrics"))
-                        if _m is not None and finished:
-                            metrics.on_stage_metrics(stage_id, req_id, _m)
-                    except Exception as e:
-                        logger.exception(
-                            f"[{self._name}] Failed to process metrics for stage {stage_id}, req {req_id}: {e}",
-                        )
-                    logger.debug(
-                        f"[{self._name}] Stage-{stage_id} completed request {req_id}; forwarding or finalizing",
-                    )
+                logger.info(
+                    f"[{self._name}] Entering scheduling loop: "
+                    f"stages={num_stages}, final_stage={final_stage_id_for_e2e}"
+                )
+                async for output in self._process_async_results(
+                    request_id,
+                    req_state,
+                    metrics,
+                    final_stage_id_for_e2e,
+                    _req_start_ts,
+                    _wall_start_ts,
+                ):
+                    yield output
+            else:
+                sp0: SamplingParams = sampling_params_list[0]  # type: ignore[index]
+                task = {
+                    "request_id": request_id,
+                    "engine_inputs": prompt,
+                    "sampling_params": sp0,
+                }
+                self.stage_list[0].submit(task)
 
-                    if getattr(stage, "final_output", False):
-                        logger.debug(
-                            f"[{self._name}] Request {req_id} finalized at stage-{stage_id}",
-                        )
+                _req_start_ts[request_id] = time.time()
+                # Mark first input time for stage-0
+                metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time()
+                logger.info(
+                    f"[{self._name}] Entering scheduling loop: "
+                    f"stages={num_stages}, final_stage={final_stage_id_for_e2e}"
+                )
 
-                        # End-to-end timing and time-per-token for final output
-                        # (only once per request at the designated final stage)
-                        try:
-                            rid_key = str(req_id)
-                            if stage_id == final_stage_id_for_e2e and rid_key not in metrics.e2e_done and finished:
-                                metrics.on_finalize_request(
-                                    stage_id,
-                                    req_id,
-                                    _req_start_ts.get(req_id, _wall_start_ts),
-                                )
-                        except Exception as e:
-                            logger.exception(
-                                f"[{self._name}] Finalize request handling error for req "
-                                f"{req_id} at stage {stage_id}: {e}",
-                            )
-
-                        # Handle diffusion outputs that already contain images
-                        if stage.final_output_type == "image":
-                            images = []
-                            if isinstance(engine_outputs, OmniRequestOutput) and engine_outputs.images:
-                                images = engine_outputs.images
-                            elif hasattr(engine_outputs, "images") and engine_outputs.images:
-                                images = engine_outputs.images
-                            yield OmniRequestOutput(
-                                stage_id=stage_id,
-                                final_output_type=stage.final_output_type,
-                                request_output=engine_outputs,
-                                images=images,
-                            )
-                        else:
-                            yield OmniRequestOutput(
-                                stage_id=stage_id,
-                                final_output_type=stage.final_output_type,
-                                request_output=engine_outputs,
-                            )
-                if not isinstance(engine_outputs, list):
-                    engine_outputs = [engine_outputs]
-                stage.set_engine_outputs(engine_outputs)
-                # Forward to next stage if there is one
-                next_stage_id = stage_id + 1
-                if next_stage_id <= final_stage_id_for_e2e and finished:
-                    next_stage: OmniStage = self.stage_list[next_stage_id]
-                    next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt)
-                    sp_next: SamplingParams = sampling_params_list[next_stage_id]
-
-                    # Check if we have a connector for this edge
-                    connector_key = (str(stage_id), str(next_stage_id))
-                    connector = self.connectors.get(connector_key)
-
-                    sent_via_connector = False
-                    if connector:
-                        sent_via_connector = try_send_via_connector(
-                            connector=connector,
-                            stage_id=stage_id,
-                            next_stage_id=next_stage_id,
-                            req_id=req_id,
-                            next_inputs=next_inputs,
-                            sampling_params=sp_next,
-                            original_prompt=prompt,
-                            next_stage_queue_submit_fn=self.stage_list[next_stage_id].submit,
-                            metrics=metrics,
-                        )
-
-                    if not sent_via_connector:
-                        # Fallback logic removed as we now enforce connector usage.
-                        # If no connector is found or send fails, we log an error and raise,
-                        # because continuing would cause the request to be silently dropped
-                        # and the orchestrator to hang waiting for completion.
-                        error_msg = (
-                            f"[{self._name}] Failed to send request {req_id} to stage-{next_stage_id} via connector. "
-                            "Configure a connector for this edge or inspect connector logs for details."
-                        )
-                        logger.error(error_msg)
-                        raise RuntimeError(error_msg)
-                    logger.debug(f"[{self._name}] Forwarded request {req_id} to stage-{next_stage_id}")
-                else:
-                    logger.debug(f"[{self._name}] Request {req_id} fully completed")
+                async for output in self._process_sequential_results(
+                    request_id,
+                    req_state,
+                    metrics,
+                    final_stage_id_for_e2e,
+                    _req_start_ts,
+                    _wall_start_ts,
+                    sampling_params_list,
+                    prompt,
+                ):
+                    yield output
 
             logger.debug(f"[{self._name}] All requests completed")
 
@@ -474,6 +411,184 @@ class AsyncOmni(OmniBase):
             logger.info("[AsyncOrchestrator] Request %s aborted.", request_id)
             raise
 
+    async def _process_async_results(
+        self,
+        request_id: str,
+        req_state: ClientRequestState,
+        metrics: OrchestratorMetrics,
+        final_stage_id_for_e2e: int,
+        req_start_ts: dict[int, float],
+        wall_start_ts: float,
+    ) -> AsyncGenerator[OmniRequestOutput, None]:
+        all_stages_finished = {stage_id: False for stage_id in range(final_stage_id_for_e2e + 1)}
+        while not all(all_stages_finished.values()):
+            for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]):
+                if all_stages_finished[stage_id]:
+                    continue
+                result = await req_state.stage_queues[stage_id].get()
+                logger.info(f"[{self._name}] Received result from stage-{stage_id}: {result}")
+                engine_outputs, finished, output_to_yield = self._process_single_result(
+                    result, stage, stage_id, metrics, req_start_ts, wall_start_ts, final_stage_id_for_e2e
+                )
+
+                all_stages_finished[stage_id] = finished
+
+                if output_to_yield:
+                    yield output_to_yield
+
+    async def _process_sequential_results(
+        self,
+        request_id: str,
+        req_state: ClientRequestState,
+        metrics: OrchestratorMetrics,
+        final_stage_id_for_e2e: int,
+        req_start_ts: dict[int, float],
+        wall_start_ts: float,
+        sampling_params_list: list[SamplingParams],
+        prompt: Any,
+    ) -> AsyncGenerator[OmniRequestOutput, None]:
+        for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]):
+            finished = False
+            while not finished:
+                result = await req_state.queue.get()
+                assert stage_id == req_state.stage_id
+                req_id = result.get("request_id")
+                engine_outputs, finished, output_to_yield = self._process_single_result(
+                    result, stage, stage_id, metrics, req_start_ts, wall_start_ts, final_stage_id_for_e2e
+                )
+                if output_to_yield:
+                    yield output_to_yield
+            if not isinstance(engine_outputs, list):
+                engine_outputs = [engine_outputs]
+            stage.set_engine_outputs(engine_outputs)
+            # Forward to next stage if there is one
+            next_stage_id = stage_id + 1
+            if next_stage_id <= final_stage_id_for_e2e and finished:
+                next_stage: OmniStage = self.stage_list[next_stage_id]
+                next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt)
+                sp_next: SamplingParams = sampling_params_list[next_stage_id]
+
+                # Check if we have a connector for this edge
+                connector_key = (str(stage_id), str(next_stage_id))
+                connector = self.connectors.get(connector_key)
+
+                sent_via_connector = False
+                if connector:
+                    sent_via_connector = try_send_via_connector(
+                        connector=connector,
+                        stage_id=stage_id,
+                        next_stage_id=next_stage_id,
+                        req_id=req_id,
+                        next_inputs=next_inputs,
+                        sampling_params=sp_next,
+                        original_prompt=prompt,
+                        next_stage_queue_submit_fn=self.stage_list[next_stage_id].submit,
+                        metrics=metrics,
+                    )
+
+                if not sent_via_connector:
+                    # Fallback logic removed as we now enforce connector usage.
+                    # If no connector is found or send fails, we log an error and raise,
+                    # because continuing would cause the request to be silently dropped
+                    # and the orchestrator to hang waiting for completion.
+                    error_msg = (
+                        f"[{self._name}] Failed to send request {req_id} to stage-{next_stage_id} via connector. "
+                        "Configure a connector for this edge or inspect connector logs for details."
+                    )
+                    logger.error(error_msg)
+                    raise RuntimeError(error_msg)
+                logger.debug(f"[{self._name}] Forwarded request {req_id} to stage-{next_stage_id}")
+            else:
+                logger.debug(f"[{self._name}] Request {req_id} fully completed")
+
+    def _process_single_result(
+        self,
+        result: dict[str, Any],
+        stage: OmniStage,
+        stage_id: int,
+        metrics: OrchestratorMetrics,
+        req_start_ts: dict[int, float],
+        wall_start_ts: float,
+        final_stage_id_for_e2e: int,
+    ) -> tuple[Any, bool, OmniRequestOutput | None]:
+        """
+        Process a single result dictionary from a stage.
+        Returns:
+            engine_outputs: The decoded outputs.
+            finished: Whether the stage processing is finished for this request.
+            output_to_yield: An OmniRequestOutput to yield, or None.
+        """
+        req_id = result.get("request_id")
+        if "error" in result:
+            logger.error(
+                f"[{self._name}] Stage {stage_id} error on request {req_id}: {result['error']}",
+            )
+            raise RuntimeError(result)
+
+        engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
+        if isinstance(engine_outputs, list):
+            engine_outputs = engine_outputs[0]
+
+        finished = engine_outputs.finished
+
+        # Mark last output time
+        metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
+
+        try:
+            _m = asdict(result.get("metrics"))
+            if _m is not None and finished:
+                metrics.on_stage_metrics(stage_id, req_id, _m)
+        except Exception as e:
+            logger.exception(
+                f"[{self._name}] Failed to process metrics for stage {stage_id}, req {req_id}: {e}",
+            )
+
+        logger.debug(
+            f"[{self._name}] Stage-{stage_id} completed request {req_id}; forwarding or finalizing",
+        )
+
+        output_to_yield = None
+        if getattr(stage, "final_output", False):
+            logger.debug(f"[{self._name}] Request {req_id} finalized at stage-{stage_id}")
+
+            # Finalize request metrics if this is the E2E final stage and it's finished
+            try:
+                rid_key = str(req_id)
+                if stage_id == final_stage_id_for_e2e and rid_key not in metrics.e2e_done and finished:
+                    metrics.on_finalize_request(
+                        stage_id,
+                        req_id,
+                        req_start_ts.get(req_id, wall_start_ts),
+                    )
+            except Exception as e:
+                logger.exception(
+                    f"[{self._name}] Finalize request handling error for req {req_id} at stage {stage_id}: {e}",
+                )
+
+            # Construct output to yield
+            images = []
+            if stage.final_output_type == "image":
+                if isinstance(engine_outputs, OmniRequestOutput) and engine_outputs.images:
+                    images = engine_outputs.images
+                elif hasattr(engine_outputs, "images") and engine_outputs.images:
+                    images = engine_outputs.images
+
+            if stage.final_output_type == "image":
+                output_to_yield = OmniRequestOutput(
+                    stage_id=stage_id,
+                    final_output_type=stage.final_output_type,
+                    request_output=engine_outputs,
+                    images=images,
+                )
+            else:
+                output_to_yield = OmniRequestOutput(
+                    stage_id=stage_id,
+                    final_output_type=stage.final_output_type,
+                    request_output=engine_outputs,
+                )
+
+        return engine_outputs, finished, output_to_yield
+
     def _run_output_handler(self) -> None:
         if self.output_handler is not None:
             return
@@ -503,8 +618,12 @@ class AsyncOmni(OmniBase):
                                 dropping output for req {req_id} at stage-{stage_id}"
                             )
                             continue
-                        await req_state.queue.put(result)
-                        req_state.stage_id = stage_id
+                        if hasattr(req_state, "stage_queues") and stage_id in req_state.stage_queues:
+                            await req_state.stage_queues[stage_id].put(result)
+                        else:
+                            # Fallback to old behavior for compatibility
+                            await req_state.queue.put(result)
+                            req_state.stage_id = stage_id
                     if idle:
                         await asyncio.sleep(0.001)  # Avoid CPU overload when idle
                     else:
@@ -512,7 +631,14 @@ class AsyncOmni(OmniBase):
             except Exception as e:
                 logger.exception("AsyncOmni output_handler failed.")
                 for req_state in request_states.values():
-                    await req_state.queue.put({"request_id": req_id, "error": str(e)})
+                    error_msg = {"request_id": req_state.request_id, "error": str(e)}
+                    # Send error to all stage queues
+                    if hasattr(req_state, "stage_queues"):
+                        for queue in req_state.stage_queues.values():
+                            await queue.put(error_msg)
+                    else:
+                        await req_state.queue.put(error_msg)
+                    error_msg = {"request_id": req_state.request_id, "error": str(e)}
                 self.output_handler = None  # Make possible for restart
 
         self.output_handler = asyncio.create_task(output_handler())

vllm_omni/entrypoints/omni.py

@@ -232,6 +232,8 @@ class OmniBase:
         self.worker_backend = worker_backend
         self.ray_address = ray_address
         self.batch_timeout = batch_timeout
+        # async chunk remains the same for each stage
+        self.async_chunk = self._is_async_chunk_enable(self.stage_configs)
 
         # Build OmniStage instances in parallel, preserve original order
         def _build_stage(idx_cfg: tuple[int, Any]) -> tuple[int, OmniStage]:
@@ -261,6 +263,11 @@ class OmniBase:
         # Wait for all stages to report readiness before seeding
         self._wait_for_stages_ready(timeout=init_timeout)
 
+    def _is_async_chunk_enable(self, stage_args: list) -> bool:
+        """get async chunk flag"""
+        engine_args = getattr(stage_args[0], "engine_args", None)
+        return bool(getattr(engine_args, "async_chunk", False))
+
     def _start_stages(self, model: str) -> None:
         """Start all stage processes."""
         if self.worker_backend == "ray":

vllm_omni/entrypoints/omni_stage.py

@@ -673,12 +673,15 @@ def _stage_worker(
                 e,
             )
     # Init engine based on stage_type
-    logger.debug(
-        "[Stage-%s] Initializing %s engine with args keys=%s",
-        stage_id,
-        stage_type,
-        list(engine_args.keys()),
-    )
+    logger.debug("[Stage-%s] Initializing %s engine with args keys=%s", stage_id, stage_type, list(engine_args.keys()))
+    if engine_args.get("async_chunk", False):
+        logger.debug("[Stage-%s] Async chunk enabled, injecting connectors config", stage_id)
+        stage_connector_spec = {}
+        for v in connectors_config.values():
+            stage_connector_spec = dict(v.get("spec", {}))
+            break
+        engine_args["stage_connector_spec"] = stage_connector_spec
+        engine_args["stage_id"] = stage_id
     try:
         if stage_type == "diffusion":
             engine_args.pop("model_stage", None)
@@ -1206,6 +1209,14 @@ async def _stage_worker_async(
         stage_type,
         list(engine_args.keys()),
     )
+    if engine_args.get("async_chunk", False):
+        logger.debug("[Stage-%s] Async chunk enabled, injecting connectors config", stage_id)
+        stage_connector_spec = {}
+        for v in connectors_config.values():
+            stage_connector_spec = dict(v.get("spec", {}))
+            break
+        engine_args["stage_connector_spec"] = stage_connector_spec
+        engine_args["stage_id"] = stage_id
     try:
         if stage_type == "diffusion":
             # For diffusion, we need to extract diffusion-specific config

vllm_omni/entrypoints/openai/serving_chat.py

@@ -18,6 +18,7 @@ try:
 except ImportError:
     soundfile = None
 
+
 from openai.types.chat.chat_completion_audio import ChatCompletionAudio as OpenAIChatCompletionAudio
 from vllm.entrypoints.chat_utils import (
     ChatCompletionMessageParam,
@@ -1213,6 +1214,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         yield f"data: {data}\n\n"
 
                 elif final_output_type == "audio":
+                    role = self.get_chat_request_role(request)
                     choices_data = self._create_audio_choice(omni_res, role, request, stream=True)
                     chunk = OmniChatCompletionStreamResponse(
                         id=request_id,
@@ -1655,7 +1657,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
     ):
         choices: list[ChatCompletionResponseChoice] = []
         final_res = omni_outputs.request_output
-        audio_tensor = final_res.multimodal_output["audio"].float().detach().cpu().numpy()
+        if stream:
+            audio_tensor = final_res.multimodal_output["audio"][-1].float().detach().cpu().numpy()
+        else:
+            audio_tensor = final_res.multimodal_output["audio"].float().detach().cpu().numpy()
 
         # Ensure audio is 1D (flatten if needed)
         if audio_tensor.ndim > 1:

vllm_omni/entrypoints/stage_utils.py

@@ -222,7 +222,8 @@ def maybe_dump_to_shm(obj: Any, threshold: int) -> tuple[bool, Any]:
     """
     payload = serialize_obj(obj)
     if len(payload) > threshold:
-        return True, shm_write_bytes(payload)
+        logger.debug(f"Dumping object to SHM with size: {len(payload)}")
+        return True, shm_write_bytes(payload, name=None)
     return False, obj
 
 

vllm_omni/entrypoints/utils.py

@@ -221,6 +221,7 @@ def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict | None
         base_engine_args = {}
     config_data = OmegaConf.load(config_path)
     stage_args = config_data.stage_args
+    global_async_chunk = config_data.get("async_chunk", False)
     # Convert any nested dataclass objects to dicts before creating OmegaConf
     base_engine_args = _convert_dataclasses_to_dict(base_engine_args)
     base_engine_args = OmegaConf.create(base_engine_args)
@@ -234,6 +235,7 @@ def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict | None
             runtime_cfg = stage_arg.runtime
             max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1)
             base_engine_args_tmp["max_num_seqs"] = max_batch_size
+            base_engine_args_tmp.async_chunk = global_async_chunk
         stage_arg.engine_args = base_engine_args_tmp
     return stage_args
 

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -180,6 +180,8 @@ class Qwen3OmniMoeForConditionalGeneration(
             self.thinker.make_empty_intermediate_tensors if self.model_stage == "thinker" else lambda: None
         )
 
+        self.chunk_segment_info = {}  # request_id -> chunk_segment_info
+
     # ==================== Device utilities ====================
 
     @staticmethod
@@ -483,12 +485,19 @@ class Qwen3OmniMoeForConditionalGeneration(
             talker_codes = talker_codes.unsqueeze(0).unsqueeze(0)
             talker_codes = talker_codes.expand(1, 16, -1)
 
-        # Use chunked decode for memory efficiency
-        audio_tensor = self.code2wav.chunked_decode(
-            talker_codes,
-            chunk_size=300,
-            left_context_size=25,
-        )
+        if self.vllm_config.model_config.async_chunk:
+            audio_tensor = self.code2wav.chunked_decode_streaming(
+                talker_codes,
+                chunk_size=25,
+                left_context_size=25,
+            )
+        else:
+            # Use chunked decode for memory efficiency
+            audio_tensor = self.code2wav.chunked_decode(
+                talker_codes,
+                chunk_size=300,
+                left_context_size=25,
+            )
 
         return audio_tensor
 
@@ -595,7 +604,6 @@ class Qwen3OmniMoeForConditionalGeneration(
             last_talker_hidden, text_step, update_dict = self.talker_preprocess_decode(
                 input_ids, input_embeds, **info_dict
             )
-
             update_dict["mtp_inputs"] = last_talker_hidden, text_step
 
         return input_ids, input_embeds, update_dict
@@ -624,6 +632,33 @@ class Qwen3OmniMoeForConditionalGeneration(
         inputs_embeds = (inputs_embeds + text_step).reshape(-1, self.talker_config.text_config.hidden_size)
         return inputs_embeds, code_predictor_codes.squeeze(-1)
 
+    def _get_tts_embed(self, thinker_embed, tts_bos_thinker, tts_eos_thinker, tts_pad_thinker):
+        """Project thinker-side TTS embeddings into talker text space."""
+        module_device = self._module_device(self.talker)
+
+        def _ensure_1x1(x: torch.Tensor) -> torch.Tensor:
+            if x.ndim == 3:
+                return x[0, -1:, :]
+            if x.ndim == 2:
+                return x[-1]
+            return x.view(1, 1, -1)
+
+        def _proj_from_thinker(x_opt: torch.Tensor | None) -> torch.Tensor:
+            if isinstance(x_opt, torch.Tensor) and x_opt.numel() > 0:
+                xin = _ensure_1x1(x_opt).to(module_device)
+            else:
+                xin = torch.zeros(
+                    (1, thinker_embed.shape[-1]),
+                    device=module_device,
+                    dtype=thinker_embed.dtype,
+                )
+            return self.talker.text_projection(xin).to(module_device)
+
+        self.tts_bos_embed = _proj_from_thinker(tts_bos_thinker)
+        self.tts_eos_embed = _proj_from_thinker(tts_eos_thinker)
+        self.tts_pad_embed = _proj_from_thinker(tts_pad_thinker)
+        return self.tts_bos_embed, self.tts_eos_embed, self.tts_pad_embed
+
     def talker_preprocess_prefill(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict):
         # Containers to return per-request updates (e.g., code_predictor_hidden_per_request)
         update_dict: dict[str, dict] = {}
@@ -647,7 +682,7 @@ class Qwen3OmniMoeForConditionalGeneration(
             if info_dict.get("thinker_input_ids") is None
             else torch.as_tensor(info_dict.get("thinker_input_ids"), device=self._module_device(self.talker))
         )
-        # Optional thinker-side TTS token embeddings (to be projected by talker)
+
         tts_bos_thinker = info_dict.get("tts_bos_embed").to(
             device=self._module_device(self.talker), dtype=torch.bfloat16
         )
@@ -700,7 +735,7 @@ class Qwen3OmniMoeForConditionalGeneration(
             tts_pad_thinker=tts_pad_thinker,
         )
 
-        # Queue tailing_text_hidden for decode (drop first for next steps),
+        # Queue trailing_text_hidden for decode (drop first for next steps),
         try:
             if isinstance(trailing_text_hidden, torch.Tensor) and trailing_text_hidden.numel() > 0:
                 if trailing_text_hidden.ndim == 2:
@@ -716,7 +751,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                     # compatible with old shape [1,S,D]
                     rem_tail = trailing_text_hidden.squeeze(0)
                 if rem_tail.shape[0] > 0:
-                    update_dict["tailing_text_hidden"] = rem_tail.detach().to("cpu").contiguous()
+                    update_dict["trailing_text_hidden"] = rem_tail.detach().to("cpu").contiguous()
             # Also persist projected tts_pad for decode fallback if needed
             if isinstance(tts_pad_thinker, torch.Tensor):
                 pad_in = tts_pad_thinker
@@ -762,31 +797,9 @@ class Qwen3OmniMoeForConditionalGeneration(
             (thinker_result_ids == self.thinker_config.video_token_id)
         ).to(input_ids.device)  # [t] # fmt: skip
 
-        # Project thinker-side TTS embeddings into talker text space (provided by thinker stage)
-        talker_dev = self._module_device(self.talker)
-
-        def _ensure_1x1(x: torch.Tensor) -> torch.Tensor:
-            if x.ndim == 3:
-                return x[0, -1:, :]
-            if x.ndim == 2:
-                return x[-1]
-            return x.view(1, 1, -1)
-
-        def _proj_from_thinker(x_opt: torch.Tensor | None) -> torch.Tensor:
-            if isinstance(x_opt, torch.Tensor) and x_opt.numel() > 0:
-                xin = _ensure_1x1(x_opt).to(talker_dev)
-            else:
-                xin = torch.zeros(
-                    (1, thinker_embed.shape[-1]),
-                    device=talker_dev,
-                    dtype=thinker_embed.dtype,
-                )
-            return self.talker.text_projection(xin).to(input_ids.device)
-
-        tts_bos_embed = _proj_from_thinker(tts_bos_thinker)
-        tts_eos_embed = _proj_from_thinker(tts_eos_thinker)
-        tts_pad_embed = _proj_from_thinker(tts_pad_thinker)
-        self.tts_pad_embed = tts_pad_embed
+        tts_bos_embed, tts_eos_embed, tts_pad_embed = self._get_tts_embed(
+            thinker_embed, tts_bos_thinker, tts_eos_thinker, tts_pad_thinker
+        )
 
         talker_input_embeds = []  # [1 t d]
         talker_input_ids = []
@@ -835,24 +848,58 @@ class Qwen3OmniMoeForConditionalGeneration(
 
         return talker_input_id, talker_input_embed, trailing_text_hidden_all
 
+    def _thinker_decode_to_talker_decode(
+        self,
+        info_dict: dict,
+        device: torch.device,
+        update_dict,
+    ):
+        """
+        Project thinker outputs to talker inputs during prefill stage.
+        Returns:
+            (input_ids, input_embeds) for talker
+        """
+        thinker_embed = info_dict.get("thinker_embeddings", None)
+        if thinker_embed is None:
+            if info_dict.get("finished_flag"):
+                return self.tts_pad_embed.to(device)
+            update_dict["finished_flag"] = True
+            return self.tts_eos_embed.to(device)
+
+        thinker_embed = thinker_embed.to(device)
+        return self.talker.text_projection(thinker_embed).to(device)
+
     def talker_preprocess_decode(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict):
         update_dict: dict[str, dict] = {}
+        last_talker_hidden = None
+        text_step = None
         try:
-            q_tail = info_dict.get("tailing_text_hidden")
-            if isinstance(q_tail, torch.Tensor) and q_tail.numel() > 0:
-                use_vec = q_tail[0:1, :]
-                new_q_tail = (
-                    q_tail[1:, :].detach().to("cpu").contiguous()
-                    if q_tail.shape[0] > 1
-                    else self.tts_pad_embed.to(input_embeds.device, dtype=input_embeds.dtype)
-                )
-                text_step = use_vec.to(input_embeds.device, dtype=input_embeds.dtype)
-                update_dict["tailing_text_hidden"] = new_q_tail
+            if self.vllm_config.model_config.async_chunk:
+                text_step = self._thinker_decode_to_talker_decode(info_dict, input_ids.device, update_dict)
             else:
-                text_step = self.tts_pad_embed.to(input_embeds.device, dtype=input_embeds.dtype)
+                q_tail = info_dict.get("trailing_text_hidden", None)
+                if isinstance(q_tail, torch.Tensor) and q_tail.numel() > 0:
+                    use_vec = q_tail[0:1, :]
+                    new_q_tail = (
+                        q_tail[1:, :].detach().to("cpu").contiguous()
+                        if q_tail.shape[0] > 1
+                        else self.tts_pad_embed.to(input_embeds.device, dtype=input_embeds.dtype)
+                    )
+                    text_step = use_vec.to(input_embeds.device, dtype=input_embeds.dtype)
+                    update_dict["trailing_text_hidden"] = new_q_tail
+                else:
+                    text_step = self.tts_pad_embed.to(input_embeds.device, dtype=input_embeds.dtype)
 
-            last_talker_hidden = info_dict.get("last_talker_hidden").to(input_embeds.device, dtype=input_embeds.dtype)
-            last_talker_hidden = last_talker_hidden.reshape(*last_talker_hidden.shape[-2:])  # [1, hidden_size]
+            last_talker_hidden_tensor = info_dict.get("last_talker_hidden")
+            if last_talker_hidden_tensor is not None:
+                last_talker_hidden = last_talker_hidden_tensor.to(input_embeds.device, dtype=input_embeds.dtype)
+                last_talker_hidden = last_talker_hidden.reshape(*last_talker_hidden.shape[-2:])  # [1, hidden_size]
+            else:
+                last_talker_hidden = torch.zeros(
+                    (1, self.talker_config.text_config.hidden_size),
+                    device=input_embeds.device,
+                    dtype=input_embeds.dtype,
+                )
         except Exception as e:
             logger.error(f"Error in decode: {e}")
 
@@ -866,7 +913,6 @@ class Qwen3OmniMoeForConditionalGeneration(
         )
 
         user_mm_mask = multimodal_mask[im_start_index:segment_end_index]
-
         # Multimodal data exists
         if user_mm_mask.any():
             user_thinker_hidden_mm = thinker_hidden[im_start_index:segment_end_index][user_mm_mask]
@@ -883,12 +929,20 @@ class Qwen3OmniMoeForConditionalGeneration(
         assistant_hidden = self.talker.text_projection(thinker_embed[im_start_index:segment_end_index]).to(
             tts_pad_embed.device
         )  # [t, d]
+
+        # [3 tokens] + [4 pad] + [1 BOS] + [1 first text] = 9 tokens
         assistant_text_hidden = torch.cat(
             (
                 assistant_hidden[:3],
                 tts_pad_embed.expand(4, -1),
                 tts_bos_embed,
-                assistant_hidden[3:4],  # First text
+                assistant_hidden[3:4]
+                if assistant_hidden.shape[0] > 3
+                else torch.zeros(
+                    (1, assistant_hidden.shape[1]),
+                    device=assistant_hidden.device,
+                    dtype=assistant_hidden.dtype,
+                ),  # First text
             ),
             dim=0,
         )
@@ -904,6 +958,9 @@ class Qwen3OmniMoeForConditionalGeneration(
             device=tts_pad_embed.device,
             dtype=torch.long,
         )
+        embed_input_ids = self.talker.embed_input_ids(codec_special_tokens).to(
+            device=tts_pad_embed.device, dtype=torch.bfloat16
+        )
         assistant_codec_hidden = torch.cat(
             (
                 torch.zeros(
@@ -911,18 +968,21 @@ class Qwen3OmniMoeForConditionalGeneration(
                     device=tts_pad_embed.device,
                     dtype=torch.bfloat16,
                 ),
-                self.talker.embed_input_ids(codec_special_tokens).to(device=tts_pad_embed.device, dtype=torch.bfloat16),
-            ),
-            dim=0,
-        )
-        trailing_text_hidden = torch.cat(
-            (
-                assistant_hidden[4:],
-                tts_eos_embed,
+                embed_input_ids,
             ),
             dim=0,
         )
 
+        if assistant_hidden.shape[0] > 4:
+            trailing_text_hidden = torch.cat(
+                (assistant_hidden[4:], tts_eos_embed),
+                dim=0,
+            )
+        else:
+            trailing_text_hidden = torch.zeros(
+                tts_eos_embed.shape, device=tts_eos_embed.device, dtype=tts_eos_embed.dtype
+            )
+
         input_embeds = assistant_text_hidden + assistant_codec_hidden
         input_ids = torch.full(
             (assistant_text_hidden.shape[0],),

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py

@@ -196,6 +196,39 @@ class Qwen3OmniMoeCode2Wav(nn.Module):
 
         return torch.cat(wavs, dim=-1)
 
+    def chunked_decode_streaming(
+        self,
+        codes: torch.Tensor,
+        chunk_size: int = 25,
+        left_context_size: int = 25,
+    ) -> torch.Tensor:
+        """
+        Decode long sequences in chunks to avoid OOM.
+
+        Uses overlapping chunks with left context to avoid boundary artifacts.
+
+        Args:
+            codes: [batch, num_quantizers, seq_len] - num_quantizers-layer RVQ codes
+            chunk_size: Number of codec frames per chunk
+            left_context_size: Number of overlapping frames for context
+
+        Returns:
+            waveform: [batch, 1, waveform_len] - Complete waveform
+        """
+        wavs = []
+        end_index = codes.shape[-1]
+        # TODO: need to optimize algorithms, current only support
+        # chunk_size = left_context_size = 25
+        if end_index == chunk_size:
+            context_size = 0
+        else:
+            context_size = left_context_size
+        # Decode chunk
+        wav_chunk = self(codes)
+        # Remove context from output (context_size * total_upsample samples)
+        wavs.append(wav_chunk[..., context_size * self.total_upsample :])
+        return torch.cat(wavs, dim=-1)
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load weights from HuggingFace checkpoint."""
         loader = AutoWeightsLoader(

vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml

@@ -4,12 +4,13 @@
 # Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
 
 # The following config has been verified on 2x H100-80G GPUs.
+async_chunk: false
 stage_args:
   - stage_id: 0
     stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
       devices: "0,1"
-      max_batch_size: 1
+      max_batch_size: 64
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
@@ -40,17 +41,16 @@ stage_args:
     stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
       devices: "1"
-      max_batch_size: 1
+      max_batch_size: 64
     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.35
+      gpu_memory_utilization: 0.3
       enforce_eager: true
       trust_remote_code: true
       engine_output_type: latent  # Output codec codes for code2wav
-      # tensor_parallel_size: 2
       enable_prefix_caching: false
       max_num_batched_tokens: 32768
       distributed_executor_backend: "mp"

vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml

@@ -0,0 +1,101 @@
+# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
+# Stage 0: Thinker (multimodal understanding + text generation)
+# Stage 1: Talker (text embeddings → 16-layer RVQ codec codes)
+# Stage 2: Code2Wav (16-layer RVQ codes → audio waveform)
+
+# The following config has been verified on 2x H100-80G GPUs.
+async_chunk: true
+stage_args:
+  - stage_id: 0
+    stage_type: llm  # Use llm stage type to launch OmniLLM
+    runtime:
+      devices: "0,1"
+      max_batch_size: 64
+    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: false
+      trust_remote_code: true
+      engine_output_type: latent  # Output hidden states for talker
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      hf_config_name: thinker_config
+      tensor_parallel_size: 2
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk
+    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: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    stage_type: llm  # Use llm stage type to launch OmniLLM
+    runtime:
+      devices: "0"
+      max_batch_size: 64
+    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.3
+      enforce_eager: false
+      trust_remote_code: true
+      engine_output_type: latent  # Output codec codes for code2wav
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      distributed_executor_backend: "mp"
+      hf_config_name: talker_config
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk
+    engine_input_source: [0]
+    # final_output: true
+    # final_output_type: text
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      detokenize: False
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+
+  - stage_id: 2
+    stage_type: llm  # Use llm stage type to launch OmniLLM
+    runtime:
+      devices: "1"
+      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
+      async_scheduling: false
+      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: 10000
+      hf_config_name: thinker_config
+    engine_input_source: [1]
+    final_output: true
+    final_output_type: audio
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.1

vllm_omni/model_executor/stage_input_processors/qwen3_omni.py

@@ -9,6 +9,7 @@ import torch
 from vllm.inputs import TextPrompt
 from vllm.platforms import current_platform
 
+from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.inputs.data import OmniTokensPrompt
 
 
@@ -48,6 +49,74 @@ def _compute_talker_prompt_ids_length(info, device: torch.device | str = "cuda")
     return sum_user_len + assistant_len
 
 
+# =========================
+# Common helpers
+# =========================
+
+
+def _ensure_list(x):
+    """Convert ConstantList / tensor-like to Python list."""
+    if hasattr(x, "_x"):
+        return list(x._x)
+    elif not isinstance(x, list):
+        return x
+    return list(x)
+
+
+def _validate_stage_inputs(stage_list, engine_input_source):
+    if not engine_input_source:
+        raise ValueError("engine_input_source cannot be empty")
+
+    stage_id = engine_input_source[0]
+    if stage_id >= len(stage_list):
+        raise IndexError(f"Invalid stage_id: {stage_id}")
+
+    stage = stage_list[stage_id]
+    if stage.engine_outputs is None:
+        raise RuntimeError(f"Stage {stage_id} has no outputs yet")
+
+    return stage.engine_outputs
+
+
+# =========================
+# Thinker -> Talker
+# =========================
+
+
+def thinker2talker_async_chunk(
+    pooling_output: dict[str, Any],
+    request: OmniEngineCoreRequest,
+) -> list[dict[str, Any]]:
+    """
+    Process thinker outputs to create talker inputs.
+    1. thinker's text generation outputs (token IDs + hidden states)
+    2. Split hidden states into: prompt embeddings + generated embeddings
+    3. Package for talker with additional information
+    """
+    all_token_ids = request.all_token_ids  # prefill + decode
+    prompt_token_ids = request.prompt_token_ids
+
+    # Convert ConstantList to regular list for OmniSerializer serialization
+    all_token_ids = _ensure_list(all_token_ids)
+    prompt_token_ids = _ensure_list(prompt_token_ids)
+
+    thinker_output = pooling_output
+
+    talker_additional_info = {
+        "thinker_embeddings": thinker_output.get("0").detach().cpu(),
+        "thinker_hidden_states": thinker_output.get("24").detach().cpu(),
+        "thinker_sequences": all_token_ids,
+        "thinker_input_ids": prompt_token_ids,
+        # Provide thinker-side TTS token embeddings for talker projection
+        "tts_bos_embed": thinker_output.get("tts_bos_embed").detach().cpu(),
+        "tts_eos_embed": thinker_output.get("tts_eos_embed").detach().cpu(),
+        "tts_pad_embed": thinker_output.get("tts_pad_embed").detach().cpu(),
+        "finished": torch.tensor(request.is_finished(), dtype=torch.bool),
+    }
+
+    return talker_additional_info
+
+
 def thinker2talker(
     stage_list: list[Any],
     engine_input_source: list[int],
@@ -71,47 +140,33 @@ def thinker2talker(
     Returns:
         List of OmniTokensPrompt for talker stage
     """
-    if not engine_input_source:
-        raise ValueError("engine_input_source cannot be empty")
-
-    source_stage_id = engine_input_source[0]
-    if source_stage_id >= len(stage_list):
-        raise IndexError(f"Invalid stage_id: {source_stage_id}")
-
-    if stage_list[source_stage_id].engine_outputs is None:
-        raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
-
-    thinker_outputs = stage_list[source_stage_id].engine_outputs
-    talker_inputs = []
+    thinker_outputs = _validate_stage_inputs(stage_list, engine_input_source)
+    talker_inputs: list[OmniTokensPrompt] = []
 
     device = torch.device(current_platform.device_type)
 
     # Process each thinker output
-    for i, thinker_output in enumerate(thinker_outputs):
+    for thinker_output in thinker_outputs:
         output = thinker_output.outputs[0]
-        thinker_embeddings = output.multimodal_output["0"].detach().to(device=device, dtype=torch.float)
 
-        thinker_hidden_states = output.multimodal_output["24"].detach().to(device=device, dtype=torch.float)
         info = {
-            "thinker_embeddings": thinker_embeddings,
-            "thinker_hidden_states": thinker_hidden_states,
-            "thinker_sequences": thinker_output.prompt_token_ids
-            + output.token_ids,  # the thinker_sequences is the whole ids
+            "thinker_embeddings": output.multimodal_output["0"].detach().to(device=device, dtype=torch.float),
+            "thinker_hidden_states": output.multimodal_output["24"].detach().to(device=device, dtype=torch.float),
+            "thinker_sequences": (
+                thinker_output.prompt_token_ids + output.token_ids
+            ),  # the thinker_sequences is the whole ids
             "thinker_input_ids": thinker_output.prompt_token_ids,
             # Provide thinker-side TTS token embeddings for talker projection
-            "tts_bos_embed": (
-                output.multimodal_output.get("tts_bos_embed").detach().to(device=device, dtype=torch.float)
-            ),
-            "tts_eos_embed": (
-                output.multimodal_output.get("tts_eos_embed").detach().to(device=device, dtype=torch.float)
-            ),
-            "tts_pad_embed": (
-                output.multimodal_output.get("tts_pad_embed").detach().to(device=device, dtype=torch.float)
-            ),
+            "tts_bos_embed": output.multimodal_output["tts_bos_embed"].detach().to(device=device, dtype=torch.float),
+            "tts_eos_embed": output.multimodal_output["tts_eos_embed"].detach().to(device=device, dtype=torch.float),
+            "tts_pad_embed": output.multimodal_output["tts_pad_embed"].detach().to(device=device, dtype=torch.float),
         }
+
+        prompt_len = _compute_talker_prompt_ids_length(info, device=device)
+
         talker_inputs.append(
             OmniTokensPrompt(
-                prompt_token_ids=[0] * _compute_talker_prompt_ids_length(info, device=device),
+                prompt_token_ids=[0] * prompt_len,
                 additional_information=info,
                 multi_modal_data=None,
                 mm_processor_kwargs=None,
@@ -121,6 +176,50 @@ def thinker2talker(
     return talker_inputs
 
 
+# =========================
+# Talker -> Code2Wav
+# =========================
+
+
+def talker2code2wav_async_chunk(
+    pooling_output: dict[str, Any],
+    request: OmniEngineCoreRequest,
+):
+    """
+    Pooling version.
+    """
+    if "code_predictor_codes" not in pooling_output:
+        return []
+
+    code_predictor_codes = pooling_output["code_predictor_codes"]
+
+    if code_predictor_codes is None:
+        return []
+    if isinstance(code_predictor_codes, torch.Tensor):
+        if code_predictor_codes.numel() == 0:
+            return []
+    elif hasattr(code_predictor_codes, "__len__"):
+        if len(code_predictor_codes) == 0:
+            return []
+
+    if isinstance(code_predictor_codes, torch.Tensor):
+        if not code_predictor_codes.any():
+            return []
+    else:
+        code_tensor = torch.tensor(code_predictor_codes, dtype=torch.long)
+        if not code_tensor.any():
+            return []
+
+    codec_codes = code_predictor_codes.to(torch.long).transpose(0, 1).cpu().to(torch.long).reshape(-1).tolist()
+    if sum(codec_codes) == 0:
+        return []
+
+    return {
+        "code_predictor_codes": codec_codes,
+        "finished": torch.tensor(request.is_finished(), dtype=torch.bool),
+    }
+
+
 def talker2code2wav(
     stage_list: list[Any],
     engine_input_source: list[int],
@@ -144,21 +243,10 @@ def talker2code2wav(
     Returns:
         List of OmniTokensPrompt for code2wav stage
     """
-    if not engine_input_source:
-        raise ValueError("engine_input_source cannot be empty")
-
-    source_stage_id = engine_input_source[0]
-    if source_stage_id >= len(stage_list):
-        raise IndexError(f"Invalid stage_id: {source_stage_id}")
-
-    if stage_list[source_stage_id].engine_outputs is None:
-        raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
-
-    talker_outputs = stage_list[source_stage_id].engine_outputs
-    code2wav_inputs = []
-
+    talker_outputs = _validate_stage_inputs(stage_list, engine_input_source)
+    code2wav_inputs: list[OmniTokensPrompt] = []
     # Process each talker output
-    for i, talker_output in enumerate(talker_outputs):
+    for talker_output in talker_outputs:
         output = talker_output.outputs[0]
         seq_len = len(output.token_ids) - 1
         # Extract codec codes from talker output

vllm_omni/worker/gpu_ar_model_runner.py

@@ -280,7 +280,6 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 aux_hidden_states = None
 
             hidden_states, multimodal_outputs = self.extract_multimodal_outputs(model_output)
-
             if multimodal_outputs is not None:
                 keys_or_type = (
                     list(multimodal_outputs.keys())
@@ -501,7 +500,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 dtype=np.int32,
             )
 
-        self._process_additional_information_updates(hidden_states, multimodal_outputs, num_scheduled_tokens_np)
+        self._process_additional_information_updates(
+            hidden_states, multimodal_outputs, num_scheduled_tokens_np, scheduler_output
+        )
 
         pooler_output: list[dict[str, object]] = []
         for rid in req_ids_output_copy:
@@ -761,8 +762,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
 
         for attempt in range(max_retries):
             try:
+                put_key = f"omni_{from_stage}_to_{to_stage}_{request_id}"
                 success, size, metadata = connector.put(
-                    from_stage=from_stage, to_stage=to_stage, request_id=request_id, data=data
+                    from_stage=from_stage, to_stage=to_stage, put_key=put_key, data=data
                 )
                 if success:
                     return success, size, metadata

vllm_omni/worker/gpu_generation_model_runner.py

@@ -50,6 +50,17 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
     - Executes generation process and returns tensors via `pooler_output`.
     """
 
+    def _update_request_states(self, scheduler_output: SchedulerOutput):
+        cached_reqs = scheduler_output.scheduled_cached_reqs
+        for _, req_id in enumerate(cached_reqs.req_ids):
+            req_state = self.requests.get(req_id)
+            assert req_state is not None
+            req_state.prompt_token_ids = cached_reqs.prompt_token_ids.get(req_id)
+            self.input_batch.remove_request(req_id)
+            # update the request state in self.input_batch
+            self.input_batch.add_request(req_state)
+            self._init_mrope_positions(req_state)
+
     @torch.inference_mode()
     def execute_model(
         self,
@@ -74,6 +85,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             record_function_or_nullcontext("gpu_model_runner: preprocess"),
             self.synchronize_input_prep(),
         ):
+            if self.model_config.async_chunk:
+                self._update_request_states(scheduler_output)
             self._update_states(scheduler_output)
             if not scheduler_output.total_num_scheduled_tokens:
                 return EMPTY_MODEL_RUNNER_OUTPUT

vllm_omni/worker/gpu_model_runner.py

@@ -808,6 +808,7 @@ class OmniGPUModelRunner(GPUModelRunner):
         hidden_states: torch.Tensor,
         multimodal_outputs: object,
         num_scheduled_tokens_np: np.ndarray,
+        scheduler_output: "SchedulerOutput",
     ) -> None:
         """Process model-provided per-request additional_information updates and merge into request state."""
         try:
@@ -815,10 +816,13 @@ class OmniGPUModelRunner(GPUModelRunner):
             # TODO(Peiqi): do we have a more elegant way to do this?
             if hasattr(self.model, "has_postprocess") and self.model.has_postprocess:
                 for req_index, req_id in enumerate(self.input_batch.req_ids):
-                    req_state = self.requests.get(req_id)
-                    req_infos = (
-                        getattr(req_state, "additional_information_cpu", None) if req_state is not None else None
-                    )
+                    if self.model_config.async_chunk:
+                        req_infos = self._get_additional_information(scheduler_output, req_id)
+                    else:
+                        req_state = self.requests.get(req_id)
+                        req_infos = (
+                            getattr(req_state, "additional_information_cpu", None) if req_state is not None else None
+                        )
                     start_offset = int(self.query_start_loc.cpu[req_index])
                     sched_tokens = int(num_scheduled_tokens_np[req_index])
                     s, e = start_offset, start_offset + sched_tokens
@@ -859,6 +863,40 @@ class OmniGPUModelRunner(GPUModelRunner):
                 start_offset = int(self.query_start_loc.cpu[req_index])
                 self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src)
 
+    def _get_additional_information(self, scheduler_output: "SchedulerOutput", req_id: str) -> dict:
+        req_infos = None
+        req_state = self.requests.get(req_id)
+        additional_information_cpu = getattr(req_state, "additional_information_cpu", None)
+        for new_req in scheduler_output.scheduled_new_reqs:
+            if new_req.req_id == req_id:
+                payload_info = getattr(new_req, "additional_information", None)
+                if payload_info is not None:
+                    return payload_info
+
+        if hasattr(scheduler_output.scheduled_cached_reqs, "additional_information"):
+            cached_infos = getattr(scheduler_output.scheduled_cached_reqs, "additional_information", {})
+            if isinstance(cached_infos, dict) and req_id in cached_infos:
+                req_infos = cached_infos[req_id]
+                if not isinstance(req_infos, dict):
+                    req_infos = None
+
+        if req_infos is None or req_infos.get("last_talker_hidden", None) is None:
+            if req_infos is None:
+                additional_information_cpu.pop("thinker_embeddings", None)
+                req_infos = additional_information_cpu
+            else:
+                req_infos["last_talker_hidden"] = additional_information_cpu.get("last_talker_hidden", None)
+                req_infos["num_processed_thinker_tokens"] = additional_information_cpu.get(
+                    "num_processed_thinker_tokens", 0
+                )
+            if not isinstance(req_infos, dict):
+                req_infos = None
+
+        if req_infos is None:
+            logger.warning(f"No additional_information found for req_id: {req_id}")
+
+        return req_infos
+
     def _preprocess(
         self,
         scheduler_output: "SchedulerOutput",
@@ -974,9 +1012,14 @@ class OmniGPUModelRunner(GPUModelRunner):
             # collect additional_information (tensor/list) for prefill portion only
             decode_req_ids = []
             for req_index, req_id in enumerate(self.input_batch.req_ids):
-                req_state = self.requests.get(req_id)
-                req_infos = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None
-
+                # Try to get additional_information from multiple sources
+                if self.vllm_config.model_config.async_chunk:
+                    req_infos = self._get_additional_information(scheduler_output, req_id)
+                else:
+                    req_state = self.requests.get(req_id)
+                    req_infos = (
+                        getattr(req_state, "additional_information_cpu", None) if req_state is not None else None
+                    )
                 start_offset = int(self.query_start_loc.cpu[req_index])
                 sched_tokens = int(num_scheduled_tokens_np[req_index])
                 s, e = start_offset, start_offset + sched_tokens