From eb1b782dd15bef20a5a762614fb56e77abc4a8c0 Mon Sep 17 00:00:00 2001 From: zzh <67720593+BLANKETusers@users.noreply.github.com> Date: Mon, 8 Jun 2026 10:58:44 +0800 Subject: [PATCH] [BugFix] fix hunyuan image3 offline cot (#4174) Signed-off-by: zzh <943967662@qq.com> Signed-off-by: zengchuang Co-authored-by: zengchuang Co-authored-by: Canlin Guo --- .../model_executor/models/hunyuan_image3/hunyuan_image3.py | 1 + vllm_omni/platforms/npu/worker/npu_model_runner.py | 1 + vllm_omni/worker/gpu_model_runner.py | 5 +++++ 3 files changed, 7 insertions(+) diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py index 7b975796f..8c983cff1 100755 --- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py +++ b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py @@ -1457,6 +1457,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo HunyuanImage3Inputs: TypeAlias = HunyuanImage3PixelInputs prefer_model_sampler = True + logitsprocs_need_output_token_ids = True # Siglip2 ViT supports data-parallel encoding (mm_encoder_tp_mode="data"): # weights are replicated and the image batch is sharded across TP ranks. diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index 4137ffe92..3108b68a8 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py @@ -33,6 +33,7 @@ class OmniNPUModelRunner(OmniGPUModelRunner, NPUModelRunner): # in _pad_for_sequence_parallelism during execute_model. # This is a workaround for vllm-ascend not passing vllm_config to enable_sp(). enable_sp(self.vllm_config) + self._maybe_enable_output_token_ids_for_model_sampler() self._init_talker_mtp() @torch.inference_mode() diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index e74997618..70f1f58dc 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -143,9 +143,14 @@ class OmniGPUModelRunner(GPUModelRunner): @instrument(span_name="Loading (GPU)") def load_model(self, *args, **kwargs) -> None: super().load_model(*args, **kwargs) + self._maybe_enable_output_token_ids_for_model_sampler() self._init_talker_mtp() self._prewarm_attention_capture_workspaces() + def _maybe_enable_output_token_ids_for_model_sampler(self) -> None: + if getattr(self.model, "logitsprocs_need_output_token_ids", False): + self.input_batch.logitsprocs_need_output_token_ids = True + def _init_talker_mtp(self) -> None: # TODO move this model specific logic to a separate class # TTS model IS the talker (no .talker sub-attr); use getattr to support both Omni and TTS.