[Refactor] Refactor async chunk and fix the shape mismatch issue (#1151)

Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>

vllm_omni/distributed/omni_connectors/adapter.py

@@ -179,6 +179,27 @@ def try_recv_via_connector(
             return None, None
 
 
+def update_request_payload(connector: "OmniConnectorBase", req_id: str, payload_data: dict[str, Any]) -> dict[str, Any]:
+    """Update the payload data for a request in the connector.
+
+    Args:
+        connector: OmniConnectorBase instance
+        req_id: Request ID to update
+        payload_data: New payload data to store
+    """
+    origin_payload = connector.request_payload[req_id]
+    for key, value in payload_data.items():
+        if key == "finished":
+            continue
+        elif isinstance(value, torch.Tensor) and key in origin_payload:
+            payload_data[key] = torch.cat([origin_payload[key], value], dim=0)
+        elif isinstance(value, list) and key in origin_payload:
+            payload_data[key] = origin_payload[key] + value
+
+    connector.request_payload[req_id] = payload_data
+    return payload_data
+
+
 def get_chunk(
     connector: "OmniConnectorBase",
     scheduler_output: SchedulerOutput,
@@ -206,6 +227,7 @@ def get_chunk(
         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
+            connector.request_payload[req_id] = payload_data
             if payload_data.get("finished"):
                 connector.finished_requests.add(req_id)
 
@@ -222,15 +244,13 @@ def get_chunk(
         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:
+            payload_data = update_request_payload(connector, req_id, 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):
@@ -280,12 +300,8 @@ def get_chunk_for_generation(
     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", [])
+    request.prompt_token_ids = payload_data.get("code_predictor_codes", [])
+    request.num_computed_tokens = 0
 
 
 def put_chunk(
@@ -308,8 +324,6 @@ def put_chunk(
     stage_id = connector.stage_id
     next_stage_id = stage_id + 1
     request_id = request.external_req_id
-    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
@@ -318,6 +332,7 @@ def put_chunk(
     if custom_process_input_func:
         try:
             payload_data = custom_process_input_func(
+                connector=connector,
                 pooling_output=pooling_output,
                 request=request,
             )
@@ -325,42 +340,8 @@ def put_chunk(
             logger.error(f"Failed to use custom_process_input_func for payload extraction: {e}")
 
         if not payload_data:
-            logger.warning("[Stage-%d] No payload data to send for request %s", stage_id, request_id)
             return
 
-        if stage_id == 0 and chunk_id == 0:
-            if connector.request_payload.get(request_id) is None:
-                if not payload_data.get("finished"):
-                    connector.request_payload[request_id] = payload_data
-                    return
-            else:
-                save_payload = connector.request_payload.pop(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.debug("[Stage-%d] Merged embeddings and hidden states for request %s", stage_id, 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
         )

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -599,11 +599,16 @@ class Qwen3OmniMoeForConditionalGeneration(
             )
             update_dict["code_predictor_codes"] = code_predictor_codes
         else:
+            # decode
+            if info_dict.get("num_processed_tokens", 0) < len(info_dict.get("thinker_input_ids", [])):
+                info_dict["num_processed_tokens"] = len(info_dict.get("thinker_input_ids", [])) + 1
+
             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
 
+        update_dict["num_processed_tokens"] = info_dict.get("num_processed_tokens", 0) + span_len
         return input_ids, input_embeds, update_dict
 
     def talker_mtp(
@@ -662,7 +667,8 @@ class Qwen3OmniMoeForConditionalGeneration(
         update_dict: dict[str, dict] = {}
         # TODO(Peiqi): add voice_type support
         voice_type = self.voice_type
-
+        start_index = info_dict.get("num_processed_tokens", 0)
+        end_index = start_index + input_embeds.shape[0]
         # Read thinker outputs for prefill
         thinker_sequence_embeds = info_dict.get("thinker_embeddings").to(
             device=self._module_device(self.talker), dtype=torch.bfloat16
@@ -762,7 +768,7 @@ class Qwen3OmniMoeForConditionalGeneration(
         except Exception:
             pass
 
-        return req_input_ids, req_embeds, update_dict
+        return req_input_ids[start_index:end_index], req_embeds[start_index:end_index], update_dict
 
     def _thinker_to_talker_prefill(
         self,
@@ -858,13 +864,14 @@ class Qwen3OmniMoeForConditionalGeneration(
             (input_ids, input_embeds) for talker
         """
         thinker_embed = info_dict.get("thinker_embeddings", None)
-        if thinker_embed is None:
+        start_index = info_dict.get("num_processed_tokens", 0)
+        if start_index >= thinker_embed.shape[0]:
             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)
+        thinker_embed = thinker_embed[start_index : start_index + 1].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):

vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml

@@ -63,7 +63,7 @@ stage_args:
     default_sampling_params:
       temperature: 0.9
       top_k: 50
-      max_tokens: 2048 # TODO: The max_tokens of the async_chunk feature cannot exceed 2048.
+      max_tokens: 4096
       seed: 42
       detokenize: False
       repetition_penalty: 1.05

vllm_omni/model_executor/stage_input_processors/qwen3_omni.py

@@ -84,6 +84,7 @@ def _validate_stage_inputs(stage_list, engine_input_source):
 
 
 def thinker2talker_async_chunk(
+    connector: Any,
     pooling_output: dict[str, Any],
     request: OmniEngineCoreRequest,
 ) -> list[dict[str, Any]]:
@@ -93,27 +94,50 @@ def thinker2talker_async_chunk(
     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),
-    }
+    request_id = request.external_req_id
+    chunk_id = connector.put_requests[request_id]
+    if chunk_id == 0:
+        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)
+        talker_additional_info = {
+            "thinker_embeddings": pooling_output.get("0").detach().cpu(),
+            "thinker_hidden_states": pooling_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": pooling_output.get("tts_bos_embed").detach().cpu(),
+            "tts_eos_embed": pooling_output.get("tts_eos_embed").detach().cpu(),
+            "tts_pad_embed": pooling_output.get("tts_pad_embed").detach().cpu(),
+            "finished": torch.tensor(request.is_finished(), dtype=torch.bool),
+        }
+        if connector.request_payload.get(request_id) is None:
+            if not request.is_finished():
+                connector.request_payload[request_id] = talker_additional_info
+                return None
+        else:
+            save_payload = connector.request_payload.pop(request_id)
+            talker_additional_info["thinker_embeddings"] = torch.cat(
+                (save_payload.get("thinker_embeddings"), talker_additional_info.get("thinker_embeddings")), dim=0
+            )
+            talker_additional_info["thinker_hidden_states"] = torch.cat(
+                (save_payload.get("thinker_hidden_states"), talker_additional_info.get("thinker_hidden_states")),
+                dim=0,
+            )
+    else:
+        output_token_ids = request.output_token_ids
+        # Convert ConstantList to regular list for OmniSerializer serialization
+        output_token_ids = _ensure_list(output_token_ids)
 
+        talker_additional_info = {
+            "thinker_embeddings": pooling_output.get("0").detach().cpu(),
+            "thinker_hidden_states": pooling_output.get("24").detach().cpu(),
+            "thinker_sequences": output_token_ids,
+            "finished": torch.tensor(request.is_finished(), dtype=torch.bool),
+        }
     return talker_additional_info
 
 
@@ -182,6 +206,7 @@ def thinker2talker(
 
 
 def talker2code2wav_async_chunk(
+    connector: Any,
     pooling_output: dict[str, Any],
     request: OmniEngineCoreRequest,
 ):
@@ -189,35 +214,49 @@ def talker2code2wav_async_chunk(
     Pooling version.
     """
     if "code_predictor_codes" not in pooling_output:
-        return []
+        return None
 
     code_predictor_codes = pooling_output["code_predictor_codes"]
 
     if code_predictor_codes is None:
-        return []
+        return None
     if isinstance(code_predictor_codes, torch.Tensor):
         if code_predictor_codes.numel() == 0:
-            return []
+            return None
     elif hasattr(code_predictor_codes, "__len__"):
         if len(code_predictor_codes) == 0:
-            return []
+            return None
 
     if isinstance(code_predictor_codes, torch.Tensor):
         if not code_predictor_codes.any():
-            return []
+            return None
     else:
         code_tensor = torch.tensor(code_predictor_codes, dtype=torch.long)
         if not code_tensor.any():
-            return []
+            return None
 
     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 None
 
-    return {
-        "code_predictor_codes": codec_codes,
+    request_id = request.external_req_id
+    chunk_size = left_context_size = 25
+    connector.code_prompt_token_ids[request_id].append(codec_codes)
+    length = len(connector.code_prompt_token_ids[request_id])
+    chunk_length = length % chunk_size
+    if chunk_length != 0 and not request.is_finished():
+        return None
+
+    context_length = chunk_length if chunk_length != 0 else chunk_size
+    end_index = min(length, left_context_size + context_length)
+
+    info = {
+        "code_predictor_codes": (
+            torch.tensor(connector.code_prompt_token_ids[request_id][-end_index:]).transpose(0, 1).reshape(-1).tolist()
+        ),
         "finished": torch.tensor(request.is_finished(), dtype=torch.bool),
     }
+    return info
 
 
 def talker2code2wav(

vllm_omni/platforms/npu/worker/npu_model_runner.py

@@ -733,13 +733,10 @@ class OmniNPUModelRunner(NPUModelRunner):
             # 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):
-                    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
-                        )
+                    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
@@ -780,39 +777,16 @@ class OmniNPUModelRunner(NPUModelRunner):
                 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)
+    def _update_additional_information(self, scheduler_output: "SchedulerOutput") -> 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
+            payload_info = getattr(new_req, "additional_information", None)
+            self._merge_additional_information_update(new_req.req_id, 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
+            if isinstance(cached_infos, dict):
+                for req_id, req_infos in cached_infos.items():
+                    self._merge_additional_information_update(req_id, req_infos)
 
     def _preprocess(
         self,
@@ -928,15 +902,12 @@ class OmniNPUModelRunner(NPUModelRunner):
             # Overlay custom prompt_embeds per request for the prompt portion;
             # collect additional_information (tensor/list) for prefill portion only
             decode_req_ids = []
+            if self.vllm_config.model_config.async_chunk:
+                self._update_additional_information(scheduler_output)
             for req_index, req_id in enumerate(self.input_batch.req_ids):
                 # 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
-                    )
+                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
@@ -946,6 +917,7 @@ class OmniNPUModelRunner(NPUModelRunner):
                 req_input_ids, req_embeds, update_dict = self.model.preprocess(
                     input_ids=input_ids[s:e], input_embeds=inputs_embeds[s:e], **req_infos
                 )
+
                 if hasattr(self.model, "talker_mtp") and span_len == 1:
                     last_talker_hidden, text_step = update_dict.pop("mtp_inputs")
                     decode_slice = slice(len(decode_req_ids), len(decode_req_ids) + 1)

vllm_omni/worker/gpu_generation_model_runner.py

@@ -58,7 +58,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             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)
+            # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
+            if self.uses_mrope:
+                self._init_mrope_positions(req_state)
 
     @torch.inference_mode()
     def execute_model(

vllm_omni/worker/gpu_model_runner.py

@@ -832,13 +832,10 @@ 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):
-                    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
-                        )
+                    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
@@ -879,39 +876,16 @@ 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)
+    def _update_additional_information(self, scheduler_output: "SchedulerOutput") -> 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
+            payload_info = getattr(new_req, "additional_information", None)
+            self._merge_additional_information_update(new_req.req_id, 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
+            if isinstance(cached_infos, dict):
+                for req_id, req_infos in cached_infos.items():
+                    self._merge_additional_information_update(req_id, req_infos)
 
     def _preprocess(
         self,
@@ -1027,15 +1001,11 @@ class OmniGPUModelRunner(GPUModelRunner):
             # Overlay custom prompt_embeds per request for the prompt portion;
             # collect additional_information (tensor/list) for prefill portion only
             decode_req_ids = []
+            if self.vllm_config.model_config.async_chunk:
+                self._update_additional_information(scheduler_output)
             for req_index, req_id in enumerate(self.input_batch.req_ids):
-                # 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
-                    )
+                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
@@ -1045,6 +1015,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                 req_input_ids, req_embeds, update_dict = self.model.preprocess(
                     input_ids=input_ids[s:e], input_embeds=inputs_embeds[s:e], **req_infos
                 )
+
                 if hasattr(self.model, "talker_mtp") and span_len == 1:
                     last_talker_hidden, text_step = update_dict.pop("mtp_inputs")
                     decode_slice = slice(len(decode_req_ids), len(decode_req_ids) + 1)