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 682474e9f..41fd1ecf5 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -15,6 +15,7 @@ from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( Qwen3OmniMoeThinkerConfig, ) from vllm.config import VllmConfig +from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP @@ -355,9 +356,18 @@ class Qwen3OmniMoeForConditionalGeneration( # ========== Stage 3: Code2Wav ========== elif self.model_stage == "code2wav": # Extract codec codes from input - codes = [] if input_ids.shape[0] % 16 == 0: - codes.append(input_ids.reshape(1, 16, -1)) + ubatch_slices = get_forward_context().ubatch_slices + if ubatch_slices is not None: + max_seq_len = max(ubatch_slices) // 16 + batch_size = len(ubatch_slices) + split_codes = torch.split(input_ids, ubatch_slices, dim=0) + codes = torch.zeros((batch_size, 16, max_seq_len), device=input_ids.device, dtype=input_ids.dtype) + for idx, code in enumerate(split_codes): + seq_len = code.shape[0] // 16 + codes[idx, :, :seq_len] = code.reshape(16, seq_len) + else: + codes = input_ids.reshape(1, 16, -1) else: logger.warning( ( @@ -373,17 +383,10 @@ class Qwen3OmniMoeForConditionalGeneration( torch.zeros(16 - input_ids.shape[0] % 16, dtype=torch.long, device=input_ids.device), ] ) - codes.append(input_ids_flatten.reshape(1, 16, -1)) + codes = input_ids_flatten.reshape(1, 16, -1) # Generate audio from codec codes - audio_tensors = [] - for code in codes: - audio_tensor = self.generate_audio(code, voice_type) - audio_tensors.append(audio_tensor) - if len(audio_tensors) > 1: - logger.warning( - "Batched input for code2wav is not supported yet, only the first audio tensor will be returned" - ) + audio_tensors = self.generate_audio(codes, voice_type) return audio_tensors @@ -448,14 +451,14 @@ class Qwen3OmniMoeForConditionalGeneration( audio_tensors = model_outputs return OmniOutput( text_hidden_states=None, - multimodal_outputs={"model_outputs": audio_tensors[0].reshape(1, -1)}, + multimodal_outputs={"model_outputs": [audio_tensor.reshape(1, -1) for audio_tensor in audio_tensors]}, ) return model_outputs # ==================== Audio Generation ==================== - def generate_audio(self, code: torch.Tensor, voice_type: str) -> torch.Tensor: + def generate_audio(self, code: torch.Tensor, voice_type: str) -> list[torch.Tensor]: """ Generate audio waveform from codec codes. @@ -484,20 +487,20 @@ class Qwen3OmniMoeForConditionalGeneration( talker_codes = talker_codes.expand(1, 16, -1) if self.vllm_config.model_config.async_chunk: - audio_tensor = self.code2wav.chunked_decode_streaming( + audio_tensors = 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( + audio_tensors = self.code2wav.chunked_decode( talker_codes, chunk_size=300, left_context_size=25, ) - return audio_tensor + return audio_tensors # ==================== Thinker-Talker Projection ==================== diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py index 7adb6f96f..f1e08c435 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py @@ -22,6 +22,7 @@ from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( SnakeBeta, ) from vllm.config import VllmConfig # type: ignore +from vllm.forward_context import get_forward_context from vllm.logger import init_logger # type: ignore from vllm.model_executor.models.utils import ( # type: ignore AutoWeightsLoader, @@ -162,7 +163,7 @@ class Qwen3OmniMoeCode2Wav(nn.Module): codes: torch.Tensor, chunk_size: int = 300, left_context_size: int = 25, - ) -> torch.Tensor: + ) -> list[torch.Tensor]: """ Decode long sequences in chunks to avoid OOM. @@ -174,7 +175,9 @@ class Qwen3OmniMoeCode2Wav(nn.Module): left_context_size: Number of overlapping frames for context Returns: - waveform: [batch, 1, waveform_len] - Complete waveform + list[torch.Tensor]: Complete waveform decoded from the input + codes. For ``batch_size == 1``, this is a list containing a + single tensor with shape ``[1, waveform_len]``. """ wavs = [] start_index = 0 @@ -194,14 +197,26 @@ class Qwen3OmniMoeCode2Wav(nn.Module): start_index = end_index - return torch.cat(wavs, dim=-1) + ubatch_slices = get_forward_context().ubatch_slices + if ubatch_slices is not None: + code_seq_lens = [seq_len // self.config.num_quantizers for seq_len in ubatch_slices] + else: + # Fallback: assume all batch elements share the same sequence length. + # Create one entry per batch so that each element is processed. + code_seq_lens = [codes.shape[-1]] * codes.shape[0] + batch_wav = torch.cat(wavs, dim=-1) + wavs = [] + for idx, code_seq_len in enumerate(code_seq_lens): + wav_chunk = batch_wav[idx, :, : code_seq_len * self.total_upsample] + wavs.append(wav_chunk) + return wavs def chunked_decode_streaming( self, codes: torch.Tensor, chunk_size: int = 25, left_context_size: int = 25, - ) -> torch.Tensor: + ) -> list[torch.Tensor]: """ Decode long sequences in chunks to avoid OOM. @@ -213,21 +228,31 @@ class Qwen3OmniMoeCode2Wav(nn.Module): left_context_size: Number of overlapping frames for context Returns: - waveform: [batch, 1, waveform_len] - Complete waveform + list[torch.Tensor]: Complete waveform decoded from the input + codes. For ``batch_size == 1``, this is a list containing a + single tensor with shape ``[1, waveform_len]``. """ - 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) + wavs = [] + batch_wav = self(codes) + ubatch_slices = get_forward_context().ubatch_slices + if ubatch_slices is not None: + code_seq_lens = [seq_len // self.config.num_quantizers for seq_len in ubatch_slices] + else: + # Fallback: assume all batch elements share the same sequence length. + # Create one entry per batch so that each element is processed. + code_seq_lens = [codes.shape[-1]] * codes.shape[0] + for idx, code_seq_len in enumerate(code_seq_lens): + # TODO: need to optimize algorithms, current only support + # chunk_size = left_context_size = 25 + if code_seq_len <= chunk_size: + context_size = 0 + else: + context_size = left_context_size + # Remove context from output (context_size * total_upsample samples) + wav_chunk = batch_wav[idx, :, context_size * self.total_upsample : code_seq_len * self.total_upsample] + wavs.append(wav_chunk) + return wavs def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights from HuggingFace checkpoint.""" 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 e3dcf940f..90f243151 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml @@ -72,7 +72,7 @@ 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: code2wav model_arch: Qwen3OmniMoeForConditionalGeneration 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 1a8b60aa1..a1b082b87 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 @@ -73,7 +73,7 @@ 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: code2wav model_arch: Qwen3OmniMoeForConditionalGeneration diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml index 0bc6e48e5..30e894ae6 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml @@ -78,7 +78,7 @@ 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: code2wav model_arch: Qwen3OmniMoeForConditionalGeneration diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 2c82b5cd4..c52ab1e8b 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -50,12 +50,26 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): """ def _update_request_states(self, scheduler_output: SchedulerOutput): + # remove requests + for req_id in scheduler_output.finished_req_ids: + self.input_batch.remove_request(req_id) + scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys() + cached_req_ids = self.input_batch.req_id_to_index.keys() + resumed_req_ids = scheduler_output.scheduled_cached_reqs.resumed_req_ids + unscheduled_req_ids = cached_req_ids - (scheduled_req_ids - resumed_req_ids) + for req_id in unscheduled_req_ids: + self.input_batch.remove_request(req_id) cached_reqs = scheduler_output.scheduled_cached_reqs - for _, req_id in enumerate(cached_reqs.req_ids): + req_states = [] + for req_id in 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) + req_states.append(req_state) + # Remove the request from the current input batch only if it is still present. + if req_id in self.input_batch.req_id_to_index: + self.input_batch.remove_request(req_id) + for req_state in req_states: # update the request state in self.input_batch self.input_batch.add_request(req_state) # Only relevant for models using M-RoPE (e.g, Qwen2-VL) @@ -244,6 +258,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): # Mark KV scales as calculated after the first forward pass self.calculate_kv_scales = False + if ubatch_slices_padded is None: + # reuse ubatch_slices_padded for code2wav batching + ubatch_slices_padded = tokens + # Run the model. # Use persistent buffers for CUDA graphs. with ( @@ -342,11 +360,22 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): {"model_outputs": out.detach().to("cpu").contiguous() if out is not None else None} ) elif isinstance(multimodal_outputs, dict): - mm_payload = {} - for key, out in multimodal_outputs.items(): - if out is not None and isinstance(out, torch.Tensor): - mm_payload[key] = out.detach().to("cpu").contiguous() - pooler_output.append(mm_payload) + num_reqs = self.input_batch.num_reqs + for i in range(num_reqs): + mm_payload = {} + for key, out in multimodal_outputs.items(): + if isinstance(out, list): + if len(out) != num_reqs: + raise ValueError( + f"Multimodal output list for key '{key}' has length {len(out)} " + f"but expected {num_reqs} (one entry per request)." + ) + mm_payload[key] = out[i].detach().to("cpu").contiguous() + elif isinstance(out, torch.Tensor): + mm_payload[key] = out.detach().to("cpu").contiguous() + else: + logger.warning(f"Unsupported multimodal output type for key '{key}': {type(out)}") + pooler_output.append(mm_payload) else: raise RuntimeError("Unsupported diffusion output type") output = OmniModelRunnerOutput(