diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py index b650712b5..d3acaebf3 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py @@ -60,6 +60,7 @@ class Qwen2_5OmniForConditionalGeneration( ): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() + self.has_process = False self.have_multimodal_outputs = True config: Qwen2_5OmniConfig = vllm_config.model_config.hf_config multimodal_config = vllm_config.model_config.multimodal_config @@ -90,6 +91,8 @@ class Qwen2_5OmniForConditionalGeneration( self.token2wav = None elif self.model_stage == "talker": + # temporal use this prepresation + self.has_process = True self.thinker = None # Initialize talker model wrapper (handles projection + LM) self.talker = init_vllm_registered_model( @@ -236,13 +239,14 @@ class Qwen2_5OmniForConditionalGeneration( if inputs_embeds is not None and inputs_embeds.device != thinker_dev: inputs_embeds = inputs_embeds.to(thinker_dev) # Run thinker - thinker_output = self.thinker( - input_ids=input_ids, - positions=positions[0], - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, - **kwargs, - ) + with torch.inference_mode(): + thinker_output = self.thinker( + input_ids=input_ids, + positions=positions[0], + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **kwargs, + ) if isinstance(thinker_output, tuple): embeds, text_hidden_states = thinker_output @@ -257,151 +261,13 @@ class Qwen2_5OmniForConditionalGeneration( # 2) Talker (if codec not provided) if self.model_stage == "talker": - # Mixed-mode support: In a single step, both Prefill*n and Decode*n are supported. - # Rules: - # - Prefill segments are wrapped with special tokens: [BOS][PAD...][EOS] - # - Decode segments consist of a single non-special token. - # - If additional_information is provided - # (can be a list split by request or a concatenated tensor plus a list of shapes), - # then for each request, reconstruct the thinker→talker input embeddings for the Prefill segments; - # - For Decode segments, if per-request auxiliary decode embeddings are provided (optional), add them; - # otherwise, keep the original embedding. - + # mock data for profile if input_ids is None and additional_information is None: input_ids = torch.zeros( - inputs_embeds.shape[0], - dtype=torch.long, - device=inputs_embeds.device, + inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device ) additional_information = {} self.thinker_reply_part = torch.zeros_like(inputs_embeds) - is_profile = True - else: - is_profile = False - - # Ensure we have base embeddings when only ids are provided - if inputs_embeds is None and input_ids is not None: - inputs_embeds = self.talker.get_input_embeddings(input_ids) - - # ------- Request-scoped additional information (no cross-request concat) ------- - request_ids: Optional[list[str]] = kwargs.get("request_ids") # ordered - request_token_spans: Optional[list[tuple[int, int]]] = kwargs.get("request_token_spans") - addi_by_req: Optional[dict] = kwargs.get("additional_information_by_req_id") - runtime_addi = kwargs.get("runtime_additional_information") - - # Normalize runtime_addi into a mapping by request_id for convenience - runtime_addi_by_req: dict[str, dict] = {} - if ( - isinstance(request_ids, list) - and isinstance(runtime_addi, list) - and len(runtime_addi) == len(request_ids) - ): - for i, rid in enumerate(request_ids): - if isinstance(rid, str) and isinstance(runtime_addi[i], dict): - runtime_addi_by_req[rid] = runtime_addi[i] - elif isinstance(request_ids, list) and isinstance(runtime_addi, dict): - for rid in request_ids: - if isinstance(rid, str) and isinstance(runtime_addi.get(rid), dict): - runtime_addi_by_req[rid] = runtime_addi[rid] - - # Containers to return per-request updates (e.g., thinker_reply_part_per_request) - update_by_req_id: dict[str, dict] = {} - - # ------- Prefill: span_len > 1 ------- - if ( - not is_profile - and isinstance(request_ids, list) - and isinstance(request_token_spans, list) - and isinstance(addi_by_req, dict) - ): - for idx_req, rid in enumerate(request_ids): - s, e = request_token_spans[idx_req] - span_len = int(e) - int(s) - if span_len <= 1: - continue - info = addi_by_req.get(rid, {}) if isinstance(rid, str) else {} - if not isinstance(info, dict): - info = {} - pe = info.get("prompt_embeds") # Tensor [P,H] - tr = info.get("thinker_result") # Tensor [K,H] - ptoks = info.get("prompt_token_ids") # list[int] - otoks = info.get("thinker_output_token_ids") # list[int] - - if not isinstance(pe, torch.Tensor): - pe = torch.zeros( - 0, - self.talker.config.hidden_size, - dtype=inputs_embeds.dtype, - device=self._module_device(self.model), - ) - if not isinstance(tr, torch.Tensor): - tr = torch.zeros( - 0, - self.talker.config.hidden_size, - dtype=inputs_embeds.dtype, - device=self._module_device(self.model), - ) - if not isinstance(ptoks, (list, torch.Tensor)): - ptoks = [] - if not isinstance(otoks, (list, torch.Tensor)): - otoks = [] - - req_input_ids, req_embeds = self._thinker_to_talker_prefill( - voice_type=voice_type, - output_prompt_embeds=tr.to(inputs_embeds.dtype).to(self._module_device(self.model)), - output_token_ids=otoks, - thinker_prompt_embeds=pe.to(inputs_embeds.dtype).to(self._module_device(self.model)), - prompt_token_ids=ptoks, - ) - seg_len = min(span_len, req_embeds.shape[0]) - inputs_embeds[s : s + seg_len] = req_embeds[:seg_len] - if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len: - input_ids[s : s + seg_len] = req_input_ids - - # Prepare per-request reply queue for subsequent decode: drop first row - if tr.ndim == 2 and tr.shape[0] > 0: - update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = ( - tr[1:].detach().to("cpu").contiguous() - ) - - # ------- Decode: span_len == 1 ------- - if not is_profile and isinstance(request_ids, list) and isinstance(request_token_spans, list): - for idx_req, rid in enumerate(request_ids): - s, e = request_token_spans[idx_req] - if (int(e) - int(s)) != 1: - continue - # choose step vector in priority order - step_vec = None - # A) runtime queue - q = None - if isinstance(rid, str): - q = runtime_addi_by_req.get(rid, {}).get("thinker_reply_part_per_request") - if isinstance(q, torch.Tensor) and q.numel() > 0: - step_vec = q[0:1] - new_q = q[1:].detach().to("cpu").contiguous() - update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = new_q - else: - # B) per-request provided decode vector (optional) - info = addi_by_req.get(rid, {}) if isinstance(addi_by_req, dict) else {} - dv = info.get("decode_output_prompt_embeds") if isinstance(info, dict) else None - if isinstance(dv, torch.Tensor) and dv.numel() > 0: - step_vec = dv[0:1] if dv.ndim == 2 else dv.view(1, -1) - elif ( - hasattr(self, "thinker_reply_part") - and isinstance(self.thinker_reply_part, torch.Tensor) - and self.thinker_reply_part.numel() > 0 - ): - # C) fallback shared pool - step_vec = self.thinker_reply_part[0:1] - self.thinker_reply_part = self.thinker_reply_part[1:] - - if isinstance(step_vec, torch.Tensor) and step_vec.numel() > 0: - one_id = input_ids[s : s + 1] - _, one_embed = self._thinker_to_talker_decode_one_step( - output_prompt_embeds=step_vec.to(inputs_embeds.dtype).to(self._module_device(self.model)), - output_token_ids=one_id, - ) - inputs_embeds[s] = one_embed[0] with torch.inference_mode(): talker_hidden = self.talker( @@ -409,10 +275,6 @@ class Qwen2_5OmniForConditionalGeneration( positions=positions[0], inputs_embeds=inputs_embeds, ) - multimodal_outputs: dict = None - # Return updates if any - if update_by_req_id: - multimodal_outputs = {"additional_information_update_by_req_id": update_by_req_id} if sampling_metadata is not None: # the padding token id is set to text model's pad token id, @@ -421,7 +283,7 @@ class Qwen2_5OmniForConditionalGeneration( return OmniOutput( text_hidden_states=talker_hidden, - multimodal_outputs=multimodal_outputs, + multimodal_outputs=None, ) if self.model_stage == "code2wav": @@ -733,6 +595,68 @@ class Qwen2_5OmniForConditionalGeneration( return talker_hf_config.tts_text_start_token_id return self.tts_text_spk_token_ids[voice_type] + def process_fn(self, + input_ids: torch.Tensor, + inputs_embeds: torch.Tensor, + **info_dict: object, + ): + # Mixed-mode support: In a single step, both Prefill*n and Decode*n are supported. + # Rules: + # - Prefill segments are wrapped with special tokens: [BOS][PAD...][EOS] + # - Decode segments consist of a single non-special token. + # - If additional_information is provided (can be a list split by request or a + # concatenated tensor plus a list of shapes), then for each request, reconstruct + # the thinker→talker input embeddings for the Prefill segments; + # - For Decode segments, if per-request auxiliary decode embeddings are provided (optional), + # add them; otherwise, keep the original embedding. + + # Ensure we have base embeddings when only ids are provided + if inputs_embeds is None and input_ids is not None: + inputs_embeds = self.talker.get_input_embeddings(input_ids) + + span_len = input_ids.shape[0] + if span_len > 1: + # prefill + return self.thinker_to_talker_process(input_ids, inputs_embeds, **info_dict) + else: + # decode + return self.thinker_to_talker_decode_one_step(input_ids, inputs_embeds, **info_dict) + + + def thinker_to_talker_process(self, + input_ids: torch.Tensor, + inputs_embeds: torch.Tensor, + **info_dict: object, + ): + update_dict = {} + + prompt_embeds = info_dict.get("prompt_embeds") # Tensor [P,H] + thinker_result = info_dict.get("thinker_result") # Tensor [K,H] + prompt_token_ids = info_dict.get("prompt_token_ids") # list[int] + thinker_output_token_ids = info_dict.get("thinker_output_token_ids") # list[int] + + if not isinstance(prompt_embeds, torch.Tensor): + prompt_embeds = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model)) + if not isinstance(thinker_result, torch.Tensor): + thinker_result = torch.zeros(0, self.talker.config.hidden_size, dtype=inputs_embeds.dtype, device=self._module_device(self.model)) + if not isinstance(prompt_token_ids, (list, torch.Tensor)): + prompt_token_ids = [] + if not isinstance(thinker_output_token_ids, (list, torch.Tensor)): + thinker_output_token_ids = [] + + req_input_ids, req_embeds = self._thinker_to_talker_prefill( + voice_type=info_dict.get("voice_type"), + output_prompt_embeds=thinker_result.to(inputs_embeds.dtype).to(self._module_device(self.model)), + output_token_ids=thinker_output_token_ids, + thinker_prompt_embeds=prompt_embeds.to(inputs_embeds.dtype).to(self._module_device(self.model)), + prompt_token_ids=prompt_token_ids, + ) + + if thinker_result.ndim == 2 and thinker_result.shape[0] > 0: + update_dict["thinker_reply_part"] = thinker_result[1:].detach().to("cpu").contiguous() + + return req_input_ids, req_embeds, update_dict + def _thinker_to_talker_prefill( self, voice_type: str, @@ -778,6 +702,38 @@ class Qwen2_5OmniForConditionalGeneration( ) return prompt_token_ids_processed, prompt_embeds + def thinker_to_talker_decode_one_step(self, + input_ids, + inputs_embeds, + **info_dict + ): + update_dict = {} + # choose step vector in priority order + step_vec = None + q = info_dict.get("thinker_reply_part", None) + if isinstance(q, torch.Tensor) and q.numel() > 0: + step_vec = q[0:1] + new_q = q[1:].detach().to("cpu").contiguous() + update_dict["thinker_reply_part"] = new_q + else: + # B) per-request provided decode vector (optional) + dv = info_dict.get("decode_output_prompt_embeds") if isinstance(info_dict, dict) else None + if isinstance(dv, torch.Tensor) and dv.numel() > 0: + step_vec = dv[0:1] if dv.ndim == 2 else dv.view(1, -1) + elif hasattr(self, "thinker_reply_part") and isinstance(self.thinker_reply_part, torch.Tensor) and self.thinker_reply_part.numel() > 0: + # C) fallback shared pool + step_vec = self.thinker_reply_part[0:1] + self.thinker_reply_part = self.thinker_reply_part[1:] + + if isinstance(step_vec, torch.Tensor) and step_vec.numel() > 0: + one_id = input_ids[0:1] + _, one_embed = self._thinker_to_talker_decode_one_step( + output_prompt_embeds=step_vec.to(inputs_embeds.dtype).to(self._module_device(self.model)), + output_token_ids=one_id, + ) + inputs_embeds[0] = one_embed[0] + return input_ids[0:1], inputs_embeds[0:1], update_dict + def _thinker_to_talker_decode_one_step( self, output_prompt_embeds, diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 7b793214b..84d312cf0 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -79,7 +79,6 @@ class GPUARModelRunner(OmniGPUModelRunner): # _prepare_inputs may reorder the batch, so we must gather multi # modal outputs after that to ensure the correct order - per_req_additional_information: dict[str, dict] | None = None if self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder: # Run the multimodal encoder if any. self._execute_mm_encoder(scheduler_output) @@ -100,14 +99,12 @@ class GPUARModelRunner(OmniGPUModelRunner): if hasattr(self, "_forward_additional_information"): self._forward_additional_information = None # New: per-request additional information for this step - per_req_additional_information = {} # Overlay custom prompt_embeds per request for the prompt portion; # collect additional_information (tensor/list) for prefill portion only for req_index, req_id in enumerate(self.input_batch.req_ids): req_state = self.requests[req_id] pe_cpu = getattr(req_state, "prompt_embeds_cpu", None) - addi_cpu = getattr(req_state, "additional_information_cpu", None) num_computed_tokens = int(self.input_batch.num_computed_tokens_cpu[req_index]) prompt_len = len(req_state.prompt_token_ids) prompt_remaining = max(0, prompt_len - num_computed_tokens) @@ -121,31 +118,6 @@ class GPUARModelRunner(OmniGPUModelRunner): ) start_offset = int(self.query_start_loc.cpu[req_index]) self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src) - # Build per-request additional information (no cross-request concat) - if addi_cpu is not None and isinstance(addi_cpu, dict): - req_info: dict[str, object] = {} - for k, v in addi_cpu.items(): - if isinstance(v, torch.Tensor): - # For prefill tokens, pass only the scheduled slice; - # for decode or no scheduled tokens, pass whole tensor - if overlay_len > 0: - try: - seg = ( - v[num_computed_tokens : num_computed_tokens + overlay_len] - .detach() - .to("cpu") - .contiguous() - ) - except Exception: - seg = v.detach().to("cpu").contiguous() - req_info[k] = seg - else: - req_info[k] = v.detach().to("cpu").contiguous() - elif isinstance(v, list): - req_info[k] = v - else: - req_info[k] = v - per_req_additional_information[req_id] = req_info input_ids = self.input_ids.gpu[:num_input_tokens] inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] @@ -200,6 +172,34 @@ class GPUARModelRunner(OmniGPUModelRunner): encoder_inputs = self._extract_encoder_inputs(scheduler_output) model_kwargs.update(encoder_inputs) + if self.model.has_process: + # Overlay custom prompt_embeds per request for the prompt portion; + # collect additional_information (tensor/list) for prefill portion only + 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) + + 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 + span_len = int(e) - int(s) + + # call the custom process function + req_input_ids, req_embeds, update_dict = self.model.process_fn( + input_ids=input_ids[s:e], + inputs_embeds=inputs_embeds[s:e], + **req_infos + ) + # TODO(Peiqi): the merge stage could move out from the critical path + self._merge_additional_information_update(req_id, update_dict) + + # update the inputs_embeds and input_ids + seg_len = min(span_len, req_embeds.shape[0]) + inputs_embeds[s : s + seg_len] = req_embeds[:seg_len] + if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len: + input_ids[s : s + seg_len] = req_input_ids + return ( num_scheduled_tokens, num_input_tokens, @@ -209,7 +209,6 @@ class GPUARModelRunner(OmniGPUModelRunner): positions, intermediate_tensors, model_kwargs, - per_req_additional_information, ) @torch.inference_mode() @@ -309,7 +308,6 @@ class GPUARModelRunner(OmniGPUModelRunner): positions, intermediate_tensors, model_kwargs, - per_req_additional_information, ) = self._preprocess( scheduler_output, num_scheduled_tokens_np, @@ -344,29 +342,9 @@ class GPUARModelRunner(OmniGPUModelRunner): record_function_or_nullcontext("Forward"), self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output, ): + # TODO(Peiqi): delete this model_kwargs_extra = {} - # Pass per-request additional information map for this step (no concat) - if per_req_additional_information: - model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information - # Always pass per-request runtime additional_information (persisted in request state) - try: - per_req_runtime_info = [] - for req_id in self.input_batch.req_ids: - req_state = self.requests.get(req_id) - info = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None - per_req_runtime_info.append(info if isinstance(info, dict) else {}) - model_kwargs_extra["runtime_additional_information"] = per_req_runtime_info - model_kwargs_extra["request_ids"] = self.input_batch.req_ids - # Pass each request's token span within the flattened sequence for this step, - # enabling the model to map decode/prefill by request - req_token_spans = [] - for req_index in range(len(self.input_batch.req_ids)): - start_offset = int(self.query_start_loc.cpu[req_index]) - sched_tokens = int(num_scheduled_tokens_np[req_index]) - req_token_spans.append((start_offset, start_offset + sched_tokens)) - model_kwargs_extra["request_token_spans"] = req_token_spans - except Exception: - pass + model_output = self.model( input_ids=input_ids, positions=positions, @@ -389,36 +367,8 @@ class GPUARModelRunner(OmniGPUModelRunner): aux_hidden_states = None hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) - # The model side may return per-request additional_information updates (model-agnostic channel). - # Convention: multimodal_outputs["additional_information_update"] is a list[dict] in batch order; - # the runner merges it into the corresponding request's additional_information_cpu for subsequent decode. - try: - if isinstance(multimodal_outputs, dict) and ( - "additional_information_update" in multimodal_outputs - or "additional_information_update_by_req_id" in multimodal_outputs - ): - # Option A: list[dict] in batch order - updates_list = multimodal_outputs.get("additional_information_update") - if isinstance(updates_list, list): - for idx, upd in enumerate(updates_list): - if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids): - continue - req_id = self.input_batch.req_ids[idx] - self._merge_additional_information_update(req_id, upd) - # Option B: dict[str, dict] keyed by req_id - updates_map = multimodal_outputs.get("additional_information_update_by_req_id") - if isinstance(updates_map, dict): - for req_id, upd in updates_map.items(): - if not isinstance(upd, dict): - continue - if req_id not in self.requests: - continue - self._merge_additional_information_update(req_id, upd) - except Exception as e: - logger.error( - f"Error merging for requests:{self.input_batch.req_ids} additional \ - information update: {e}, with the multimodal_outputs as {multimodal_outputs}" - ) + # TODO(Peiqi): move the merge stage out here + if not self.broadcast_pp_output: # Common case. if not get_pp_group().is_last_rank: