From 5304bc1e16bc18573dacd053d89d3b475aa69552 Mon Sep 17 00:00:00 2001 From: amy-why-3459 Date: Wed, 4 Feb 2026 00:00:56 +0800 Subject: [PATCH] [Refactor] Refactor async chunk and fix the shape mismatch issue (#1151) Signed-off-by: amy-why-3459 --- .../distributed/omni_connectors/adapter.py | 71 +++++--------- .../models/qwen3_omni/qwen3_omni.py | 15 ++- .../qwen3_omni_moe_async_chunk.yaml | 2 +- .../stage_input_processors/qwen3_omni.py | 95 +++++++++++++------ .../platforms/npu/worker/npu_model_runner.py | 58 +++-------- .../worker/gpu_generation_model_runner.py | 4 +- vllm_omni/worker/gpu_model_runner.py | 59 +++--------- 7 files changed, 138 insertions(+), 166 deletions(-) diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py index 977ecc88b..4cd225078 100644 --- a/vllm_omni/distributed/omni_connectors/adapter.py +++ b/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 ) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 37753ee47..e2967cf9b 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/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): diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml index c8cb67dd5..1a8b60aa1 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml +++ b/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 diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index 03daa8e42..701923309 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/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( diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index a9a376f95..83f476b3c 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/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) diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 9125cc52a..2c82b5cd4 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/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( diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index cb14ed49a..d52a2014f 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/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)