From bbb5a728714bbe3ca3c2d05261b125bbec08fd4f Mon Sep 17 00:00:00 2001 From: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com> Date: Mon, 8 Dec 2025 12:44:13 +0800 Subject: [PATCH] [BugFix] Fix qwen3omni thinker batching. (#207) Signed-off-by: yinpeiqi --- vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py | 6 +++--- vllm_omni/worker/gpu_ar_model_runner.py | 9 ++++++++- 2 files changed, 11 insertions(+), 4 deletions(-) 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 11aeb44d1..d9aeb42b0 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -342,9 +342,9 @@ class Qwen3OmniMoeForConditionalGeneration( 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] + 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 diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index c78aade3d..c8d434328 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -291,6 +291,11 @@ class GPUARModelRunner(OmniGPUModelRunner): # Case 1: tensor aligned on token dimension if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]: mm_payload[k] = v.detach().to("cpu")[prev_logits_index : logits_index + 1].contiguous() + elif isinstance(v, torch.Tensor) and v.shape[0] != hidden_states_cpu.shape[0]: + logger.error( + f"Error in merge multimodal outputs: Tensor dimension mismatch, \ + {v.shape} != {hidden_states_cpu.shape} for {k}" + ) # Case 2: nested dict of tensors aligned on token dimension (e.g., selected_hidden_layers) elif isinstance(v, dict): sub_dict: dict[str, torch.Tensor] = {} @@ -302,7 +307,9 @@ class GPUARModelRunner(OmniGPUModelRunner): if sub_dict: mm_payload[k] = sub_dict elif isinstance(v, list): - element: torch.Tensor = v[0] + element = v[0] + if isinstance(element, torch.Tensor): + element = element.detach().to("cpu").contiguous() multimodal_outputs[k] = v[1:] if len(v) > 1 else v mm_payload[k] = element except Exception as e: