diff --git a/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml index 5e3be865e..f8eb59973 100644 --- a/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml @@ -45,7 +45,7 @@ stage_args: 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: true + enforce_eager: false trust_remote_code: true engine_output_type: latent # Output codec codes for code2wav # tensor_parallel_size: 2 diff --git a/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml index 7420e2745..3f727e0d9 100644 --- a/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml +++ b/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml @@ -45,7 +45,7 @@ stage_args: 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: true + enforce_eager: false trust_remote_code: true engine_output_type: latent # Output codec codes for code2wav # tensor_parallel_size: 2 diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 01c6ef7dd..4fe6b0083 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -1226,6 +1226,7 @@ async def _stage_worker_async( gen_output = res _gen_t1 = _time.time() _gen_ms = (_gen_t1 - _gen_t0) * 1000.0 + _gen_t0 = _gen_t1 await generation_out_q.put((rid, gen_output, _gen_ms)) except Exception as e: logger.exception("Failed on request %s: %s", rid, e) 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 05169303a..6ff339f58 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -8,6 +8,13 @@ from functools import cached_property import torch import torch.nn as nn +from transformers.generation.logits_process import ( + LogitsProcessorList, + RepetitionPenaltyLogitsProcessor, + SuppressTokensLogitsProcessor, + TemperatureLogitsWarper, + TopKLogitsWarper, +) from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( Qwen3OmniMoeCode2WavConfig, Qwen3OmniMoeConfig, @@ -146,13 +153,6 @@ class Qwen3OmniMoeForConditionalGeneration( self.talker.init_multi_modal(thinker_config) self.model = self.talker self.code2wav = None - from transformers.generation.logits_process import ( - LogitsProcessorList, - RepetitionPenaltyLogitsProcessor, - SuppressTokensLogitsProcessor, - TemperatureLogitsWarper, - TopKLogitsWarper, - ) self.logits_processors = LogitsProcessorList( [ @@ -367,46 +367,17 @@ class Qwen3OmniMoeForConditionalGeneration( **capture_kwargs, **kwargs, ) - # Compute thinker-side TTS token embeddings for BOS/EOS/PAD and expose via multimodal outputs. - # These will later be projected into talker text space by the talker stage. - multimodal_outputs = captured_layer_dict if captured_layer_dict is not None else {} - try: - thinker_tts_embeds = self.thinker.embed_input_ids(self.tts_tokens) # [1,3,thinker_hidden] - if ( - isinstance(thinker_tts_embeds, torch.Tensor) - and thinker_tts_embeds.ndim == 3 - and thinker_tts_embeds.shape[1] == 3 - ): - bos_eos_pad = thinker_tts_embeds.to(text_hidden_states.device).chunk(3, dim=1) # 3 * [1,1,H] - multimodal_outputs["tts_bos_embed"] = [bos_eos_pad[0]] - multimodal_outputs["tts_eos_embed"] = [bos_eos_pad[1]] - multimodal_outputs["tts_pad_embed"] = [bos_eos_pad[2]] - except Exception: - # Best-effort; absence will be handled by talker with fallbacks - pass - - # Return text-only output (with multimodal sidecar) - return OmniOutput( - text_hidden_states=(text_hidden_states.reshape(-1, text_hidden_states.shape[-1])), - multimodal_outputs=multimodal_outputs, - ) + return text_hidden_states, captured_layer_dict # ========== Stage 2.1: Talker ========== elif self.model_stage == "talker": if input_ids is None: input_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device) - self.code_predictor_hidden_pool = 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.embed_input_ids(input_ids) - if not is_profile and input_ids.shape[0] > 1: - self.input_ids_list = input_ids.tolist() - # TODO(Peiqi): temporal hack here to support voice_type. if not hasattr(self, "voice_type"): self.voice_type = voice_type @@ -418,19 +389,7 @@ class Qwen3OmniMoeForConditionalGeneration( positions=positions, inputs_embeds=inputs_embeds, ) - - # merge the code_predictor_codes from the info_dict list into a single tensor - multimodal_outputs: dict = None - if not is_profile: - # Here is the only place to use runtime_additional_information. After MTP in the - # preprocess function, the code_predictor_codes are stored in the info_dict list. - # We need to merge the tensors from different requests into a single tensor. - # In the future, we may allow user to custom an aggregated function. - info_dicts = kwargs.get("runtime_additional_information") - code_predictor_codes = [info.get("code_predictor_codes") for info in info_dicts] - multimodal_outputs = {"code_predictor_codes": torch.cat(code_predictor_codes, dim=0)} - - return OmniOutput(text_hidden_states=talker_hidden, multimodal_outputs=multimodal_outputs) + return talker_hidden # ========== Stage 3: Code2Wav ========== elif self.model_stage == "code2wav": @@ -459,10 +418,7 @@ class Qwen3OmniMoeForConditionalGeneration( "Batched input for code2wav is not supported yet, only the first audio tensor will be returned" ) - return OmniOutput( - text_hidden_states=None, - multimodal_outputs={"model_outputs": audio_tensors[0].reshape(1, -1)}, - ) + return audio_tensors # Fallback (shouldn't reach here) return OmniOutput( @@ -473,6 +429,63 @@ class Qwen3OmniMoeForConditionalGeneration( multimodal_outputs=None, ) + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs) -> OmniOutput: + """ + Make an OmniOutput object from model outputs. + Args: + model_outputs: Model outputs + """ + if isinstance(model_outputs, OmniOutput): + return model_outputs + + if self.model_stage == "thinker": + text_hidden_states, captured_layer_dict = model_outputs + # Compute thinker-side TTS token embeddings for BOS/EOS/PAD and expose via multimodal outputs. + # These will later be projected into talker text space by the talker stage. + multimodal_outputs = captured_layer_dict if captured_layer_dict is not None else {} + try: + thinker_tts_embeds = self.thinker.embed_input_ids(self.tts_tokens) # [1,3,thinker_hidden] + if ( + isinstance(thinker_tts_embeds, torch.Tensor) + and thinker_tts_embeds.ndim == 3 + and thinker_tts_embeds.shape[1] == 3 + ): + bos_eos_pad = thinker_tts_embeds.to(text_hidden_states.device).chunk(3, dim=1) # 3 * [1,1,H] + multimodal_outputs["tts_bos_embed"] = [bos_eos_pad[0]] + multimodal_outputs["tts_eos_embed"] = [bos_eos_pad[1]] + multimodal_outputs["tts_pad_embed"] = [bos_eos_pad[2]] + except Exception: + # Best-effort; absence will be handled by talker with fallbacks + pass + + # Return text-only output (with multimodal sidecar) + return OmniOutput( + text_hidden_states=(text_hidden_states.reshape(-1, text_hidden_states.shape[-1])), + multimodal_outputs=multimodal_outputs, + ) + elif self.model_stage == "talker": + talker_hidden = model_outputs + # merge the code_predictor_codes from the info_dict list into a single tensor + multimodal_outputs: dict = None + # Here is the only place to use runtime_additional_information. After MTP in the + # preprocess function, the code_predictor_codes are stored in the info_dict list. + # We need to merge the tensors from different requests into a single tensor. + # In the future, we may allow user to custom an aggregated function. + info_dicts = kwargs.get("runtime_additional_information") + code_predictor_codes = [info.get("code_predictor_codes") for info in info_dicts] + multimodal_outputs = {"code_predictor_codes": torch.cat(code_predictor_codes, dim=0)} + span_len = multimodal_outputs["code_predictor_codes"].shape[0] + talker_hidden = talker_hidden[:span_len] + return OmniOutput(text_hidden_states=talker_hidden, multimodal_outputs=multimodal_outputs) + elif self.model_stage == "code2wav": + audio_tensors = model_outputs + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": audio_tensors[0].reshape(1, -1)}, + ) + + return model_outputs + # ==================== Audio Generation ==================== def generate_audio(self, code: torch.Tensor, voice_type: str) -> torch.Tensor: @@ -624,9 +637,7 @@ class Qwen3OmniMoeForConditionalGeneration( input_ids, input_embeds, **info_dict ) - # execute talker MTP - input_embeds, code_predictor_codes = self.talker_mtp(input_ids, input_embeds, last_talker_hidden, text_step) - update_dict["code_predictor_codes"] = code_predictor_codes + update_dict["mtp_inputs"] = last_talker_hidden, text_step return input_ids, input_embeds, update_dict @@ -657,7 +668,7 @@ class Qwen3OmniMoeForConditionalGeneration( dtype=torch.long, ) inputs_embeds = (inputs_embeds + text_step).reshape(-1, self.talker_config.text_config.hidden_size) - return inputs_embeds, code_predictor_codes.squeeze(-1).detach().to("cpu").contiguous() + return inputs_embeds, code_predictor_codes.squeeze(-1) 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) @@ -872,7 +883,6 @@ class Qwen3OmniMoeForConditionalGeneration( def talker_preprocess_decode(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict): update_dict: dict[str, dict] = {} - try: q_tail = info_dict.get("tailing_text_hidden") if isinstance(q_tail, torch.Tensor) and q_tail.numel() > 0: diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py index bd990977a..fd5689976 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py @@ -14,6 +14,7 @@ import torch.nn as nn import torch.nn.functional as F from transformers import Cache, PretrainedConfig from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import Qwen3OmniMoeRotaryEmbedding +from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, ModelConfig, VllmConfig from vllm.logger import init_logger from vllm.model_executor.layers.layernorm import RMSNorm @@ -463,6 +464,7 @@ class Qwen3OmniCodePredictorBaseModel(nn.Module): return self.codec_embedding +@support_torch_compile class Qwen3OmniMoeTalkerCodePredictor(nn.Module): """ Code predictor wrapper matching HF structure. diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index d68d1e65f..255d45b45 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -4,6 +4,11 @@ from typing import Any import torch import torch.nn as nn import torch.nn.functional as F +from transformers.generation.logits_process import ( + LogitsProcessorList, + TopKLogitsWarper, + TopPLogitsWarper, +) from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( Qwen3OmniMoeTalkerConfig, ) @@ -134,6 +139,7 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( self.code_predictor = Qwen3OmniMoeTalkerCodePredictor( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "code_predictor") ) + self.empty_code = torch.empty((1, 0), dtype=torch.long) def code_predictor_forward( self, @@ -185,6 +191,12 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( all_codes_per_position = [] middle_hidden_states = [] # Collect hidden states for each position + logits_processors = LogitsProcessorList( + [ + TopKLogitsWarper(top_k=top_k), + TopPLogitsWarper(top_p=top_p), + ] + ) # Generate residual layers for each position for pos in range(seq_len): layer0_code = input_ids[:, pos : pos + 1] # [batch, 1] @@ -220,19 +232,13 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( # Use the corresponding lm_head for this layer logits = self.code_predictor.lm_head[layer_idx](hidden_state[:, -1:, :]) # [batch, 1, vocab_size] - from transformers.generation.logits_process import ( - LogitsProcessorList, - TopKLogitsWarper, - TopPLogitsWarper, - ) - logits_processors = LogitsProcessorList( - [ - TopKLogitsWarper(top_k=top_k), - TopPLogitsWarper(top_p=top_p), - ] - ) - input_ids_for_logits_processors = torch.tensor([pos_codes[1:]]).to(logits.device, dtype=torch.long) + if len(pos_codes) > 1: + input_ids_for_logits_processors = torch.cat(pos_codes[1:], dim=1).to( + device=logits.device, dtype=torch.long + ) + else: + input_ids_for_logits_processors = self.empty_code logits = logits_processors(input_ids_for_logits_processors, logits.squeeze(0)).unsqueeze(0) # Sample from the filtered distribution diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml index e077cf529..c63dc5638 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml @@ -16,7 +16,7 @@ stage_args: 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: true + enforce_eager: false trust_remote_code: true engine_output_type: latent # Output hidden states for talker distributed_executor_backend: "mp" @@ -47,7 +47,7 @@ stage_args: 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: true + enforce_eager: false trust_remote_code: true engine_output_type: latent # Output codec codes for code2wav # tensor_parallel_size: 2 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 6887fd969..246ea2996 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -160,11 +160,11 @@ def talker2code2wav( # Process each talker output for i, talker_output in enumerate(talker_outputs): output = talker_output.outputs[0] - + seq_len = len(output.token_ids) # Extract codec codes from talker output # Expected shape: [8, seq_len] (8-layer RVQ codes) codec_codes = ( - output.multimodal_output["code_predictor_codes"] + output.multimodal_output["code_predictor_codes"][-seq_len:] .to(torch.long) .transpose(0, 1) .cpu() diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 859027a8f..77b32c681 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -29,7 +29,6 @@ from vllm.v1.worker.gpu_model_runner import ( ) from vllm.v1.worker.utils import is_residual_scattered_for_sp -from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner @@ -185,8 +184,6 @@ class GPUARModelRunner(OmniGPUModelRunner): logits_index=logits_indices, sampler=self.sampler, ) - if isinstance(model_output, tuple): - model_output = OmniOutput(*model_output) with record_function_or_nullcontext("gpu_model_runner: postprocess"): if self.use_aux_hidden_state_outputs: diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 9e155f4ec..3ceed5bde 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Any, cast import numpy as np import torch +from vllm.compilation.cuda_graph import CUDAGraphWrapper from vllm.config import CUDAGraphMode from vllm.distributed.parallel_state import get_pp_group from vllm.forward_context import set_forward_context @@ -38,6 +39,25 @@ class OmniGPUModelRunner(GPUModelRunner): self._omni_num_scheduled_tokens_np: np.ndarray | None = None self._omni_last_model_output: object | None = None + def load_model(self, *args, **kwargs) -> None: + super().load_model(*args, **kwargs) + # TODO move this model specific logic to a separate class + if hasattr(self.model, "talker_mtp") and self.model.talker is not None: + self.talker_mtp = self.model.talker_mtp + cudagraph_mode = self.compilation_config.cudagraph_mode + assert cudagraph_mode is not None + if cudagraph_mode.has_full_cudagraphs(): + self.talker_mtp = CUDAGraphWrapper( + self.model.talker_mtp, self.vllm_config, runtime_mode=CUDAGraphMode.FULL + ) + hidden_size = self.model_config.hf_config.talker_config.text_config.hidden_size + self.talker_mtp_input_ids = self._make_buffer(self.max_num_reqs, dtype=torch.int32) + self.talker_mtp_inputs_embeds = self._make_buffer( + self.max_num_reqs, hidden_size, dtype=self.dtype, numpy=False + ) + self.last_talker_hidden = self._make_buffer(self.max_num_reqs, hidden_size, dtype=self.dtype, numpy=False) + self.text_step = self._make_buffer(self.max_num_reqs, hidden_size, dtype=self.dtype, numpy=False) + def _init_mrope_positions(self, req_state: CachedRequestState): """Initialize M-RoPE positions for multimodal inputs. @@ -543,6 +563,13 @@ class OmniGPUModelRunner(GPUModelRunner): ubatch_slices=ubatch_slices, ), ): + if self.model.talker is not None and hasattr(self.model, "talker_mtp") and num_tokens_padded == 1: + outputs = self.talker_mtp( + self.talker_mtp_input_ids.gpu[:num_tokens_padded], + self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded], + self.last_talker_hidden.gpu[:num_tokens_padded], + self.text_step.gpu[:num_tokens_padded], + ) outputs = self.model( input_ids=input_ids, positions=positions, @@ -866,6 +893,34 @@ 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 ) + # run talker mtp decode + if hasattr(self.model, "talker_mtp"): + _cudagraph_mode, batch_desc, _, _ = self._determine_batch_execution_and_padding( + num_tokens=span_len, + num_reqs=1, + num_scheduled_tokens_np=num_scheduled_tokens_np[req_index], + max_num_scheduled_tokens=1, + force_eager=span_len > 1, + use_cascade_attn=False, + ) + last_talker_hidden, text_step = update_dict.pop("mtp_inputs") + if _cudagraph_mode != CUDAGraphMode.NONE: + self.talker_mtp_input_ids.gpu[:span_len].copy_(req_input_ids) + self.talker_mtp_inputs_embeds.gpu[:span_len].copy_(req_embeds) + self.last_talker_hidden.gpu[:span_len].copy_(last_talker_hidden) + self.text_step.gpu[:span_len].copy_(text_step) + req_input_ids = self.talker_mtp_input_ids.gpu[:span_len] + req_embeds = self.talker_mtp_inputs_embeds.gpu[:span_len] + last_talker_hidden = self.last_talker_hidden.gpu[:span_len] + text_step = self.text_step.gpu[:span_len] + with set_forward_context( + None, self.vllm_config, cudagraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc + ): + req_embeds, code_predictor_codes = self.talker_mtp( + req_input_ids, req_embeds, last_talker_hidden, text_step + ) + update_dict["code_predictor_codes"] = code_predictor_codes + # TODO(Peiqi): the merge stage could move out from the critical path self._merge_additional_information_update(req_id, update_dict) @@ -909,6 +964,8 @@ class OmniGPUModelRunner(GPUModelRunner): **model_kwargs, **model_kwargs_extra, ) + if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"): + model_output = self.model.make_omni_output(model_output, **model_kwargs_extra) # Cache model output so later sample_tokens can consume multimodal results. self._omni_last_model_output = model_output return model_output